diff --git a/primitives/state-machine/src/in_memory_backend.rs b/primitives/state-machine/src/in_memory_backend.rs index f211f60202730..ca300aec919c4 100644 --- a/primitives/state-machine/src/in_memory_backend.rs +++ b/primitives/state-machine/src/in_memory_backend.rs @@ -17,56 +17,21 @@ //! State machine in memory backend. use crate::{ - StorageKey, StorageValue, StorageCollection, - trie_backend::TrieBackend, + StorageKey, StorageValue, StorageCollection, trie_backend::TrieBackend, backend::Backend, }; -use std::{collections::{BTreeMap, HashMap}}; +use std::collections::{BTreeMap, HashMap}; use hash_db::Hasher; -use sp_trie::{ - MemoryDB, TrieMut, - trie_types::TrieDBMut, -}; +use sp_trie::{MemoryDB, empty_trie_root, Layout}; use codec::Codec; use sp_core::storage::{ChildInfo, Storage}; -/// Insert input pairs into memory db. -fn insert_into_memory_db(mut root: H::Out, mdb: &mut MemoryDB, input: I) -> H::Out -where - H: Hasher, - I: IntoIterator)>, -{ - { - let mut trie = if root == Default::default() { - TrieDBMut::::new(mdb, &mut root) - } else { - TrieDBMut::::from_existing(mdb, &mut root).unwrap() - }; - for (key, value) in input { - if let Err(e) = match value { - Some(value) => { - trie.insert(&key, &value) - }, - None => { - trie.remove(&key) - }, - } { - panic!("Failed to write to trie: {}", e); - } - } - trie.commit(); - } - root -} - /// Create a new empty instance of in-memory backend. pub fn new_in_mem() -> TrieBackend, H> where H::Out: Codec + Ord, { let db = MemoryDB::default(); - let mut backend = TrieBackend::new(db, Default::default()); - backend.insert(std::iter::empty()); - backend + TrieBackend::new(db, empty_trie_root::>()) } impl TrieBackend, H> @@ -92,32 +57,16 @@ where &mut self, changes: T, ) { - let mut new_child_roots = Vec::new(); - let mut root_map = None; - let root = self.root().clone(); - for (child_info, map) in changes { - if let Some(child_info) = child_info.as_ref() { - let prefix_storage_key = child_info.prefixed_storage_key(); - let ch = insert_into_memory_db::(root, self.backend_storage_mut(), map.clone().into_iter()); - new_child_roots.push((prefix_storage_key.into_inner(), Some(ch.as_ref().into()))); - } else { - root_map = Some(map); - } - } + let (top, child) = changes.into_iter().partition::, _>(|v| v.0.is_none()); + let (root, transaction) = self.full_storage_root( + top.iter().map(|(_, v)| v).flatten().map(|(k, v)| (&k[..], v.as_deref())), + child.iter() + .filter_map(|v| + v.0.as_ref().map(|c| (c, v.1.iter().map(|(k, v)| (&k[..], v.as_deref())))) + ), + ); - let root = match root_map { - Some(map) => insert_into_memory_db::( - root, - self.backend_storage_mut(), - map.into_iter().chain(new_child_roots.into_iter()), - ), - None => insert_into_memory_db::( - root, - self.backend_storage_mut(), - new_child_roots.into_iter(), - ), - }; - self.essence.set_root(root); + self.apply_transaction(root, transaction); } /// Merge trie nodes into this backend. @@ -127,6 +76,12 @@ where Self::new(clone, root) } + /// Apply the given transaction to this backend and set the root to the given value. + pub fn apply_transaction(&mut self, root: H::Out, transaction: MemoryDB) { + self.backend_storage_mut().consolidate(transaction); + self.essence.set_root(root); + } + /// Compare with another in-memory backend. pub fn eq(&self, other: &Self) -> bool { self.root() == other.root() @@ -158,7 +113,9 @@ where { fn from(inner: HashMap, BTreeMap>) -> Self { let mut backend = new_in_mem(); - backend.insert(inner.into_iter().map(|(k, m)| (k, m.into_iter().map(|(k, v)| (k, Some(v))).collect()))); + backend.insert( + inner.into_iter().map(|(k, m)| (k, m.into_iter().map(|(k, v)| (k, Some(v))).collect())), + ); backend } } @@ -232,4 +189,16 @@ mod tests { let storage_key = child_info.prefixed_storage_key(); assert!(trie_backend.storage(storage_key.as_slice()).unwrap().is_some()); } + + #[test] + fn insert_multiple_times_child_data_works() { + let mut storage = new_in_mem::(); + let child_info = ChildInfo::new_default(b"1"); + + storage.insert(vec![(Some(child_info.clone()), vec![(b"2".to_vec(), Some(b"3".to_vec()))])]); + storage.insert(vec![(Some(child_info.clone()), vec![(b"1".to_vec(), Some(b"3".to_vec()))])]); + + assert_eq!(storage.child_storage(&child_info, &b"2"[..]), Ok(Some(b"3".to_vec()))); + assert_eq!(storage.child_storage(&child_info, &b"1"[..]), Ok(Some(b"3".to_vec()))); + } } diff --git a/primitives/state-machine/src/testing.rs b/primitives/state-machine/src/testing.rs index 4dcd308285625..23c3abe4910c5 100644 --- a/primitives/state-machine/src/testing.rs +++ b/primitives/state-machine/src/testing.rs @@ -153,8 +153,11 @@ impl TestExternalities &mut self.changes_trie_storage } - /// Return a new backend with all pending value. - pub fn commit_all(&self) -> InMemoryBackend { + /// Return a new backend with all pending changes. + /// + /// In contrast to [`commit_all`](Self::commit_all) this will not panic if there are open + /// transactions. + fn as_backend(&self) -> InMemoryBackend { let top: Vec<_> = self.overlay.changes() .map(|(k, v)| (k.clone(), v.value().cloned())) .collect(); @@ -172,6 +175,23 @@ impl TestExternalities self.backend.update(transaction) } + /// Commit all pending changes to the underlying backend. + /// + /// # Panic + /// + /// This will panic if there are still open transactions. + pub fn commit_all(&mut self) -> Result<(), String> { + let changes = self.overlay.drain_storage_changes::<_, _, N>( + &self.backend, + None, + Default::default(), + &mut Default::default(), + )?; + + self.backend.apply_transaction(changes.transaction_storage_root, changes.transaction); + Ok(()) + } + /// Execute the given closure while `self` is set as externalities. /// /// Returns the result of the given closure. @@ -209,7 +229,7 @@ impl PartialEq for TestExternalities /// This doesn't test if they are in the same state, only if they contains the /// same data at this state fn eq(&self, other: &TestExternalities) -> bool { - self.commit_all().eq(&other.commit_all()) + self.as_backend().eq(&other.as_backend()) } } @@ -258,7 +278,7 @@ impl sp_externalities::ExtensionStore for TestExternalities where #[cfg(test)] mod tests { use super::*; - use sp_core::{H256, traits::Externalities}; + use sp_core::{H256, traits::Externalities, storage::ChildInfo}; use sp_runtime::traits::BlakeTwo256; use hex_literal::hex; @@ -289,4 +309,45 @@ mod tests { fn assert_send() {} assert_send::>(); } + + #[test] + fn commit_all_and_kill_child_storage() { + let mut ext = TestExternalities::::default(); + let child_info = ChildInfo::new_default(&b"test_child"[..]); + + { + let mut ext = ext.ext(); + ext.place_child_storage(&child_info, b"doe".to_vec(), Some(b"reindeer".to_vec())); + ext.place_child_storage(&child_info, b"dog".to_vec(), Some(b"puppy".to_vec())); + ext.place_child_storage(&child_info, b"dog2".to_vec(), Some(b"puppy2".to_vec())); + } + + ext.commit_all().unwrap(); + + { + let mut ext = ext.ext(); + + assert!(!ext.kill_child_storage(&child_info, Some(2)), "Should not delete all keys"); + + assert!(ext.child_storage(&child_info, &b"doe"[..]).is_none()); + assert!(ext.child_storage(&child_info, &b"dog"[..]).is_none()); + assert!(ext.child_storage(&child_info, &b"dog2"[..]).is_some()); + } + } + + #[test] + fn as_backend_generates_same_backend_as_commit_all() { + let mut ext = TestExternalities::::default(); + { + let mut ext = ext.ext(); + ext.set_storage(b"doe".to_vec(), b"reindeer".to_vec()); + ext.set_storage(b"dog".to_vec(), b"puppy".to_vec()); + ext.set_storage(b"dogglesworth".to_vec(), b"cat".to_vec()); + } + + let backend = ext.as_backend(); + + ext.commit_all().unwrap(); + assert!(ext.backend.eq(&backend), "Both backend should be equal."); + } } diff --git a/primitives/trie/src/lib.rs b/primitives/trie/src/lib.rs index 2687d8e422796..4914d85f5811f 100644 --- a/primitives/trie/src/lib.rs +++ b/primitives/trie/src/lib.rs @@ -184,7 +184,7 @@ pub fn delta_trie_root( DB: hash_db::HashDB, { { - let mut trie = TrieDBMut::::from_existing(&mut *db, &mut root)?; + let mut trie = TrieDBMut::::from_existing(db, &mut root)?; let mut delta = delta.into_iter().collect::>(); delta.sort_by(|l, r| l.0.borrow().cmp(r.0.borrow())); @@ -223,9 +223,13 @@ pub fn read_trie_value_with< Ok(TrieDB::::new(&*db, root)?.get_with(key, query).map(|x| x.map(|val| val.to_vec()))?) } +/// Determine the empty trie root. +pub fn empty_trie_root() -> ::Out { + L::trie_root::<_, Vec, Vec>(core::iter::empty()) +} + /// Determine the empty child trie root. -pub fn empty_child_trie_root( -) -> ::Out { +pub fn empty_child_trie_root() -> ::Out { L::trie_root::<_, Vec, Vec>(core::iter::empty()) }