Make triecache generic and work for no-std environments (#14403)

* Remove `AsLocalTrieCache` trait

* Introduce new trait AsTrieDbCache

* Use AsTrieDbCache trait

* Make it compile

* Docs

* Improve naming of associated type, implement cache usage for no-std

* Improve naming

* Improve docs

* Allow construction with optional cache

* FMT

* Remove unused variable

* Revert unwanted change

* Apply suggestions from code review

Co-authored-by: Bastian Köcher <git@kchr.de>

* More comment adjustments

* Docs

* Trigger CI

* ".git/.scripts/commands/fmt/fmt.sh"

* Apply suggestions from code review

Co-authored-by: Koute <koute@users.noreply.github.com>

* Review comments

* Review comments

* Apply suggestions from code review

Co-authored-by: Michal Kucharczyk <1728078+michalkucharczyk@users.noreply.github.com>

* fmt

* Bump trie-db again

---------

Co-authored-by: Bastian Köcher <git@kchr.de>
Co-authored-by: command-bot <>
Co-authored-by: Koute <koute@users.noreply.github.com>
Co-authored-by: Michal Kucharczyk <1728078+michalkucharczyk@users.noreply.github.com>

substrate/primitives/state-machine/Cargo.toml

@@ -27,13 +27,13 @@ sp-externalities = { version = "0.19.0", default-features = false, path = "../ex
 sp-panic-handler = { version = "8.0.0", optional = true, path = "../panic-handler" }
 sp-std = { version = "8.0.0", default-features = false, path = "../std" }
 sp-trie = { version = "22.0.0", default-features = false, path = "../trie" }
+trie-db = { version = "0.27.1", default-features = false }
 
 [dev-dependencies]
 array-bytes = "4.1"
 pretty_assertions = "1.2.1"
 rand = "0.8.5"
 sp-runtime = { version = "24.0.0", path = "../runtime" }
-trie-db = "0.27.1"
 assert_matches = "1.5"
 
 [features]
@@ -49,6 +49,7 @@ std = [
 	"sp-panic-handler",
 	"sp-std/std",
 	"sp-trie/std",
+	"trie-db/std",
 	"thiserror",
 	"tracing",
 ]

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

@@ -36,6 +36,8 @@ mod testing;
 mod trie_backend;
 mod trie_backend_essence;
 
+pub use trie_backend::TrieCacheProvider;
+
 #[cfg(feature = "std")]
 pub use std_reexport::*;
 

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

@@ -24,61 +24,142 @@ use crate::{
 	trie_backend_essence::{RawIter, TrieBackendEssence, TrieBackendStorage},
 	Backend, StorageKey, StorageValue,
 };
+
 use codec::Codec;
 #[cfg(feature = "std")]
 use hash_db::HashDB;
 use hash_db::Hasher;
 use sp_core::storage::{ChildInfo, StateVersion};
 #[cfg(feature = "std")]
-use sp_trie::{cache::LocalTrieCache, recorder::Recorder};
+use sp_trie::{
-#[cfg(feature = "std")]
+	cache::{LocalTrieCache, TrieCache},
-use sp_trie::{MemoryDB, StorageProof};
+	recorder::Recorder,
-
+	MemoryDB, StorageProof,
-/// Dummy type to be used in `no_std`.
+};
-///
-/// This is required to have the type available for [`TrieBackendBuilder`] and [`TrieBackend`].
 #[cfg(not(feature = "std"))]
-pub struct LocalTrieCache<H>(sp_std::marker::PhantomData<H>);
+use sp_trie::{Error, NodeCodec};
+use trie_db::TrieCache as TrieCacheT;
+#[cfg(not(feature = "std"))]
+use trie_db::{node::NodeOwned, CachedValue};
 
-/// Special trait to support taking the [`LocalTrieCache`] by value or by reference.
+/// A provider of trie caches that are compatible with [`trie_db::TrieDB`].
-///
+pub trait TrieCacheProvider<H: Hasher> {
-/// This trait is internal use only and to emphasize this, the trait is sealed.
+	/// Cache type that implements [`trie_db::TrieCache`].
-pub trait AsLocalTrieCache<H: Hasher>: sealed::Sealed {
+	type Cache<'a>: TrieCacheT<sp_trie::NodeCodec<H>> + 'a
-	/// Returns `self` as [`LocalTrieCache`].
+	where
-	#[cfg(feature = "std")]
+		Self: 'a;
-	fn as_local_trie_cache(&self) -> &LocalTrieCache<H>;
+
+	/// Return a [`trie_db::TrieDB`] compatible cache.
+	///
+	/// The `storage_root` parameter should be the storage root of the used trie.
+	fn as_trie_db_cache(&self, storage_root: H::Out) -> Self::Cache<'_>;
+
+	/// Returns a cache that can be used with a [`trie_db::TrieDBMut`].
+	///
+	/// When finished with the operation on the trie, it is required to call [`Self::merge`] to
+	/// merge the cached items for the correct `storage_root`.
+	fn as_trie_db_mut_cache(&self) -> Self::Cache<'_>;
+
+	/// Merge the cached data in `other` into the provider using the given `new_root`.
+	///
+	/// This must be used for the cache returned by [`Self::as_trie_db_mut_cache`] as otherwise the
+	/// cached data is just thrown away.
+	fn merge<'a>(&'a self, other: Self::Cache<'a>, new_root: H::Out);
 }
 
-impl<H: Hasher> AsLocalTrieCache<H> for LocalTrieCache<H> {
+#[cfg(feature = "std")]
-	#[cfg(feature = "std")]
+impl<H: Hasher> TrieCacheProvider<H> for LocalTrieCache<H> {
-	#[inline]
+	type Cache<'a> = TrieCache<'a, H> where H: 'a;
-	fn as_local_trie_cache(&self) -> &LocalTrieCache<H> {
+
-		self
+	fn as_trie_db_cache(&self, storage_root: H::Out) -> Self::Cache<'_> {
+		self.as_trie_db_cache(storage_root)
+	}
+
+	fn as_trie_db_mut_cache(&self) -> Self::Cache<'_> {
+		self.as_trie_db_mut_cache()
+	}
+
+	fn merge<'a>(&'a self, other: Self::Cache<'a>, new_root: H::Out) {
+		other.merge_into(self, new_root)
 	}
 }
 
 #[cfg(feature = "std")]
-impl<H: Hasher> AsLocalTrieCache<H> for &LocalTrieCache<H> {
+impl<H: Hasher> TrieCacheProvider<H> for &LocalTrieCache<H> {
-	#[inline]
+	type Cache<'a> = TrieCache<'a, H> where Self: 'a;
-	fn as_local_trie_cache(&self) -> &LocalTrieCache<H> {
+
-		self
+	fn as_trie_db_cache(&self, storage_root: H::Out) -> Self::Cache<'_> {
+		(*self).as_trie_db_cache(storage_root)
+	}
+
+	fn as_trie_db_mut_cache(&self) -> Self::Cache<'_> {
+		(*self).as_trie_db_mut_cache()
+	}
+
+	fn merge<'a>(&'a self, other: Self::Cache<'a>, new_root: H::Out) {
+		other.merge_into(self, new_root)
 	}
 }
 
-/// Special module that contains the `Sealed` trait.
+/// Cache provider that allows construction of a [`TrieBackend`] and satisfies the requirements, but
-mod sealed {
+/// can never be instantiated.
-	use super::*;
+#[cfg(not(feature = "std"))]
-
+pub struct UnimplementedCacheProvider<H> {
-	/// A special trait which prevents externals to implement the [`AsLocalTrieCache`] outside
+	// Not strictly necessary, but the H bound allows to use this as a drop-in
-	/// of this crate.
+	// replacement for the `LocalTrieCache` in no-std contexts.
-	pub trait Sealed {}
+	_phantom: core::marker::PhantomData<H>,
-
+	// Statically prevents construction.
-	impl<H: Hasher> Sealed for LocalTrieCache<H> {}
+	_infallible: core::convert::Infallible,
-	impl<H: Hasher> Sealed for &LocalTrieCache<H> {}
 }
 
+#[cfg(not(feature = "std"))]
+impl<H: Hasher> trie_db::TrieCache<NodeCodec<H>> for UnimplementedCacheProvider<H> {
+	fn lookup_value_for_key(&mut self, _key: &[u8]) -> Option<&CachedValue<H::Out>> {
+		unimplemented!()
+	}
+
+	fn cache_value_for_key(&mut self, _key: &[u8], _value: CachedValue<H::Out>) {
+		unimplemented!()
+	}
+
+	fn get_or_insert_node(
+		&mut self,
+		_hash: H::Out,
+		_fetch_node: &mut dyn FnMut() -> trie_db::Result<NodeOwned<H::Out>, H::Out, Error<H::Out>>,
+	) -> trie_db::Result<&NodeOwned<H::Out>, H::Out, Error<H::Out>> {
+		unimplemented!()
+	}
+
+	fn get_node(&mut self, _hash: &H::Out) -> Option<&NodeOwned<H::Out>> {
+		unimplemented!()
+	}
+}
+
+#[cfg(not(feature = "std"))]
+impl<H: Hasher> TrieCacheProvider<H> for UnimplementedCacheProvider<H> {
+	type Cache<'a> = UnimplementedCacheProvider<H> where H: 'a;
+
+	fn as_trie_db_cache(&self, _storage_root: <H as Hasher>::Out) -> Self::Cache<'_> {
+		unimplemented!()
+	}
+
+	fn as_trie_db_mut_cache(&self) -> Self::Cache<'_> {
+		unimplemented!()
+	}
+
+	fn merge<'a>(&'a self, _other: Self::Cache<'a>, _new_root: <H as Hasher>::Out) {
+		unimplemented!()
+	}
+}
+
+#[cfg(feature = "std")]
+type DefaultCache<H> = LocalTrieCache<H>;
+
+#[cfg(not(feature = "std"))]
+type DefaultCache<H> = UnimplementedCacheProvider<H>;
+
 /// Builder for creating a [`TrieBackend`].
-pub struct TrieBackendBuilder<S: TrieBackendStorage<H>, H: Hasher, C = LocalTrieCache<H>> {
+pub struct TrieBackendBuilder<S: TrieBackendStorage<H>, H: Hasher, C = DefaultCache<H>> {
 	storage: S,
 	root: H::Out,
 	#[cfg(feature = "std")]
@@ -86,7 +167,7 @@ pub struct TrieBackendBuilder<S: TrieBackendStorage<H>, H: Hasher, C = LocalTrie
 	cache: Option<C>,
 }
 
-impl<S, H> TrieBackendBuilder<S, H, LocalTrieCache<H>>
+impl<S, H> TrieBackendBuilder<S, H, DefaultCache<H>>
 where
 	S: TrieBackendStorage<H>,
 	H: Hasher,
@@ -108,6 +189,16 @@ where
 	S: TrieBackendStorage<H>,
 	H: Hasher,
 {
+	/// Create a new builder instance.
+	pub fn new_with_cache(storage: S, root: H::Out, cache: C) -> Self {
+		Self {
+			storage,
+			root,
+			#[cfg(feature = "std")]
+			recorder: None,
+			cache: Some(cache),
+		}
+	}
 	/// Wrap the given [`TrieBackend`].
 	///
 	/// This can be used for example if all accesses to the trie should
@@ -121,10 +212,7 @@ where
 			root: *other.essence.root(),
 			#[cfg(feature = "std")]
 			recorder: None,
-			#[cfg(feature = "std")]
 			cache: other.essence.trie_node_cache.as_ref(),
-			#[cfg(not(feature = "std"))]
-			cache: None,
 		}
 	}
 
@@ -141,23 +229,23 @@ where
 	}
 
 	/// Use the given optional `cache` for the to be configured [`TrieBackend`].
-	#[cfg(feature = "std")]
 	pub fn with_optional_cache<LC>(self, cache: Option<LC>) -> TrieBackendBuilder<S, H, LC> {
 		TrieBackendBuilder {
 			cache,
 			root: self.root,
 			storage: self.storage,
+			#[cfg(feature = "std")]
 			recorder: self.recorder,
 		}
 	}
 
 	/// Use the given `cache` for the to be configured [`TrieBackend`].
-	#[cfg(feature = "std")]
 	pub fn with_cache<LC>(self, cache: LC) -> TrieBackendBuilder<S, H, LC> {
 		TrieBackendBuilder {
 			cache: Some(cache),
 			root: self.root,
 			storage: self.storage,
+			#[cfg(feature = "std")]
 			recorder: self.recorder,
 		}
 	}
@@ -179,10 +267,8 @@ where
 	/// Build the configured [`TrieBackend`].
 	#[cfg(not(feature = "std"))]
 	pub fn build(self) -> TrieBackend<S, H, C> {
-		let _ = self.cache;
-
 		TrieBackend {
-			essence: TrieBackendEssence::new(self.storage, self.root),
+			essence: TrieBackendEssence::new_with_cache(self.storage, self.root, self.cache),
 			next_storage_key_cache: Default::default(),
 		}
 	}
@@ -223,12 +309,13 @@ fn access_cache<T, R>(cell: &CacheCell<T>, callback: impl FnOnce(&mut T) -> R) -
 }
 
 /// Patricia trie-based backend. Transaction type is an overlay of changes to commit.
-pub struct TrieBackend<S: TrieBackendStorage<H>, H: Hasher, C = LocalTrieCache<H>> {
+pub struct TrieBackend<S: TrieBackendStorage<H>, H: Hasher, C = DefaultCache<H>> {
 	pub(crate) essence: TrieBackendEssence<S, H, C>,
 	next_storage_key_cache: CacheCell<Option<CachedIter<S, H, C>>>,
 }
 
-impl<S: TrieBackendStorage<H>, H: Hasher, C: AsLocalTrieCache<H> + Send + Sync> TrieBackend<S, H, C>
+impl<S: TrieBackendStorage<H>, H: Hasher, C: TrieCacheProvider<H> + Send + Sync>
+	TrieBackend<S, H, C>
 where
 	H::Out: Codec,
 {
@@ -276,7 +363,7 @@ where
 	}
 }
 
-impl<S: TrieBackendStorage<H>, H: Hasher, C: AsLocalTrieCache<H>> sp_std::fmt::Debug
+impl<S: TrieBackendStorage<H>, H: Hasher, C: TrieCacheProvider<H>> sp_std::fmt::Debug
 	for TrieBackend<S, H, C>
 {
 	fn fmt(&self, f: &mut sp_std::fmt::Formatter<'_>) -> sp_std::fmt::Result {
@@ -284,7 +371,7 @@ impl<S: TrieBackendStorage<H>, H: Hasher, C: AsLocalTrieCache<H>> sp_std::fmt::D
 	}
 }
 
-impl<S: TrieBackendStorage<H>, H: Hasher, C: AsLocalTrieCache<H> + Send + Sync> Backend<H>
+impl<S: TrieBackendStorage<H>, H: Hasher, C: TrieCacheProvider<H> + Send + Sync> Backend<H>
 	for TrieBackend<S, H, C>
 where
 	H::Out: Ord + Codec,

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

@@ -20,7 +20,7 @@
 
 use crate::{
 	backend::{Consolidate, IterArgs, StorageIterator},
-	trie_backend::AsLocalTrieCache,
+	trie_backend::TrieCacheProvider,
 	warn, StorageKey, StorageValue,
 };
 use codec::Codec;
@@ -39,7 +39,6 @@ use sp_trie::{
 };
 #[cfg(feature = "std")]
 use std::{collections::HashMap, sync::Arc};
-
 // In this module, we only use layout for read operation and empty root,
 // where V1 and V0 are equivalent.
 use sp_trie::LayoutV1 as Layout;
@@ -100,7 +99,7 @@ where
 	H: Hasher,
 	S: TrieBackendStorage<H>,
 	H::Out: Codec + Ord,
-	C: AsLocalTrieCache<H> + Send + Sync,
+	C: TrieCacheProvider<H> + Send + Sync,
 {
 	#[inline]
 	fn prepare<R>(
@@ -160,7 +159,7 @@ where
 	H: Hasher,
 	S: TrieBackendStorage<H>,
 	H::Out: Codec + Ord,
-	C: AsLocalTrieCache<H> + Send + Sync,
+	C: TrieCacheProvider<H> + Send + Sync,
 {
 	type Backend = crate::TrieBackend<S, H, C>;
 	type Error = crate::DefaultError;
@@ -209,29 +208,28 @@ pub struct TrieBackendEssence<S: TrieBackendStorage<H>, H: Hasher, C> {
 	empty: H::Out,
 	#[cfg(feature = "std")]
 	pub(crate) cache: Arc<RwLock<Cache<H::Out>>>,
-	#[cfg(feature = "std")]
 	pub(crate) trie_node_cache: Option<C>,
 	#[cfg(feature = "std")]
 	pub(crate) recorder: Option<Recorder<H>>,
-	#[cfg(not(feature = "std"))]
-	_phantom: PhantomData<C>,
 }
 
 impl<S: TrieBackendStorage<H>, H: Hasher, C> TrieBackendEssence<S, H, C> {
 	/// Create new trie-based backend.
 	pub fn new(storage: S, root: H::Out) -> Self {
+		Self::new_with_cache(storage, root, None)
+	}
+
+	/// Create new trie-based backend.
+	pub fn new_with_cache(storage: S, root: H::Out, cache: Option<C>) -> Self {
 		TrieBackendEssence {
 			storage,
 			root,
 			empty: H::hash(&[0u8]),
 			#[cfg(feature = "std")]
 			cache: Arc::new(RwLock::new(Cache::new())),
-			#[cfg(feature = "std")]
+			trie_node_cache: cache,
-			trie_node_cache: None,
 			#[cfg(feature = "std")]
 			recorder: None,
-			#[cfg(not(feature = "std"))]
-			_phantom: PhantomData,
 		}
 	}
 
@@ -289,11 +287,10 @@ impl<S: TrieBackendStorage<H>, H: Hasher, C> TrieBackendEssence<S, H, C> {
 	}
 }
 
-impl<S: TrieBackendStorage<H>, H: Hasher, C: AsLocalTrieCache<H>> TrieBackendEssence<S, H, C> {
+impl<S: TrieBackendStorage<H>, H: Hasher, C: TrieCacheProvider<H>> TrieBackendEssence<S, H, C> {
 	/// Call the given closure passing it the recorder and the cache.
 	///
 	/// If the given `storage_root` is `None`, `self.root` will be used.
-	#[cfg(feature = "std")]
 	#[inline]
 	fn with_recorder_and_cache<R>(
 		&self,
@@ -304,32 +301,23 @@ impl<S: TrieBackendStorage<H>, H: Hasher, C: AsLocalTrieCache<H>> TrieBackendEss
 		) -> R,
 	) -> R {
 		let storage_root = storage_root.unwrap_or_else(|| self.root);
-		let mut recorder = self.recorder.as_ref().map(|r| r.as_trie_recorder(storage_root));
+		let mut cache = self.trie_node_cache.as_ref().map(|c| c.as_trie_db_cache(storage_root));
-		let recorder = match recorder.as_mut() {
-			Some(recorder) => Some(recorder as &mut dyn TrieRecorder<H::Out>),
-			None => None,
-		};
-
-		let mut cache = self
-			.trie_node_cache
-			.as_ref()
-			.map(|c| c.as_local_trie_cache().as_trie_db_cache(storage_root));
 		let cache = cache.as_mut().map(|c| c as _);
 
-		callback(recorder, cache)
+		#[cfg(feature = "std")]
-	}
+		{
+			let mut recorder = self.recorder.as_ref().map(|r| r.as_trie_recorder(storage_root));
+			let recorder = match recorder.as_mut() {
+				Some(recorder) => Some(recorder as &mut dyn TrieRecorder<H::Out>),
+				None => None,
+			};
+			callback(recorder, cache)
+		}
 
-	#[cfg(not(feature = "std"))]
+		#[cfg(not(feature = "std"))]
-	#[inline]
+		{
-	fn with_recorder_and_cache<R>(
+			callback(None, cache)
-		&self,
+		}
-		_: Option<H::Out>,
-		callback: impl FnOnce(
-			Option<&mut dyn TrieRecorder<H::Out>>,
-			Option<&mut dyn TrieCache<NodeCodec<H>>>,
-		) -> R,
-	) -> R {
-		callback(None, None)
 	}
 
 	/// Call the given closure passing it the recorder and the cache.
@@ -356,12 +344,12 @@ impl<S: TrieBackendStorage<H>, H: Hasher, C: AsLocalTrieCache<H>> TrieBackendEss
 		};
 
 		let result = if let Some(local_cache) = self.trie_node_cache.as_ref() {
-			let mut cache = local_cache.as_local_trie_cache().as_trie_db_mut_cache();
+			let mut cache = local_cache.as_trie_db_mut_cache();
 
 			let (new_root, r) = callback(recorder, Some(&mut cache));
 
 			if let Some(new_root) = new_root {
-				cache.merge_into(local_cache.as_local_trie_cache(), new_root);
+				local_cache.merge(cache, new_root);
 			}
 
 			r
@@ -375,17 +363,29 @@ impl<S: TrieBackendStorage<H>, H: Hasher, C: AsLocalTrieCache<H>> TrieBackendEss
 	#[cfg(not(feature = "std"))]
 	fn with_recorder_and_cache_for_storage_root<R>(
 		&self,
-		_: Option<H::Out>,
+		_storage_root: Option<H::Out>,
 		callback: impl FnOnce(
 			Option<&mut dyn TrieRecorder<H::Out>>,
 			Option<&mut dyn TrieCache<NodeCodec<H>>>,
 		) -> (Option<H::Out>, R),
 	) -> R {
-		callback(None, None).1
+		if let Some(local_cache) = self.trie_node_cache.as_ref() {
+			let mut cache = local_cache.as_trie_db_mut_cache();
+
+			let (new_root, r) = callback(None, Some(&mut cache));
+
+			if let Some(new_root) = new_root {
+				local_cache.merge(cache, new_root);
+			}
+
+			r
+		} else {
+			callback(None, None).1
+		}
 	}
 }
 
-impl<S: TrieBackendStorage<H>, H: Hasher, C: AsLocalTrieCache<H> + Send + Sync>
+impl<S: TrieBackendStorage<H>, H: Hasher, C: TrieCacheProvider<H> + Send + Sync>
 	TrieBackendEssence<S, H, C>
 where
 	H::Out: Codec + Ord,
@@ -805,8 +805,8 @@ where
 	}
 }
 
-impl<S: TrieBackendStorage<H>, H: Hasher, C: AsLocalTrieCache<H> + Send + Sync> AsHashDB<H, DBValue>
+impl<S: TrieBackendStorage<H>, H: Hasher, C: TrieCacheProvider<H> + Send + Sync>
-	for TrieBackendEssence<S, H, C>
+	AsHashDB<H, DBValue> for TrieBackendEssence<S, H, C>
 {
 	fn as_hash_db<'b>(&'b self) -> &'b (dyn HashDB<H, DBValue> + 'b) {
 		self
@@ -816,7 +816,7 @@ impl<S: TrieBackendStorage<H>, H: Hasher, C: AsLocalTrieCache<H> + Send + Sync>
 	}
 }
 
-impl<S: TrieBackendStorage<H>, H: Hasher, C: AsLocalTrieCache<H> + Send + Sync> HashDB<H, DBValue>
+impl<S: TrieBackendStorage<H>, H: Hasher, C: TrieCacheProvider<H> + Send + Sync> HashDB<H, DBValue>
 	for TrieBackendEssence<S, H, C>
 {
 	fn get(&self, key: &H::Out, prefix: Prefix) -> Option<DBValue> {
@@ -849,7 +849,7 @@ impl<S: TrieBackendStorage<H>, H: Hasher, C: AsLocalTrieCache<H> + Send + Sync>
 	}
 }
 
-impl<S: TrieBackendStorage<H>, H: Hasher, C: AsLocalTrieCache<H> + Send + Sync>
+impl<S: TrieBackendStorage<H>, H: Hasher, C: TrieCacheProvider<H> + Send + Sync>
 	HashDBRef<H, DBValue> for TrieBackendEssence<S, H, C>
 {
 	fn get(&self, key: &H::Out, prefix: Prefix) -> Option<DBValue> {

substrate/primitives/trie/src/cache/mod.rs

@@ -529,7 +529,7 @@ impl<'a, H: Hasher> TrieCache<'a, H> {
 	/// `storage_root` is the new storage root that was obtained after finishing all operations
 	/// using the [`TrieDBMut`](trie_db::TrieDBMut).
 	pub fn merge_into(self, local: &LocalTrieCache<H>, storage_root: H::Out) {
-		let cache = if let ValueCache::Fresh(cache) = self.value_cache { cache } else { return };
+		let ValueCache::Fresh(cache) = self.value_cache else { return };
 
 		if !cache.is_empty() {
 			let mut value_cache = local.value_cache.lock();