diff --git a/core/state-machine/src/backend.rs b/core/state-machine/src/backend.rs index e2f398ef7ccae..5fbda1aa32f45 100644 --- a/core/state-machine/src/backend.rs +++ b/core/state-machine/src/backend.rs @@ -246,10 +246,11 @@ impl error::Error for Void { fn description(&self) -> &str { "unreachable error" } } -/// In-memory backend. Fully recomputes tries on each commit but useful for -/// tests. +/// In-memory backend. Fully recomputes tries each time `as_trie_backend` is called but useful for +/// tests and proof checking. pub struct InMemory { inner: HashMap>, HashMap, Vec>>, + // This field is only needed for returning reference in `as_trie_backend`. trie: Option, H>>, _hasher: PhantomData, } @@ -467,7 +468,6 @@ impl Backend for InMemory { fn as_trie_backend(&mut self)-> Option<&TrieBackend> { let mut mdb = MemoryDB::default(); - let mut root = None; let mut new_child_roots = Vec::new(); let mut root_map = None; for (storage_key, map) in &self.inner { @@ -478,16 +478,15 @@ impl Backend for InMemory { root_map = Some(map); } } - // root handling - if let Some(map) = root_map.take() { - root = Some(insert_into_memory_db::( + let root = match root_map { + Some(map) => insert_into_memory_db::( &mut mdb, - map.clone().into_iter().chain(new_child_roots.into_iter()) - )?); - } - let root = match root { - Some(root) => root, - None => insert_into_memory_db::(&mut mdb, ::std::iter::empty())?, + map.clone().into_iter().chain(new_child_roots.into_iter()), + )?, + None => insert_into_memory_db::( + &mut mdb, + new_child_roots.into_iter(), + )?, }; self.trie = Some(TrieBackend::new(mdb, root)); self.trie.as_ref() @@ -513,3 +512,20 @@ pub(crate) fn insert_into_memory_db(mdb: &mut MemoryDB, input: I) -> Op Some(root) } + +#[cfg(test)] +mod tests { + use super::*; + + /// Assert in memory backend with only child trie keys works as trie backend. + #[test] + fn in_memory_with_child_trie_only() { + let storage = InMemory::::default(); + let mut storage = storage.update( + vec![(Some(b"1".to_vec()), b"2".to_vec(), Some(b"3".to_vec()))] + ); + let trie_backend = storage.as_trie_backend().unwrap(); + assert_eq!(trie_backend.child_storage(b"1", b"2").unwrap(), Some(b"3".to_vec())); + assert!(trie_backend.storage(b"1").unwrap().is_some()); + } +}