state-machine: Move all functionality from trie backend to the essence (#10904)

* state-machine: Move all functionality from trie backend to the essence

This is required for some future changes of me and it also makes more sense to have all the
functionality inside the essence. Besides that it changes the child root cache to directly
cache the hash.

* Update primitives/state-machine/src/trie_backend_essence.rs

Co-authored-by: cheme <emericchevalier.pro@gmail.com>

* FMT

Co-authored-by: cheme <emericchevalier.pro@gmail.com>

substrate/primitives/state-machine/src/trie_backend.rs

@@ -18,19 +18,13 @@
 //! Trie-based state machine backend.
 
 use crate::{
-	debug,
-	trie_backend_essence::{Ephemeral, TrieBackendEssence, TrieBackendStorage},
-	warn, Backend, StorageKey, StorageValue,
+	trie_backend_essence::{TrieBackendEssence, TrieBackendStorage},
+	Backend, StorageKey, StorageValue,
 };
-use codec::{Codec, Decode};
+use codec::Codec;
 use hash_db::Hasher;
-use sp_core::storage::{ChildInfo, ChildType, StateVersion};
-use sp_std::{boxed::Box, vec::Vec};
-use sp_trie::{
-	child_delta_trie_root, delta_trie_root, empty_child_trie_root,
-	trie_types::{TrieDB, TrieError},
-	LayoutV0, LayoutV1, Trie,
-};
+use sp_core::storage::{ChildInfo, StateVersion};
+use sp_std::vec::Vec;
 
 /// Patricia trie-based backend. Transaction type is an overlay of changes to commit.
 pub struct TrieBackend<S: TrieBackendStorage<H>, H: Hasher> {
@@ -144,43 +138,11 @@ where
 	}
 
 	fn pairs(&self) -> Vec<(StorageKey, StorageValue)> {
-		let collect_all = || -> Result<_, Box<TrieError<H::Out>>> {
-			let trie = TrieDB::<H>::new(self.essence(), self.essence.root())?;
-			let mut v = Vec::new();
-			for x in trie.iter()? {
-				let (key, value) = x?;
-				v.push((key.to_vec(), value.to_vec()));
-			}
-
-			Ok(v)
-		};
-
-		match collect_all() {
-			Ok(v) => v,
-			Err(e) => {
-				debug!(target: "trie", "Error extracting trie values: {}", e);
-				Vec::new()
-			},
-		}
+		self.essence.pairs()
 	}
 
 	fn keys(&self, prefix: &[u8]) -> Vec<StorageKey> {
-		let collect_all = || -> Result<_, Box<TrieError<H::Out>>> {
-			let trie = TrieDB::<H>::new(self.essence(), self.essence.root())?;
-			let mut v = Vec::new();
-			for x in trie.iter()? {
-				let (key, _) = x?;
-				if key.starts_with(prefix) {
-					v.push(key.to_vec());
-				}
-			}
-
-			Ok(v)
-		};
-
-		collect_all()
-			.map_err(|e| debug!(target: "trie", "Error extracting trie keys: {}", e))
-			.unwrap_or_default()
+		self.essence.keys(prefix)
 	}
 
 	fn storage_root<'a>(
@@ -191,25 +153,7 @@ where
 	where
 		H::Out: Ord,
 	{
-		let mut write_overlay = S::Overlay::default();
-		let mut root = *self.essence.root();
-
-		{
-			let mut eph = Ephemeral::new(self.essence.backend_storage(), &mut write_overlay);
-			let res = match state_version {
-				StateVersion::V0 =>
-					delta_trie_root::<LayoutV0<H>, _, _, _, _, _>(&mut eph, root, delta),
-				StateVersion::V1 =>
-					delta_trie_root::<LayoutV1<H>, _, _, _, _, _>(&mut eph, root, delta),
-			};
-
-			match res {
-				Ok(ret) => root = ret,
-				Err(e) => warn!(target: "trie", "Failed to write to trie: {}", e),
-			}
-		}
-
-		(root, write_overlay)
+		self.essence.storage_root(delta, state_version)
 	}
 
 	fn child_storage_root<'a>(
@@ -221,45 +165,7 @@ where
 	where
 		H::Out: Ord,
 	{
-		let default_root = match child_info.child_type() {
-			ChildType::ParentKeyId => empty_child_trie_root::<LayoutV1<H>>(),
-		};
-		let mut write_overlay = S::Overlay::default();
-		let prefixed_storage_key = child_info.prefixed_storage_key();
-		let mut root = match self.storage(prefixed_storage_key.as_slice()) {
-			Ok(value) => value
-				.and_then(|r| Decode::decode(&mut &r[..]).ok())
-				.unwrap_or_else(|| default_root.clone()),
-			Err(e) => {
-				warn!(target: "trie", "Failed to read child storage root: {}", e);
-				default_root.clone()
-			},
-		};
-
-		{
-			let mut eph = Ephemeral::new(self.essence.backend_storage(), &mut write_overlay);
-			match match state_version {
-				StateVersion::V0 => child_delta_trie_root::<LayoutV0<H>, _, _, _, _, _, _>(
-					child_info.keyspace(),
-					&mut eph,
-					root,
-					delta,
-				),
-				StateVersion::V1 => child_delta_trie_root::<LayoutV1<H>, _, _, _, _, _, _>(
-					child_info.keyspace(),
-					&mut eph,
-					root,
-					delta,
-				),
-			} {
-				Ok(ret) => root = ret,
-				Err(e) => warn!(target: "trie", "Failed to write to trie: {}", e),
-			}
-		}
-
-		let is_default = root == default_root;
-
-		(root, is_default, write_overlay)
+		self.essence.child_storage_root(child_info, delta, state_version)
 	}
 
 	fn as_trie_backend(&self) -> Option<&TrieBackend<Self::TrieBackendStorage, H>> {

substrate/primitives/state-machine/src/trie_backend_essence.rs

@@ -23,10 +23,11 @@ use codec::Encode;
 use hash_db::{self, AsHashDB, HashDB, HashDBRef, Hasher, Prefix};
 #[cfg(feature = "std")]
 use parking_lot::RwLock;
-use sp_core::storage::ChildInfo;
+use sp_core::storage::{ChildInfo, ChildType, StateVersion};
 use sp_std::{boxed::Box, vec::Vec};
 use sp_trie::{
-	empty_child_trie_root, read_child_trie_value, read_trie_value,
+	child_delta_trie_root, delta_trie_root, empty_child_trie_root, read_child_trie_value,
+	read_trie_value,
 	trie_types::{TrieDB, TrieError},
 	DBValue, KeySpacedDB, PrefixedMemoryDB, Trie, TrieDBIterator, TrieDBKeyIterator,
 };
@@ -58,12 +59,12 @@ pub trait Storage<H: Hasher>: Send + Sync {
 
 /// Local cache for child root.
 #[cfg(feature = "std")]
-pub(crate) struct Cache {
-	pub child_root: HashMap<Vec<u8>, Option<Vec<u8>>>,
+pub(crate) struct Cache<H> {
+	pub child_root: HashMap<Vec<u8>, Option<H>>,
 }
 
 #[cfg(feature = "std")]
-impl Cache {
+impl<H> Cache<H> {
 	fn new() -> Self {
 		Cache { child_root: HashMap::new() }
 	}
@@ -75,7 +76,7 @@ pub struct TrieBackendEssence<S: TrieBackendStorage<H>, H: Hasher> {
 	root: H::Out,
 	empty: H::Out,
 	#[cfg(feature = "std")]
-	pub(crate) cache: Arc<RwLock<Cache>>,
+	pub(crate) cache: Arc<RwLock<Cache<H::Out>>>,
 }
 
 impl<S: TrieBackendStorage<H>, H: Hasher> TrieBackendEssence<S, H>
@@ -130,22 +131,26 @@ where
 	}
 
 	/// Access the root of the child storage in its parent trie
-	fn child_root(&self, child_info: &ChildInfo) -> Result<Option<StorageValue>> {
+	fn child_root(&self, child_info: &ChildInfo) -> Result<Option<H::Out>> {
 		#[cfg(feature = "std")]
 		{
 			if let Some(result) = self.cache.read().child_root.get(child_info.storage_key()) {
-				return Ok(result.clone())
+				return Ok(*result)
 			}
 		}
 
-		let result = self.storage(child_info.prefixed_storage_key().as_slice())?;
+		let result = self.storage(child_info.prefixed_storage_key().as_slice())?.map(|r| {
+			let mut hash = H::Out::default();
+
+			// root is fetched from DB, not writable by runtime, so it's always valid.
+			hash.as_mut().copy_from_slice(&r[..]);
+
+			hash
+		});
 
 		#[cfg(feature = "std")]
 		{
-			self.cache
-				.write()
-				.child_root
-				.insert(child_info.storage_key().to_vec(), result.clone());
+			self.cache.write().child_root.insert(child_info.storage_key().to_vec(), result);
 		}
 
 		Ok(result)
@@ -163,15 +168,7 @@ where
 			None => return Ok(None),
 		};
 
-		let mut hash = H::Out::default();
-
-		if child_root.len() != hash.as_ref().len() {
-			return Err(format!("Invalid child storage hash at {:?}", child_info.storage_key()))
-		}
-		// note: child_root and hash must be same size, panics otherwise.
-		hash.as_mut().copy_from_slice(&child_root[..]);
-
-		self.next_storage_key_from_root(&hash, Some(child_info), key)
+		self.next_storage_key_from_root(&child_root, Some(child_info), key)
 	}
 
 	/// Return next key from main trie or child trie by providing corresponding root.
@@ -231,9 +228,10 @@ where
 		child_info: &ChildInfo,
 		key: &[u8],
 	) -> Result<Option<StorageValue>> {
-		let root = self
-			.child_root(child_info)?
-			.unwrap_or_else(|| empty_child_trie_root::<Layout<H>>().encode());
+		let root = match self.child_root(child_info)? {
+			Some(root) => root,
+			None => return Ok(None),
+		};
 
 		let map_e = |e| format!("Trie lookup error: {}", e);
 
@@ -253,19 +251,13 @@ where
 		f: impl FnMut(Vec<u8>, Vec<u8>) -> bool,
 		allow_missing_nodes: bool,
 	) -> Result<bool> {
-		let mut child_root;
 		let root = if let Some(child_info) = child_info.as_ref() {
-			if let Some(fetched_child_root) = self.child_root(child_info)? {
-				child_root = H::Out::default();
-				// root is fetched from DB, not writable by runtime, so it's always valid.
-				child_root.as_mut().copy_from_slice(fetched_child_root.as_slice());
-
-				&child_root
-			} else {
-				return Ok(true)
+			match self.child_root(child_info)? {
+				Some(child_root) => child_root,
+				None => return Ok(true),
 			}
 		} else {
-			&self.root
+			self.root
 		};
 
 		self.trie_iter_inner(&root, prefix, f, child_info, start_at, allow_missing_nodes)
@@ -279,22 +271,21 @@ where
 		prefix: Option<&[u8]>,
 		mut f: F,
 	) {
-		let mut child_root = H::Out::default();
 		let root = if let Some(child_info) = child_info.as_ref() {
-			let root_vec = match self.child_root(child_info) {
-				Ok(v) => v.unwrap_or_else(|| empty_child_trie_root::<Layout<H>>().encode()),
+			match self.child_root(child_info) {
+				Ok(Some(v)) => v,
+				// If the child trie doesn't exist, there is no need to continue.
+				Ok(None) => return,
 				Err(e) => {
 					debug!(target: "trie", "Error while iterating child storage: {}", e);
 					return
 				},
-			};
-			child_root.as_mut().copy_from_slice(&root_vec);
-			&child_root
+			}
 		} else {
-			&self.root
+			self.root
 		};
 
-		self.trie_iter_key_inner(root, prefix, |k| f(k), child_info)
+		self.trie_iter_key_inner(&root, prefix, |k| f(k), child_info)
 	}
 
 	/// Execute given closure for all keys starting with prefix.
@@ -304,15 +295,16 @@ where
 		prefix: &[u8],
 		mut f: impl FnMut(&[u8]),
 	) {
-		let root_vec = match self.child_root(child_info) {
-			Ok(v) => v.unwrap_or_else(|| empty_child_trie_root::<Layout<H>>().encode()),
+		let root = match self.child_root(child_info) {
+			Ok(Some(v)) => v,
+			// If the child trie doesn't exist, there is no need to continue.
+			Ok(None) => return,
 			Err(e) => {
 				debug!(target: "trie", "Error while iterating child storage: {}", e);
 				return
 			},
 		};
-		let mut root = H::Out::default();
-		root.as_mut().copy_from_slice(&root_vec);
+
 		self.trie_iter_key_inner(
 			&root,
 			Some(prefix),
@@ -438,6 +430,130 @@ where
 			false,
 		);
 	}
+
+	/// Returns all `(key, value)` pairs in the trie.
+	pub fn pairs(&self) -> Vec<(StorageKey, StorageValue)> {
+		let collect_all = || -> sp_std::result::Result<_, Box<TrieError<H::Out>>> {
+			let trie = TrieDB::<H>::new(self, &self.root)?;
+			let mut v = Vec::new();
+			for x in trie.iter()? {
+				let (key, value) = x?;
+				v.push((key.to_vec(), value.to_vec()));
+			}
+
+			Ok(v)
+		};
+
+		match collect_all() {
+			Ok(v) => v,
+			Err(e) => {
+				debug!(target: "trie", "Error extracting trie values: {}", e);
+				Vec::new()
+			},
+		}
+	}
+
+	/// Returns all keys that start with the given `prefix`.
+	pub fn keys(&self, prefix: &[u8]) -> Vec<StorageKey> {
+		let collect_all = || -> sp_std::result::Result<_, Box<TrieError<H::Out>>> {
+			let trie = TrieDB::<H>::new(self, &self.root)?;
+			let mut v = Vec::new();
+			for x in trie.iter()? {
+				let (key, _) = x?;
+				if key.starts_with(prefix) {
+					v.push(key.to_vec());
+				}
+			}
+
+			Ok(v)
+		};
+
+		collect_all()
+			.map_err(|e| debug!(target: "trie", "Error extracting trie keys: {}", e))
+			.unwrap_or_default()
+	}
+
+	/// Return the storage root after applying the given `delta`.
+	pub fn storage_root<'a>(
+		&self,
+		delta: impl Iterator<Item = (&'a [u8], Option<&'a [u8]>)>,
+		state_version: StateVersion,
+	) -> (H::Out, S::Overlay)
+	where
+		H::Out: Ord,
+	{
+		let mut write_overlay = S::Overlay::default();
+		let mut root = self.root;
+
+		{
+			let mut eph = Ephemeral::new(self.backend_storage(), &mut write_overlay);
+			let res = match state_version {
+				StateVersion::V0 =>
+					delta_trie_root::<sp_trie::LayoutV0<H>, _, _, _, _, _>(&mut eph, root, delta),
+				StateVersion::V1 =>
+					delta_trie_root::<sp_trie::LayoutV1<H>, _, _, _, _, _>(&mut eph, root, delta),
+			};
+
+			match res {
+				Ok(ret) => root = ret,
+				Err(e) => warn!(target: "trie", "Failed to write to trie: {}", e),
+			}
+		}
+
+		(root, write_overlay)
+	}
+
+	/// Returns the child storage root for the child trie `child_info` after applying the given
+	/// `delta`.
+	pub fn child_storage_root<'a>(
+		&self,
+		child_info: &ChildInfo,
+		delta: impl Iterator<Item = (&'a [u8], Option<&'a [u8]>)>,
+		state_version: StateVersion,
+	) -> (H::Out, bool, S::Overlay)
+	where
+		H::Out: Ord,
+	{
+		let default_root = match child_info.child_type() {
+			ChildType::ParentKeyId => empty_child_trie_root::<sp_trie::LayoutV1<H>>(),
+		};
+		let mut write_overlay = S::Overlay::default();
+		let mut root = match self.child_root(child_info) {
+			Ok(Some(hash)) => hash,
+			Ok(None) => default_root,
+			Err(e) => {
+				warn!(target: "trie", "Failed to read child storage root: {}", e);
+				default_root.clone()
+			},
+		};
+
+		{
+			let mut eph = Ephemeral::new(self.backend_storage(), &mut write_overlay);
+			match match state_version {
+				StateVersion::V0 =>
+					child_delta_trie_root::<sp_trie::LayoutV0<H>, _, _, _, _, _, _>(
+						child_info.keyspace(),
+						&mut eph,
+						root,
+						delta,
+					),
+				StateVersion::V1 =>
+					child_delta_trie_root::<sp_trie::LayoutV1<H>, _, _, _, _, _, _>(
+						child_info.keyspace(),
+						&mut eph,
+						root,
+						delta,
+					),
+			} {
+				Ok(ret) => root = ret,
+				Err(e) => warn!(target: "trie", "Failed to write to trie: {}", e),
+			}
+		}
+
+		let is_default = root == default_root;
+
+		(root, is_default, write_overlay)
+	}
 }
 
 pub(crate) struct Ephemeral<'a, S: 'a + TrieBackendStorage<H>, H: 'a + Hasher> {