Switch contract storage to child trie (#2002)

substrate/srml/contract/src/account_db.rs

@@ -16,14 +16,15 @@
 
 //! Auxilliaries to help with managing partial changes to accounts state.
 
-use super::{CodeHash, CodeHashOf, StorageOf, Trait};
+use super::{CodeHash, CodeHashOf, Trait, AccountInfo, TrieId, AccountInfoOf};
 use {balances, system};
 use rstd::cell::RefCell;
+use rstd::rc::Rc;
 use rstd::collections::btree_map::{BTreeMap, Entry};
 use rstd::prelude::*;
 use runtime_primitives::traits::Zero;
-use srml_support::{StorageMap, StorageDoubleMap, traits::{UpdateBalanceOutcome,
-	SignedImbalance, Currency, Imbalance}};
+use srml_support::{StorageMap, traits::{UpdateBalanceOutcome,
+	SignedImbalance, Currency, Imbalance}, storage::child};
 
 pub struct ChangeEntry<T: Trait> {
 	balance: Option<T::Balance>,
@@ -45,8 +46,50 @@ impl<T: Trait> Default for ChangeEntry<T> {
 
 pub type ChangeSet<T> = BTreeMap<<T as system::Trait>::AccountId, ChangeEntry<T>>;
 
+#[derive(Clone, Default)]
+pub struct AccountTrieIdMapping<A: Ord> {
+	to_account: BTreeMap<TrieId, A>,
+	to_key: BTreeMap<A, TrieId>,
+	// this lock is related to the way overlaydb stack
+	// if set it must be unset at the lower level
+	lock: bool,
+}
+
+impl<A: Clone + Ord> AccountTrieIdMapping<A> {
+
+	pub fn new() -> Self {
+		AccountTrieIdMapping {
+			to_account: BTreeMap::new(),
+			to_key: BTreeMap::new(),
+			lock: false,
+		}
+	}
+
+	pub fn lock(&mut self) {
+		self.lock = true;
+	}
+	pub fn unlock(&mut self) {
+		self.lock = false;
+	}
+	pub fn insert(&mut self, account: A, ks: TrieId) {
+		self.to_account.insert(ks.clone(), account.clone());
+		self.to_key.insert(account, ks);
+	}
+	pub fn get_trieid(&self, account: &A) -> Option<&TrieId> {
+		if self.lock { return None }
+		self.to_key.get(account)
+	}
+	pub fn get_account(&self, ks: &TrieId) -> Option<&A> {
+		if self.lock { return None }
+		self.to_account.get(ks)
+	}
+
+}
+
 pub trait AccountDb<T: Trait> {
-	fn get_storage(&self, account: &T::AccountId, location: &[u8]) -> Option<Vec<u8>>;
+	fn get_account_info(&self, account: &T::AccountId) -> Option<AccountInfo>;
+	fn get_or_create_trieid(&self, account: &T::AccountId) -> TrieId;
+	fn get_storage(&self, trie_id: &TrieId, location: &[u8]) -> Option<Vec<u8>>;
 	fn get_code(&self, account: &T::AccountId) -> Option<CodeHash<T>>;
 	fn get_balance(&self, account: &T::AccountId) -> T::Balance;
 
@@ -55,8 +98,18 @@ pub trait AccountDb<T: Trait> {
 
 pub struct DirectAccountDb;
 impl<T: Trait> AccountDb<T> for DirectAccountDb {
-	fn get_storage(&self, account: &T::AccountId, location: &[u8]) -> Option<Vec<u8>> {
-		<StorageOf<T>>::get(account, &location.to_vec())
+	fn get_account_info(&self, account: &T::AccountId) -> Option<AccountInfo> {
+		let res: Option<AccountInfo> = AccountInfoOf::<T>::get(account);
+		res
+	}
+	fn get_or_create_trieid(&self, account: &T::AccountId) -> TrieId {
+		use super::TrieIdGenerator;
+		<Self as AccountDb<T>>::get_account_info(self, account)
+			.map(|s|s.trie_id)
+			.unwrap_or_else(||<T as Trait>::TrieIdGenerator::trie_id(account))
+	}
+	fn get_storage(&self, trie_id: &TrieId, location: &[u8]) -> Option<Vec<u8>> {
+		child::get_raw(trie_id, location)
 	}
 	fn get_code(&self, account: &T::AccountId) -> Option<CodeHash<T>> {
 		<CodeHashOf<T>>::get(account)
@@ -67,12 +120,13 @@ impl<T: Trait> AccountDb<T> for DirectAccountDb {
 	fn commit(&mut self, s: ChangeSet<T>) {
 		let mut total_imbalance = SignedImbalance::zero();
 		for (address, changed) in s.into_iter() {
+			let trieid = <Self as AccountDb<T>>::get_or_create_trieid(&self, &address);
 			if let Some(balance) = changed.balance {
 				let (imbalance, outcome) = balances::Module::<T>::ensure_free_balance_is(&address, balance);
 				total_imbalance = total_imbalance.merge(imbalance);
 				if let UpdateBalanceOutcome::AccountKilled = outcome {
 					// Account killed. This will ultimately lead to calling `OnFreeBalanceZero` callback
-					// which will make removal of CodeHashOf and StorageOf for this account.
+					// which will make removal of CodeHashOf and AccountStorage for this account.
 					// In order to avoid writing over the deleted properties we `continue` here.
 					continue;
 				}
@@ -86,9 +140,9 @@ impl<T: Trait> AccountDb<T> for DirectAccountDb {
 			}
 			for (k, v) in changed.storage.into_iter() {
 				if let Some(value) = v {
-					<StorageOf<T>>::insert(&address, &k, value);
+					child::put_raw(&trieid[..], &k, &value[..]);
 				} else {
-					<StorageOf<T>>::remove(&address, &k);
+					child::kill(&trieid[..], &k);
 				}
 			}
 		}
@@ -104,19 +158,30 @@ impl<T: Trait> AccountDb<T> for DirectAccountDb {
 		}
 	}
 }
-
 pub struct OverlayAccountDb<'a, T: Trait + 'a> {
 	local: RefCell<ChangeSet<T>>,
+	trie_account: Rc<RefCell<AccountTrieIdMapping<<T as system::Trait>::AccountId>>>,
+	trie_account_cache: bool,
 	underlying: &'a AccountDb<T>,
 }
 impl<'a, T: Trait> OverlayAccountDb<'a, T> {
-	pub fn new(underlying: &'a AccountDb<T>) -> OverlayAccountDb<'a, T> {
+	pub fn new(
+		underlying: &'a AccountDb<T>,
+		trie_account: Rc<RefCell<AccountTrieIdMapping<<T as system::Trait>::AccountId>>>,
+		trie_account_cache: bool,
+	) -> OverlayAccountDb<'a, T> {
 		OverlayAccountDb {
 			local: RefCell::new(ChangeSet::new()),
+			trie_account,
+			trie_account_cache,
 			underlying,
 		}
 	}
 
+	pub fn reg_cache_new_rc(&self) -> Rc<RefCell<AccountTrieIdMapping<<T as system::Trait>::AccountId>>> {
+		self.trie_account.clone()
+	}
+
 	pub fn into_change_set(self) -> ChangeSet<T> {
 		self.local.into_inner()
 	}
@@ -127,13 +192,13 @@ impl<'a, T: Trait> OverlayAccountDb<'a, T> {
 		location: Vec<u8>,
 		value: Option<Vec<u8>>,
 	) {
-		self.local
-			.borrow_mut()
+		self.local.borrow_mut()
 			.entry(account.clone())
 			.or_insert(Default::default())
 			.storage
 			.insert(location, value);
 	}
+
 	pub fn set_code(&mut self, account: &T::AccountId, code: Option<CodeHash<T>>) {
 		self.local
 			.borrow_mut()
@@ -151,13 +216,39 @@ impl<'a, T: Trait> OverlayAccountDb<'a, T> {
 }
 
 impl<'a, T: Trait> AccountDb<T> for OverlayAccountDb<'a, T> {
-	fn get_storage(&self, account: &T::AccountId, location: &[u8]) -> Option<Vec<u8>> {
-		self.local
+	fn get_account_info(&self, account: &T::AccountId) -> Option<AccountInfo> {
+		let v = self.underlying.get_account_info(account);
+		if self.trie_account_cache {
+			v.as_ref().map(|v|self.trie_account.as_ref().borrow_mut().insert(account.clone(), v.trie_id.clone()));
+		}
+		v
+	}
+	fn get_or_create_trieid(&self, account: &T::AccountId) -> TrieId {
+		if self.trie_account_cache {
+			let mut ka_mut = self.trie_account.as_ref().borrow_mut();
+			if let Some(v) = ka_mut.get_trieid(account) {
+				v.clone()
+			} else {
+				ka_mut.unlock();
+				let v = self.underlying.get_or_create_trieid(account);
+				ka_mut.insert(account.clone(), v.clone());
+				v
+			}
+		} else {
+			let res = self.trie_account.as_ref().borrow().get_trieid(account).map(|v|v.clone());
+			res.unwrap_or_else(|| {
+				self.trie_account.as_ref().borrow_mut().lock();
+				self.underlying.get_or_create_trieid(account)
+			})
+		}
+	}
+	fn get_storage(&self, ks: &TrieId, location: &[u8]) -> Option<Vec<u8>> {
+		self.trie_account.as_ref().borrow().get_account(ks).and_then(|account| self.local
 			.borrow()
-			.get(account)
+			.get(&account)
 			.and_then(|a| a.storage.get(location))
 			.cloned()
-			.unwrap_or_else(|| self.underlying.get_storage(account, location))
+			.unwrap_or_else(|| self.underlying.get_storage(ks, location)))
 	}
 	fn get_code(&self, account: &T::AccountId) -> Option<CodeHash<T>> {
 		self.local

substrate/srml/contract/src/exec.rs

@@ -14,11 +14,13 @@
 // You should have received a copy of the GNU General Public License
 // along with Substrate. If not, see <http://www.gnu.org/licenses/>.
 
-use super::{CodeHash, Config, ContractAddressFor, Event, RawEvent, Trait};
-use crate::account_db::{AccountDb, DirectAccountDb, OverlayAccountDb};
+use super::{CodeHash, Config, ContractAddressFor, Event, RawEvent, Trait, TrieId};
+use crate::account_db::{AccountDb, DirectAccountDb, OverlayAccountDb, AccountTrieIdMapping};
 use crate::gas::{GasMeter, Token, approx_gas_for_balance};
 
 use rstd::prelude::*;
+use rstd::cell::RefCell;
+use rstd::rc::Rc;
 use runtime_primitives::traits::{CheckedAdd, CheckedSub, Zero};
 use srml_support::traits::{WithdrawReason, Currency};
 use timestamp;
@@ -225,6 +227,7 @@ impl<T: Trait> Token<T> for ExecFeeToken {
 
 pub struct ExecutionContext<'a, T: Trait + 'a, V, L> {
 	pub self_account: T::AccountId,
+	pub self_trieid: TrieId,
 	pub overlay: OverlayAccountDb<'a, T>,
 	pub depth: usize,
 	pub events: Vec<Event<T>>,
@@ -244,9 +247,11 @@ where
 	///
 	/// The specified `origin` address will be used as `sender` for
 	pub fn top_level(origin: T::AccountId, cfg: &'a Config<T>, vm: &'a V, loader: &'a L) -> Self {
-		let overlay = OverlayAccountDb::<T>::new(&DirectAccountDb);
+		let overlay = OverlayAccountDb::<T>::new(&DirectAccountDb, Rc::new(RefCell::new(AccountTrieIdMapping::new())), true);
+		let self_trieid = overlay.get_or_create_trieid(&origin);
 		ExecutionContext {
 			self_account: origin,
+			self_trieid,
 			depth: 0,
 			overlay,
 			events: Vec::new(),
@@ -258,9 +263,11 @@ where
 	}
 
 	fn nested(&self, overlay: OverlayAccountDb<'a, T>, dest: T::AccountId) -> Self {
+		let self_trieid = overlay.get_or_create_trieid(&dest);
 		ExecutionContext {
-			overlay: overlay,
+			overlay,
 			self_account: dest,
+			self_trieid,
 			depth: self.depth + 1,
 			events: Vec::new(),
 			calls: Vec::new(),
@@ -295,7 +302,7 @@ where
 
 		let (change_set, events, calls) = {
 			let mut nested = self.nested(
-				OverlayAccountDb::new(&self.overlay),
+				OverlayAccountDb::new(&self.overlay, self.overlay.reg_cache_new_rc(), false),
 				dest.clone()
 			);
 
@@ -370,7 +377,8 @@ where
 		}
 
 		let (change_set, events, calls) = {
-			let mut overlay = OverlayAccountDb::new(&self.overlay);
+			let mut overlay = OverlayAccountDb::new(&self.overlay, self.overlay.reg_cache_new_rc(), false);
+				
 			overlay.set_code(&dest, Some(code_hash.clone()));
 			let mut nested = self.nested(overlay, dest.clone());
 
@@ -554,7 +562,7 @@ where
 	type T = T;
 
 	fn get_storage(&self, key: &[u8]) -> Option<Vec<u8>> {
-		self.ctx.overlay.get_storage(&self.ctx.self_account, key)
+		self.ctx.overlay.get_storage(&self.ctx.self_trieid, key)
 	}
 
 	fn set_storage(&mut self, key: &[u8], value: Option<Vec<u8>>) {
@@ -639,9 +647,9 @@ mod tests {
 	use crate::{CodeHash, Config};
 	use runtime_io::with_externalities;
 	use std::cell::RefCell;
+	use std::rc::Rc;
 	use std::collections::HashMap;
 	use std::marker::PhantomData;
-	use std::rc::Rc;
 	use assert_matches::assert_matches;
 
 	const ALICE: u64 = 1;

substrate/srml/contract/src/lib.rs

@@ -63,7 +63,7 @@ mod wasm;
 mod tests;
 
 use crate::exec::ExecutionContext;
-use crate::account_db::AccountDb;
+use crate::account_db::{AccountDb, DirectAccountDb};
 
 #[cfg(feature = "std")]
 use serde_derive::{Serialize, Deserialize};
@@ -73,16 +73,16 @@ use rstd::marker::PhantomData;
 use parity_codec::{Codec, Encode, Decode};
 use runtime_primitives::traits::{Hash, As, SimpleArithmetic,Bounded, StaticLookup};
 use srml_support::dispatch::{Result, Dispatchable};
-use srml_support::{Parameter, StorageMap, StorageValue, StorageDoubleMap, decl_module, decl_event, decl_storage};
+use srml_support::{Parameter, StorageMap, StorageValue, decl_module, decl_event, decl_storage, storage::child};
 use srml_support::traits::{OnFreeBalanceZero, OnUnbalanced};
 use system::{ensure_signed, RawOrigin};
-use runtime_io::{blake2_256, twox_128};
 use timestamp;
 
 pub type CodeHash<T> = <T as system::Trait>::Hash;
+pub type TrieId = Vec<u8>;
 
 /// A function that generates an `AccountId` for a contract upon instantiation.
-pub trait ContractAddressFor<CodeHash, AccountId: Sized> {
+pub trait ContractAddressFor<CodeHash, AccountId> {
 	fn contract_address_for(code_hash: &CodeHash, data: &[u8], origin: &AccountId) -> AccountId;
 }
 
@@ -91,6 +91,48 @@ pub trait ComputeDispatchFee<Call, Balance> {
 	fn compute_dispatch_fee(call: &Call) -> Balance;
 }
 
+#[derive(Encode,Decode,Clone,Debug)]
+/// Information for managing an acocunt and its sub trie abstraction.
+/// This is the required info to cache for an account
+pub struct AccountInfo {
+	/// unique ID for the subtree encoded as a byte
+	pub trie_id: TrieId,
+	/// the size of stored value in octet
+	pub current_mem_stored: u64,
+}
+
+/// Get a trie id (trie id must be unique and collision resistant depending upon its context)
+/// Note that it is different than encode because trie id should have collision resistance
+/// property (being a proper uniqueid).
+pub trait TrieIdGenerator<AccountId> {
+	/// get a trie id for an account, using reference to parent account trie id to ensure
+	/// uniqueness of trie id 
+	/// The implementation must ensure every new trie id is unique: two consecutive call with the
+	/// same parameter needs to return different trie id values.
+	fn trie_id(account_id: &AccountId) -> TrieId;
+}
+
+/// Get trie id from `account_id`
+pub struct TrieIdFromParentCounter<T: Trait>(PhantomData<T>);
+
+/// This generator use inner counter for account id and apply hash over `AccountId +
+/// accountid_counter`
+impl<T: Trait> TrieIdGenerator<T::AccountId> for TrieIdFromParentCounter<T>
+where
+	T::AccountId: AsRef<[u8]>
+{
+	fn trie_id(account_id: &T::AccountId) -> TrieId {
+		// note that skipping a value due to error is not an issue here.
+		// we only need uniqueness, not sequence.
+		let new_seed = <AccountCounter<T>>::mutate(|v| v.wrapping_add(1));
+
+		let mut buf = Vec::new();
+		buf.extend_from_slice(account_id.as_ref());
+		buf.extend_from_slice(&new_seed.to_le_bytes()[..]);
+		T::Hashing::hash(&buf[..]).as_ref().into()
+	}
+}
+
 pub trait Trait: balances::Trait + timestamp::Trait {
 	/// The outer call dispatch type.
 	type Call: Parameter + Dispatchable<Origin=<Self as system::Trait>::Origin>;
@@ -109,6 +151,8 @@ pub trait Trait: balances::Trait + timestamp::Trait {
 	/// It is recommended (though not required) for this function to return a fee that would be taken
 	/// by executive module for regular dispatch.
 	type ComputeDispatchFee: ComputeDispatchFee<Self::Call, <Self as balances::Trait>::Balance>;
+	/// trieid id generator
+	type TrieIdGenerator: TrieIdGenerator<Self::AccountId>;
 
 	/// Handler for the unbalanced reduction when making a gas payment.
 	type GasPayment: OnUnbalanced<balances::NegativeImbalance<Self>>;
@@ -214,8 +258,8 @@ decl_module! {
 			let result = ctx.call(dest, value, &mut gas_meter, &data, exec::EmptyOutputBuf::new());
 
 			if let Ok(_) = result {
-				// Commit all changes that made it thus far into the persistent storage.
-				account_db::DirectAccountDb.commit(ctx.overlay.into_change_set());
+				// Commit all changes that made it thus far into the persistant storage.
+				DirectAccountDb.commit(ctx.overlay.into_change_set());
 
 				// Then deposit all events produced.
 				ctx.events.into_iter().for_each(Self::deposit_event);
@@ -268,7 +312,7 @@ decl_module! {
 
 			if let Ok(_) = result {
 				// Commit all changes that made it thus far into the persistant storage.
-				account_db::DirectAccountDb.commit(ctx.overlay.into_change_set());
+				DirectAccountDb.commit(ctx.overlay.into_change_set());
 
 				// Then deposit all events produced.
 				ctx.events.into_iter().for_each(Self::deposit_event);
@@ -344,34 +388,19 @@ decl_storage! {
 		pub PristineCode: map CodeHash<T> => Option<Vec<u8>>;
 		/// A mapping between an original code hash and instrumented wasm code, ready for the execution.
 		pub CodeStorage: map CodeHash<T> => Option<wasm::PrefabWasmModule>;
-	}
-}
-
-/// The storage items associated with an account/key.
-///
-pub(crate) struct StorageOf<T>(rstd::marker::PhantomData<T>);
-impl<T: Trait> StorageDoubleMap for StorageOf<T> {
-	const PREFIX: &'static [u8] = b"con:sto:";
-	type Key1 = T::AccountId;
-	type Key2 = Vec<u8>;
-	type Value = Vec<u8>;
-
-	/// Hashed by XX
-	fn derive_key1(key1_data: Vec<u8>) -> Vec<u8> {
-		twox_128(&key1_data).to_vec()
-	}
-
-	/// Blake2 is used for `Key2` is because it will be used as a key for contract's storage and
-	/// thus will be susceptible for a untrusted input.
-	fn derive_key2(key2_data: Vec<u8>) -> Vec<u8> {
-		blake2_256(&key2_data).to_vec()
+		/// The subtrie counter
+		pub AccountCounter: u64 = 0;
+		/// The code associated with a given account.
+		pub AccountInfoOf: map T::AccountId => Option<AccountInfo>;
 	}
 }
 
 impl<T: Trait> OnFreeBalanceZero<T::AccountId> for Module<T> {
 	fn on_free_balance_zero(who: &T::AccountId) {
 		<CodeHashOf<T>>::remove(who);
-		<StorageOf<T>>::remove_prefix(who);
+		<DirectAccountDb as AccountDb<T>>::get_account_info(&DirectAccountDb, who).map(|subtrie| {
+			child::kill_storage(&subtrie.trie_id);
+		});
 	}
 }
 

substrate/srml/contract/src/tests.rs

@@ -24,7 +24,7 @@ use runtime_primitives::testing::{Digest, DigestItem, H256, Header, UintAuthorit
 use runtime_primitives::traits::{BlakeTwo256, IdentityLookup};
 use runtime_primitives::BuildStorage;
 use runtime_io;
-use srml_support::{StorageMap, StorageDoubleMap, assert_ok, impl_outer_event, impl_outer_dispatch,
+use srml_support::{storage::child, StorageMap, assert_ok, impl_outer_event, impl_outer_dispatch,
 	impl_outer_origin, traits::Currency};
 use substrate_primitives::Blake2Hasher;
 use system::{self, Phase, EventRecord};
@@ -32,9 +32,13 @@ use {wabt, balances, consensus};
 use hex_literal::*;
 use assert_matches::assert_matches;
 use crate::{
-	ContractAddressFor, GenesisConfig, Module, RawEvent, StorageOf,
-	Trait, ComputeDispatchFee
+	ContractAddressFor, GenesisConfig, Module, RawEvent,
+	Trait, ComputeDispatchFee, TrieIdGenerator, TrieId,
+	AccountInfo, AccountInfoOf,
 };
+use substrate_primitives::storage::well_known_keys;
+use parity_codec::{Encode, Decode, KeyedVec};
+use std::sync::atomic::{AtomicUsize, Ordering};
 
 mod contract {
 	// Re-export contents of the root. This basically
@@ -68,7 +72,7 @@ impl system::Trait for Test {
 	type Hashing = BlakeTwo256;
 	type Digest = Digest;
 	type AccountId = u64;
-	type Lookup = IdentityLookup<u64>;
+	type Lookup = IdentityLookup<Self::AccountId>;
 	type Header = Header;
 	type Event = MetaEvent;
 	type Log = DigestItem;
@@ -97,6 +101,7 @@ impl Trait for Test {
 	type DetermineContractAddress = DummyContractAddressFor;
 	type Event = MetaEvent;
 	type ComputeDispatchFee = DummyComputeDispatchFee;
+	type TrieIdGenerator = DummyTrieIdGenerator;
 	type GasPayment = ();
 }
 
@@ -111,6 +116,17 @@ impl ContractAddressFor<H256, u64> for DummyContractAddressFor {
 	}
 }
 
+static KEY_COUNTER: AtomicUsize = AtomicUsize::new(0);
+
+pub struct DummyTrieIdGenerator;
+impl TrieIdGenerator<u64> for DummyTrieIdGenerator {
+	fn trie_id(account_id: &u64) -> TrieId {
+		let mut res = KEY_COUNTER.fetch_add(1, Ordering::Relaxed).to_le_bytes().to_vec();
+		res.extend_from_slice(&account_id.to_le_bytes());
+		res
+	}
+}
+
 pub struct DummyComputeDispatchFee;
 impl ComputeDispatchFee<Call, u64> for DummyComputeDispatchFee {
 	fn compute_dispatch_fee(call: &Call) -> u64 {
@@ -217,18 +233,29 @@ fn refunds_unused_gas() {
 
 #[test]
 fn account_removal_removes_storage() {
+	let unique_id1 = b"unique_id1";
+	let unique_id2 = b"unique_id2";
 	with_externalities(
 		&mut ExtBuilder::default().existential_deposit(100).build(),
 		|| {
 			// Setup two accounts with free balance above than exsistential threshold.
 			{
 				Balances::deposit_creating(&1, 110);
-				<StorageOf<Test>>::insert(&1, &b"foo".to_vec(), b"1".to_vec());
-				<StorageOf<Test>>::insert(&1, &b"bar".to_vec(), b"2".to_vec());
+				AccountInfoOf::<Test>::insert(1, &AccountInfo {
+					trie_id: unique_id1.to_vec(),
+					current_mem_stored: 0,
+				});
+				child::put(&unique_id1[..], &b"foo".to_vec(), &b"1".to_vec());
+				assert_eq!(child::get(&unique_id1[..], &b"foo".to_vec()), Some(b"1".to_vec()));
+				child::put(&unique_id1[..], &b"bar".to_vec(), &b"2".to_vec());
 
 				Balances::deposit_creating(&2, 110);
-				<StorageOf<Test>>::insert(&2, &b"hello".to_vec(), b"3".to_vec());
-				<StorageOf<Test>>::insert(&2, &b"world".to_vec(), b"4".to_vec());
+				AccountInfoOf::<Test>::insert(2, &AccountInfo {
+					trie_id: unique_id2.to_vec(),
+					current_mem_stored: 0,
+				});
+				child::put(&unique_id2[..], &b"hello".to_vec(), &b"3".to_vec());
+				child::put(&unique_id2[..], &b"world".to_vec(), &b"4".to_vec());
 			}
 
 			// Transfer funds from account 1 of such amount that after this transfer
@@ -240,15 +267,15 @@ fn account_removal_removes_storage() {
 			// Verify that all entries from account 1 is removed, while
 			// entries from account 2 is in place.
 			{
-				assert_eq!(<StorageOf<Test>>::get(&1, &b"foo".to_vec()), None);
-				assert_eq!(<StorageOf<Test>>::get(&1, &b"bar".to_vec()), None);
+				assert_eq!(child::get_raw(&unique_id1[..], &b"foo".to_vec()), None);
+				assert_eq!(child::get_raw(&unique_id1[..], &b"bar".to_vec()), None);
 
 				assert_eq!(
-					<StorageOf<Test>>::get(&2, &b"hello".to_vec()),
+					child::get(&unique_id2[..], &b"hello".to_vec()),
 					Some(b"3".to_vec())
 				);
 				assert_eq!(
-					<StorageOf<Test>>::get(&2, &b"world".to_vec()),
+					child::get(&unique_id2[..], &b"world".to_vec()),
 					Some(b"4".to_vec())
 				);
 			}