feat: Rebrand Polkadot/Substrate references to PezkuwiChain

This commit systematically rebrands various references from Parity Technologies'
Polkadot/Substrate ecosystem to PezkuwiChain within the kurdistan-sdk.

Key changes include:
- Updated external repository URLs (zombienet-sdk, parity-db, parity-scale-codec, wasm-instrument) to point to pezkuwichain forks.
- Modified internal documentation and code comments to reflect PezkuwiChain naming and structure.
- Replaced direct references to  with  or specific paths within the  for XCM, Pezkuwi, and other modules.
- Cleaned up deprecated  issue and PR references in various  and  files, particularly in  and  modules.
- Adjusted image and logo URLs in documentation to point to PezkuwiChain assets.
- Removed or rephrased comments related to external Polkadot/Substrate PRs and issues.

This is a significant step towards fully customizing the SDK for the PezkuwiChain ecosystem.

bizinikiwi/primitives/trie/Cargo.toml

@@ -0,0 +1,72 @@
+[package]
+name = "pezsp-trie"
+version = "29.0.0"
+authors.workspace = true
+description = "Patricia trie stuff using a parity-scale-codec node format"
+repository.workspace = true
+license = "Apache-2.0"
+edition.workspace = true
+homepage.workspace = true
+documentation = "https://docs.rs/pezsp-trie"
+readme = "README.md"
+
+[lints]
+workspace = true
+
+[package.metadata.docs.rs]
+targets = ["x86_64-unknown-linux-gnu"]
+
+[[bench]]
+name = "bench"
+harness = false
+
+[dependencies]
+ahash = { optional = true, workspace = true }
+codec = { workspace = true }
+foldhash = { workspace = true }
+hash-db = { workspace = true }
+hashbrown = { workspace = true }
+memory-db = { workspace = true }
+nohash-hasher = { optional = true, workspace = true }
+parking_lot = { optional = true, workspace = true, default-features = true }
+prometheus-endpoint = { optional = true, workspace = true, default-features = true }
+rand = { optional = true, workspace = true, default-features = true }
+scale-info = { features = ["derive"], workspace = true }
+schnellru = { optional = true, workspace = true }
+pezsp-core = { workspace = true }
+pezsp-externalities = { workspace = true }
+thiserror = { optional = true, workspace = true }
+tracing = { optional = true, workspace = true, default-features = true }
+trie-db = { workspace = true }
+trie-root = { workspace = true }
+
+[dev-dependencies]
+array-bytes = { workspace = true, default-features = true }
+criterion = { workspace = true, default-features = true }
+pezsp-runtime = { workspace = true, default-features = true }
+trie-bench = { workspace = true }
+trie-standardmap = { workspace = true }
+
+[features]
+default = ["std"]
+std = [
+	"ahash",
+	"codec/std",
+	"foldhash/std",
+	"hash-db/std",
+	"memory-db/std",
+	"nohash-hasher",
+	"parking_lot",
+	"prometheus-endpoint",
+	"rand",
+	"scale-info/std",
+	"schnellru",
+	"pezsp-core/std",
+	"pezsp-externalities/std",
+	"pezsp-runtime/std",
+	"thiserror",
+	"tracing",
+	"trie-db/std",
+	"trie-root/std",
+]
+runtime-benchmarks = ["pezsp-runtime/runtime-benchmarks"]

bizinikiwi/primitives/trie/README.md

@@ -0,0 +1,3 @@
+Utility functions to interact with Bizinikiwi's Base-16 Modified Merkle Patricia tree ("trie").
+
+License: Apache-2.0

bizinikiwi/primitives/trie/benches/bench.rs

@@ -0,0 +1,31 @@
+// This file is part of Bizinikiwi.
+
+// Copyright (C) Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use criterion::{criterion_group, criterion_main, Criterion};
+criterion_group!(benches, benchmark);
+criterion_main!(benches);
+
+fn benchmark(c: &mut Criterion) {
+	trie_bench::standard_benchmark::<
+		pezsp_trie::LayoutV1<pezsp_runtime::traits::BlakeTwo256>,
+		pezsp_trie::TrieStream,
+	>(c, "bizinikiwi-blake2");
+	trie_bench::standard_benchmark::<
+		pezsp_trie::LayoutV1<pezsp_runtime::traits::BlakeTwo256>,
+		pezsp_trie::TrieStream,
+	>(c, "bizinikiwi-keccak");
+}

bizinikiwi/primitives/trie/src/accessed_nodes_tracker.rs

@@ -0,0 +1,119 @@
+// This file is part of Bizinikiwi.
+
+// Copyright (C) Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Helpers for checking for duplicate nodes.
+
+use alloc::collections::BTreeSet;
+use core::hash::Hash;
+use scale_info::TypeInfo;
+use pezsp_core::{Decode, Encode};
+use trie_db::{RecordedForKey, TrieAccess, TrieRecorder};
+
+/// Error associated with the `AccessedNodesTracker` module.
+#[derive(Encode, Decode, Clone, Eq, PartialEq, Debug, TypeInfo)]
+pub enum Error {
+	/// The proof contains unused nodes.
+	UnusedNodes,
+}
+
+/// Helper struct used to ensure that a storage proof doesn't contain duplicate or unused nodes.
+///
+/// The struct needs to be used as a `TrieRecorder` and `ensure_no_unused_nodes()` has to be called
+/// to actually perform the check.
+pub struct AccessedNodesTracker<H: Hash> {
+	proof_nodes_count: usize,
+	recorder: BTreeSet<H>,
+}
+
+impl<H: Hash> AccessedNodesTracker<H> {
+	/// Create a new instance of `RedundantNodesChecker`, starting from a `RawStorageProof`.
+	pub fn new(proof_nodes_count: usize) -> Self {
+		Self { proof_nodes_count, recorder: BTreeSet::new() }
+	}
+
+	/// Ensure that all the nodes in the proof have been accessed.
+	pub fn ensure_no_unused_nodes(self) -> Result<(), Error> {
+		if self.proof_nodes_count != self.recorder.len() {
+			return Err(Error::UnusedNodes);
+		}
+
+		Ok(())
+	}
+}
+
+impl<H: Hash + Ord> TrieRecorder<H> for AccessedNodesTracker<H> {
+	fn record(&mut self, access: TrieAccess<H>) {
+		match access {
+			TrieAccess::NodeOwned { hash, .. } |
+			TrieAccess::EncodedNode { hash, .. } |
+			TrieAccess::Value { hash, .. } => {
+				self.recorder.insert(hash);
+			},
+			_ => {},
+		}
+	}
+
+	fn trie_nodes_recorded_for_key(&self, _key: &[u8]) -> RecordedForKey {
+		RecordedForKey::None
+	}
+}
+
+#[cfg(test)]
+pub mod tests {
+	use super::*;
+	use crate::{tests::create_storage_proof, StorageProof};
+	use hash_db::Hasher;
+	use trie_db::{Trie, TrieDBBuilder};
+
+	type Hash = <pezsp_core::Blake2Hasher as Hasher>::Out;
+	type Layout = crate::LayoutV1<pezsp_core::Blake2Hasher>;
+
+	const TEST_DATA: &[(&[u8], &[u8])] =
+		&[(b"key1", &[1; 64]), (b"key2", &[2; 64]), (b"key3", &[3; 64])];
+
+	#[test]
+	fn proof_with_unused_nodes_is_rejected() {
+		let (raw_proof, root) = create_storage_proof::<Layout>(TEST_DATA);
+		let proof = StorageProof::new(raw_proof.clone());
+		let proof_nodes_count = proof.len();
+
+		let mut accessed_nodes_tracker = AccessedNodesTracker::<Hash>::new(proof_nodes_count);
+		{
+			let db = proof.clone().into_memory_db();
+			let trie = TrieDBBuilder::<Layout>::new(&db, &root)
+				.with_recorder(&mut accessed_nodes_tracker)
+				.build();
+
+			trie.get(b"key1").unwrap().unwrap();
+			trie.get(b"key2").unwrap().unwrap();
+			trie.get(b"key3").unwrap().unwrap();
+		}
+		assert_eq!(accessed_nodes_tracker.ensure_no_unused_nodes(), Ok(()));
+
+		let mut accessed_nodes_tracker = AccessedNodesTracker::<Hash>::new(proof_nodes_count);
+		{
+			let db = proof.into_memory_db();
+			let trie = TrieDBBuilder::<Layout>::new(&db, &root)
+				.with_recorder(&mut accessed_nodes_tracker)
+				.build();
+
+			trie.get(b"key1").unwrap().unwrap();
+			trie.get(b"key2").unwrap().unwrap();
+		}
+		assert_eq!(accessed_nodes_tracker.ensure_no_unused_nodes(), Err(Error::UnusedNodes));
+	}
+}

bizinikiwi/primitives/trie/src/cache/metrics.rs

@@ -0,0 +1,269 @@
+// This file is part of Bizinikiwi.
+
+// Copyright (C) Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Metrics for the trie cache.
+
+use prometheus_endpoint::{
+	exponential_buckets,
+	prometheus::{core::Collector, HistogramTimer},
+	CounterVec, GaugeVec, HistogramOpts, HistogramVec, Opts, PrometheusError, Registry, U64,
+};
+
+// Register a metric with the given registry.
+fn register<T: Clone + Collector + 'static>(
+	metric: T,
+	registry: &Registry,
+) -> Result<T, PrometheusError> {
+	registry.register(Box::new(metric.clone()))?;
+	Ok(metric)
+}
+
+/// Metrics for the trie cache.
+/// This struct is used to track the performance of the trie cache.
+/// It contains histograms and counters for the shared and local caches.
+#[derive(Clone)]
+pub struct Metrics {
+	// The duration in seconds to update the shared trie caches from local to shared cache.
+	shared_update_duration: HistogramVec,
+	// Number of attempts hitting the shared trie caches.
+	shared_hits: CounterVec<U64>,
+	// Number of attempts to the shared trie caches.
+	shared_fetch_attempts: CounterVec<U64>,
+	// Number of attempts hitting the local trie caches.
+	local_hits: CounterVec<U64>,
+	// Number of attempts to the local caches.
+	local_fetch_attempts: CounterVec<U64>,
+	// Length of the local caches.
+	local_cache_lengths: HistogramVec,
+	// The inline size of the shared caches.
+	shared_cache_inline_size: GaugeVec<U64>,
+	// The heap size of the shared caches.
+	shared_cache_heap_size: GaugeVec<U64>,
+}
+
+impl Metrics {
+	/// Create a new instance of the metrics.
+	pub(crate) fn register(registry: &Registry) -> Result<Self, PrometheusError> {
+		Ok(Self {
+			shared_update_duration: register(
+				HistogramVec::new(
+					HistogramOpts {
+						common_opts: Opts::new(
+							"trie_cache_shared_update_duration",
+							"Duration in seconds to update the shared trie caches from local cache to shared cache",
+						),
+						buckets: exponential_buckets(0.001, 4.0, 9)
+							.expect("function parameters are constant and always valid; qed"),
+					},
+					&["cache_type"], // node or value
+				)?,
+				registry,
+			)?,
+			shared_hits: register(
+				CounterVec::new(
+					Opts::new(
+						"trie_cache_shared_hits",
+						"Number of attempts hitting the shared trie cache",
+					),
+					&["cache_type"], // node or value
+				)?,
+				registry,
+			)?,
+			shared_fetch_attempts: register(
+				CounterVec::new(
+					Opts::new(
+						"trie_cache_shared_fetch_attempts",
+						"Number of attempts to the shared trie cache",
+					),
+					&["cache_type"],
+				)?,
+				registry,
+			)?,
+			local_hits: register(
+				CounterVec::new(
+					Opts::new(
+						"trie_cache_local_hits",
+						"Number of attempts hitting the local trie cache",
+					),
+					&["cache_type"],
+				)?,
+				registry,
+			)?,
+			local_fetch_attempts: register(
+				CounterVec::new(
+					Opts::new(
+						"trie_cache_local_fetch_attempts",
+						"Number of attempts to the local cache",
+					),
+					&["cache_type"],
+				)?,
+				registry,
+			)?,
+			local_cache_lengths: register(
+				HistogramVec::new(
+					HistogramOpts {
+						common_opts: Opts::new(
+							"trie_cache_local_cache_lengths",
+							"Histogram of length of the local cache",
+						),
+						buckets: exponential_buckets(1.0, 4.0, 9)
+							.expect("function parameters are constant and always valid; qed"),
+					},
+					&["cache_type"],
+				)?,
+				registry,
+			)?,
+			shared_cache_inline_size: register(
+				GaugeVec::new(
+					Opts::new(
+						"trie_cache_shared_cache_inline_size",
+						"The inline size of the shared caches",
+					),
+					&["cache_type"],
+				)?,
+				registry,
+			)?,
+			shared_cache_heap_size: register(
+				GaugeVec::new(
+					Opts::new(
+						"trie_cache_shared_cache_heap_size",
+						"The heap size of the shared caches",
+					),
+					&["cache_type"],
+				)?,
+				registry,
+			)?,
+		})
+	}
+
+	/// Start a timer for the shared node cache update duration.
+	pub(crate) fn start_shared_node_update_timer(&self) -> HistogramTimer {
+		self.shared_update_duration.with_label_values(&["node"]).start_timer()
+	}
+
+	/// Start a timer for the shared value cache update duration.
+	pub(crate) fn start_shared_value_update_timer(&self) -> HistogramTimer {
+		self.shared_update_duration.with_label_values(&["value"]).start_timer()
+	}
+
+	/// Observe the shared node cache length.
+	pub(crate) fn observe_local_node_cache_length(&self, node_cache_len: usize) {
+		self.local_cache_lengths
+			.with_label_values(&["node"])
+			.observe(node_cache_len as f64);
+	}
+
+	/// Observe the shared value cache length.
+	pub(crate) fn observe_local_value_cache_length(&self, value_cache_len: usize) {
+		self.local_cache_lengths
+			.with_label_values(&["value"])
+			.observe(value_cache_len as f64);
+	}
+
+	/// Observe the shared node cache inline size.
+	pub(crate) fn observe_node_cache_inline_size(&self, cache_size: usize) {
+		self.shared_cache_inline_size
+			.with_label_values(&["node"])
+			.set(cache_size as u64);
+	}
+
+	/// Observe the shared value cache inline size.
+	pub(crate) fn observe_value_cache_inline_size(&self, cache_size: usize) {
+		self.shared_cache_inline_size
+			.with_label_values(&["value"])
+			.set(cache_size as u64);
+	}
+
+	/// Observe the shared node cache heap size.
+	pub(crate) fn observe_node_cache_heap_size(&self, cache_size: usize) {
+		self.shared_cache_heap_size.with_label_values(&["node"]).set(cache_size as u64);
+	}
+
+	/// Observe the shared value cache heap size.
+	pub(crate) fn observe_value_cache_heap_size(&self, cache_size: usize) {
+		self.shared_cache_heap_size.with_label_values(&["value"]).set(cache_size as u64);
+	}
+
+	/// Observe the hit stats from an instance of a local cache.
+	pub(crate) fn observe_hits_stats(&self, stats: &TrieHitStatsSnapshot) {
+		self.shared_hits
+			.with_label_values(&["node"])
+			.inc_by(stats.node_cache.shared_hits);
+		self.shared_fetch_attempts
+			.with_label_values(&["node"])
+			.inc_by(stats.node_cache.shared_fetch_attempts);
+		self.local_hits.with_label_values(&["node"]).inc_by(stats.node_cache.local_hits);
+		self.local_fetch_attempts
+			.with_label_values(&["node"])
+			.inc_by(stats.node_cache.local_fetch_attempts);
+
+		self.shared_hits
+			.with_label_values(&["value"])
+			.inc_by(stats.value_cache.shared_hits);
+		self.shared_fetch_attempts
+			.with_label_values(&["value"])
+			.inc_by(stats.value_cache.shared_fetch_attempts);
+		self.local_hits
+			.with_label_values(&["value"])
+			.inc_by(stats.value_cache.local_hits);
+		self.local_fetch_attempts
+			.with_label_values(&["value"])
+			.inc_by(stats.value_cache.local_fetch_attempts);
+	}
+}
+
+/// A snapshot of the hit/miss stats.
+#[derive(Default, Copy, Clone, Debug)]
+pub(crate) struct HitStatsSnapshot {
+	pub(crate) shared_hits: u64,
+	pub(crate) shared_fetch_attempts: u64,
+	pub(crate) local_hits: u64,
+	pub(crate) local_fetch_attempts: u64,
+}
+
+impl std::fmt::Display for HitStatsSnapshot {
+	fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
+		let shared_hits = self.shared_hits;
+		let shared_fetch_attempts = self.shared_fetch_attempts;
+		let local_hits = self.local_hits;
+		let local_fetch_attempts = self.local_fetch_attempts;
+
+		if shared_fetch_attempts == 0 && local_hits == 0 {
+			write!(fmt, "empty")
+		} else {
+			let percent_local = (local_hits as f32 / local_fetch_attempts as f32) * 100.0;
+			let percent_shared = (shared_hits as f32 / shared_fetch_attempts as f32) * 100.0;
+			write!(
+				fmt,
+				"local hit rate = {}% [{}/{}], shared hit rate = {}% [{}/{}]",
+				percent_local as u32,
+				local_hits,
+				local_fetch_attempts,
+				percent_shared as u32,
+				shared_hits,
+				shared_fetch_attempts
+			)
+		}
+	}
+}
+
+/// Snapshot of the hit/miss stats for the node cache and the value cache.
+#[derive(Default, Debug, Clone, Copy)]
+pub(crate) struct TrieHitStatsSnapshot {
+	pub(crate) node_cache: HitStatsSnapshot,
+	pub(crate) value_cache: HitStatsSnapshot,
+}

bizinikiwi/primitives/trie/src/cache/shared_cache.rs

@@ -0,0 +1,984 @@
+// This file is part of Bizinikiwi.
+
+// Copyright (C) Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+///! Provides the [`SharedNodeCache`], the [`SharedValueCache`] and the [`SharedTrieCache`]
+///! that combines both caches and is exported to the outside.
+use super::{
+	metrics::Metrics, CacheSize, LocalNodeCacheConfig, LocalNodeCacheLimiter,
+	LocalValueCacheConfig, LocalValueCacheLimiter, NodeCached, TrieHitStats, TrieHitStatsSnapshot,
+};
+use crate::cache::LOG_TARGET;
+use core::{hash::Hash, time::Duration};
+use hash_db::Hasher;
+use nohash_hasher::BuildNoHashHasher;
+use parking_lot::{Mutex, RwLock, RwLockWriteGuard};
+use prometheus_endpoint::Registry;
+use schnellru::LruMap;
+use std::{
+	collections::{hash_map::Entry as SetEntry, HashMap},
+	hash::{BuildHasher, Hasher as _},
+	sync::{Arc, LazyLock},
+	time::Instant,
+};
+use trie_db::{node::NodeOwned, CachedValue};
+static RANDOM_STATE: LazyLock<ahash::RandomState> = LazyLock::new(|| {
+	use rand::Rng;
+	let mut rng = rand::thread_rng();
+	ahash::RandomState::generate_with(rng.gen(), rng.gen(), rng.gen(), rng.gen())
+});
+
+pub struct SharedNodeCacheLimiter {
+	/// The maximum size (in bytes) the cache can hold inline.
+	///
+	/// This space is always consumed whether there are any items in the map or not.
+	max_inline_size: usize,
+
+	/// The maximum size (in bytes) the cache can hold on the heap.
+	max_heap_size: usize,
+
+	/// The current size (in bytes) of data allocated by this cache on the heap.
+	///
+	/// This doesn't include the size of the map itself.
+	heap_size: usize,
+
+	/// A counter with the number of elements that got evicted from the cache.
+	///
+	/// Reset to zero before every update.
+	items_evicted: usize,
+
+	/// The maximum number of elements that we allow to be evicted.
+	///
+	/// Reset on every update.
+	max_items_evicted: usize,
+}
+
+impl<H> schnellru::Limiter<H, NodeOwned<H>> for SharedNodeCacheLimiter
+where
+	H: AsRef<[u8]>,
+{
+	type KeyToInsert<'a> = H;
+	type LinkType = u32;
+
+	#[inline]
+	fn is_over_the_limit(&self, _length: usize) -> bool {
+		// Once we hit the limit of max items evicted this will return `false` and prevent
+		// any further evictions, but this is fine because the outer loop which inserts
+		// items into this cache will just detect this and stop inserting new items.
+		self.items_evicted <= self.max_items_evicted && self.heap_size > self.max_heap_size
+	}
+
+	#[inline]
+	fn on_insert(
+		&mut self,
+		_length: usize,
+		key: Self::KeyToInsert<'_>,
+		node: NodeOwned<H>,
+	) -> Option<(H, NodeOwned<H>)> {
+		let new_item_heap_size = node.size_in_bytes() - std::mem::size_of::<NodeOwned<H>>();
+		if new_item_heap_size > self.max_heap_size {
+			// Item's too big to add even if the cache's empty; bail.
+			return None;
+		}
+
+		self.heap_size += new_item_heap_size;
+		Some((key, node))
+	}
+
+	#[inline]
+	fn on_replace(
+		&mut self,
+		_length: usize,
+		_old_key: &mut H,
+		_new_key: H,
+		old_node: &mut NodeOwned<H>,
+		new_node: &mut NodeOwned<H>,
+	) -> bool {
+		debug_assert_eq!(_old_key.as_ref(), _new_key.as_ref());
+
+		let new_item_heap_size = new_node.size_in_bytes() - std::mem::size_of::<NodeOwned<H>>();
+		if new_item_heap_size > self.max_heap_size {
+			// Item's too big to add even if the cache's empty; bail.
+			return false;
+		}
+
+		let old_item_heap_size = old_node.size_in_bytes() - std::mem::size_of::<NodeOwned<H>>();
+		self.heap_size = self.heap_size - old_item_heap_size + new_item_heap_size;
+		true
+	}
+
+	#[inline]
+	fn on_cleared(&mut self) {
+		self.heap_size = 0;
+	}
+
+	#[inline]
+	fn on_removed(&mut self, _: &mut H, node: &mut NodeOwned<H>) {
+		self.heap_size -= node.size_in_bytes() - std::mem::size_of::<NodeOwned<H>>();
+		self.items_evicted += 1;
+	}
+
+	#[inline]
+	fn on_grow(&mut self, new_memory_usage: usize) -> bool {
+		new_memory_usage <= self.max_inline_size
+	}
+}
+
+pub struct SharedValueCacheLimiter {
+	/// The maximum size (in bytes) the cache can hold inline.
+	///
+	/// This space is always consumed whether there are any items in the map or not.
+	max_inline_size: usize,
+
+	/// The maximum size (in bytes) the cache can hold on the heap.
+	max_heap_size: usize,
+
+	/// The current size (in bytes) of data allocated by this cache on the heap.
+	///
+	/// This doesn't include the size of the map itself.
+	heap_size: usize,
+
+	/// A set with all of the keys deduplicated to save on memory.
+	known_storage_keys: HashMap<Arc<[u8]>, (), ahash::RandomState>,
+
+	/// A counter with the number of elements that got evicted from the cache.
+	///
+	/// Reset to zero before every update.
+	items_evicted: usize,
+
+	/// The maximum number of elements that we allow to be evicted.
+	///
+	/// Reset on every update.
+	max_items_evicted: usize,
+}
+
+impl<H> schnellru::Limiter<ValueCacheKey<H>, CachedValue<H>> for SharedValueCacheLimiter
+where
+	H: AsRef<[u8]>,
+{
+	type KeyToInsert<'a> = ValueCacheKey<H>;
+	type LinkType = u32;
+
+	#[inline]
+	fn is_over_the_limit(&self, _length: usize) -> bool {
+		self.items_evicted <= self.max_items_evicted && self.heap_size > self.max_heap_size
+	}
+
+	#[inline]
+	fn on_insert(
+		&mut self,
+		_length: usize,
+		mut key: Self::KeyToInsert<'_>,
+		value: CachedValue<H>,
+	) -> Option<(ValueCacheKey<H>, CachedValue<H>)> {
+		match self.known_storage_keys.entry(key.storage_key.clone()) {
+			SetEntry::Vacant(entry) => {
+				let new_item_heap_size = key.storage_key.len();
+				if new_item_heap_size > self.max_heap_size {
+					// Item's too big to add even if the cache's empty; bail.
+					return None;
+				}
+
+				self.heap_size += new_item_heap_size;
+				entry.insert(());
+			},
+			SetEntry::Occupied(entry) => {
+				key.storage_key = entry.key().clone();
+			},
+		}
+
+		Some((key, value))
+	}
+
+	#[inline]
+	fn on_replace(
+		&mut self,
+		_length: usize,
+		_old_key: &mut ValueCacheKey<H>,
+		_new_key: ValueCacheKey<H>,
+		_old_value: &mut CachedValue<H>,
+		_new_value: &mut CachedValue<H>,
+	) -> bool {
+		debug_assert_eq!(_new_key.storage_key, _old_key.storage_key);
+		true
+	}
+
+	#[inline]
+	fn on_removed(&mut self, key: &mut ValueCacheKey<H>, _: &mut CachedValue<H>) {
+		if Arc::strong_count(&key.storage_key) == 2 {
+			// There are only two instances of this key:
+			//   1) one memoized in `known_storage_keys`,
+			//   2) one inside the map.
+			//
+			// This means that after this remove goes through the `Arc` will be deallocated.
+			self.heap_size -= key.storage_key.len();
+			self.known_storage_keys.remove(&key.storage_key);
+		}
+		self.items_evicted += 1;
+	}
+
+	#[inline]
+	fn on_cleared(&mut self) {
+		self.heap_size = 0;
+		self.known_storage_keys.clear();
+	}
+
+	#[inline]
+	fn on_grow(&mut self, new_memory_usage: usize) -> bool {
+		new_memory_usage <= self.max_inline_size
+	}
+}
+
+type SharedNodeCacheMap<H> =
+	LruMap<H, NodeOwned<H>, SharedNodeCacheLimiter, schnellru::RandomState>;
+
+/// The shared node cache.
+///
+/// Internally this stores all cached nodes in a [`LruMap`]. It ensures that when updating the
+/// cache, that the cache stays within its allowed bounds.
+pub(super) struct SharedNodeCache<H>
+where
+	H: AsRef<[u8]>,
+{
+	/// The cached nodes, ordered by least recently used.
+	pub(super) lru: SharedNodeCacheMap<H>,
+}
+
+impl<H: AsRef<[u8]> + Eq + std::hash::Hash> SharedNodeCache<H> {
+	/// Create a new instance.
+	fn new(max_inline_size: usize, max_heap_size: usize) -> Self {
+		Self {
+			lru: LruMap::new(SharedNodeCacheLimiter {
+				max_inline_size,
+				max_heap_size,
+				heap_size: 0,
+				items_evicted: 0,
+				max_items_evicted: 0, // Will be set during `update`.
+			}),
+		}
+	}
+
+	/// Update the cache with the `list` of nodes which were either newly added or accessed.
+	pub fn update(
+		&mut self,
+		list: impl IntoIterator<Item = (H, NodeCached<H>)>,
+		config: &LocalNodeCacheConfig,
+		metrics: &Option<Metrics>,
+	) {
+		let mut access_count = 0;
+		let mut add_count = 0;
+
+		self.lru.limiter_mut().items_evicted = 0;
+		self.lru.limiter_mut().max_items_evicted =
+			self.lru.len() * 100 / config.shared_node_cache_max_replace_percent;
+
+		for (key, cached_node) in list {
+			if cached_node.is_from_shared_cache {
+				if self.lru.get(&key).is_some() {
+					access_count += 1;
+
+					if access_count >= config.shared_node_cache_max_promoted_keys {
+						// Stop when we've promoted a large enough number of items.
+						break;
+					}
+
+					continue;
+				}
+			}
+
+			self.lru.insert(key, cached_node.node);
+			add_count += 1;
+
+			if self.lru.limiter().items_evicted > self.lru.limiter().max_items_evicted {
+				// Stop when we've evicted a big enough chunk of the shared cache.
+				break;
+			}
+		}
+
+		metrics.as_ref().map(|m| {
+			m.observe_node_cache_inline_size(self.lru.memory_usage());
+			m.observe_node_cache_heap_size(self.lru.limiter().heap_size);
+		});
+
+		tracing::debug!(
+			target: super::LOG_TARGET,
+			"Updated the shared node cache: {} accesses, {} new values, {}/{} evicted (length = {}, inline size={}/{}, heap size={}/{})",
+			access_count,
+			add_count,
+			self.lru.limiter().items_evicted,
+			self.lru.limiter().max_items_evicted,
+			self.lru.len(),
+			self.lru.memory_usage(),
+			self.lru.limiter().max_inline_size,
+			self.lru.limiter().heap_size,
+			self.lru.limiter().max_heap_size,
+		);
+	}
+
+	/// Reset the cache.
+	fn reset(&mut self) {
+		self.lru.clear();
+	}
+}
+
+/// The hash of [`ValueCacheKey`].
+#[derive(PartialEq, Eq, Clone, Copy, Hash)]
+#[repr(transparent)]
+pub struct ValueCacheKeyHash(u64);
+
+impl ValueCacheKeyHash {
+	pub fn raw(self) -> u64 {
+		self.0
+	}
+}
+
+impl ValueCacheKeyHash {
+	pub fn from_hasher_and_storage_key(
+		mut hasher: impl std::hash::Hasher,
+		storage_key: &[u8],
+	) -> Self {
+		hasher.write(storage_key);
+
+		Self(hasher.finish())
+	}
+}
+
+impl nohash_hasher::IsEnabled for ValueCacheKeyHash {}
+
+/// The key type that is being used to address a [`CachedValue`].
+#[derive(Eq)]
+pub(super) struct ValueCacheKey<H> {
+	/// The storage root of the trie this key belongs to.
+	pub storage_root: H,
+	/// The key to access the value in the storage.
+	pub storage_key: Arc<[u8]>,
+	/// The hash that identifies this instance of `storage_root` and `storage_key`.
+	pub hash: ValueCacheKeyHash,
+}
+
+/// A borrowed variant of [`ValueCacheKey`].
+pub(super) struct ValueCacheRef<'a, H> {
+	/// The storage root of the trie this key belongs to.
+	pub storage_root: H,
+	/// The key to access the value in the storage.
+	pub storage_key: &'a [u8],
+	/// The hash that identifies this instance of `storage_root` and `storage_key`.
+	pub hash: ValueCacheKeyHash,
+}
+
+impl<'a, H> ValueCacheRef<'a, H> {
+	pub fn new(storage_key: &'a [u8], storage_root: H) -> Self
+	where
+		H: AsRef<[u8]>,
+	{
+		let hash = ValueCacheKey::<H>::hash_data(&storage_key, &storage_root);
+		Self { storage_root, storage_key, hash }
+	}
+}
+
+impl<'a, H> From<ValueCacheRef<'a, H>> for ValueCacheKey<H> {
+	fn from(value: ValueCacheRef<'a, H>) -> Self {
+		ValueCacheKey {
+			storage_root: value.storage_root,
+			storage_key: value.storage_key.into(),
+			hash: value.hash,
+		}
+	}
+}
+
+impl<'a, H: std::hash::Hash> std::hash::Hash for ValueCacheRef<'a, H> {
+	fn hash<Hasher: std::hash::Hasher>(&self, state: &mut Hasher) {
+		self.hash.hash(state)
+	}
+}
+
+impl<'a, H> PartialEq<ValueCacheKey<H>> for ValueCacheRef<'a, H>
+where
+	H: AsRef<[u8]>,
+{
+	fn eq(&self, rhs: &ValueCacheKey<H>) -> bool {
+		self.storage_root.as_ref() == rhs.storage_root.as_ref() &&
+			self.storage_key == &*rhs.storage_key
+	}
+}
+
+impl<H> ValueCacheKey<H> {
+	/// Constructs [`Self::Value`].
+	#[cfg(test)] // Only used in tests.
+	pub fn new_value(storage_key: impl Into<Arc<[u8]>>, storage_root: H) -> Self
+	where
+		H: AsRef<[u8]>,
+	{
+		let storage_key = storage_key.into();
+		let hash = Self::hash_data(&storage_key, &storage_root);
+		Self { storage_root, storage_key, hash }
+	}
+
+	/// Returns a hasher prepared to build the final hash to identify [`Self`].
+	///
+	/// See [`Self::hash_data`] for building the hash directly.
+	pub fn hash_partial_data(storage_root: &H) -> impl std::hash::Hasher + Clone
+	where
+		H: AsRef<[u8]>,
+	{
+		let mut hasher = RANDOM_STATE.build_hasher();
+		hasher.write(storage_root.as_ref());
+		hasher
+	}
+
+	/// Hash the `key` and `storage_root` that identify [`Self`].
+	///
+	/// Returns a `u64` which represents the unique hash for the given inputs.
+	pub fn hash_data(key: &[u8], storage_root: &H) -> ValueCacheKeyHash
+	where
+		H: AsRef<[u8]>,
+	{
+		let hasher = Self::hash_partial_data(storage_root);
+
+		ValueCacheKeyHash::from_hasher_and_storage_key(hasher, key)
+	}
+
+	/// Checks whether the key is equal to the given `storage_key` and `storage_root`.
+	#[inline]
+	pub fn is_eq(&self, storage_root: &H, storage_key: &[u8]) -> bool
+	where
+		H: PartialEq,
+	{
+		self.storage_root == *storage_root && *self.storage_key == *storage_key
+	}
+}
+
+// Implement manually so that only `hash` is accessed.
+impl<H: std::hash::Hash> std::hash::Hash for ValueCacheKey<H> {
+	fn hash<Hasher: std::hash::Hasher>(&self, state: &mut Hasher) {
+		self.hash.hash(state)
+	}
+}
+
+impl<H> nohash_hasher::IsEnabled for ValueCacheKey<H> {}
+
+// Implement manually to not have to compare `hash`.
+impl<H: PartialEq> PartialEq for ValueCacheKey<H> {
+	#[inline]
+	fn eq(&self, other: &Self) -> bool {
+		self.is_eq(&other.storage_root, &other.storage_key)
+	}
+}
+
+type SharedValueCacheMap<H> = schnellru::LruMap<
+	ValueCacheKey<H>,
+	CachedValue<H>,
+	SharedValueCacheLimiter,
+	BuildNoHashHasher<ValueCacheKey<H>>,
+>;
+
+/// The shared value cache.
+///
+/// The cache ensures that it stays in the configured size bounds.
+pub(super) struct SharedValueCache<H>
+where
+	H: AsRef<[u8]>,
+{
+	/// The cached nodes, ordered by least recently used.
+	pub(super) lru: SharedValueCacheMap<H>,
+}
+
+impl<H: Eq + std::hash::Hash + Clone + Copy + AsRef<[u8]>> SharedValueCache<H> {
+	/// Create a new instance.
+	fn new(max_inline_size: usize, max_heap_size: usize) -> Self {
+		Self {
+			lru: schnellru::LruMap::with_hasher(
+				SharedValueCacheLimiter {
+					max_inline_size,
+					max_heap_size,
+					heap_size: 0,
+					known_storage_keys: HashMap::with_hasher(RANDOM_STATE.clone()),
+					items_evicted: 0,
+					max_items_evicted: 0, // Will be set during `update`.
+				},
+				Default::default(),
+			),
+		}
+	}
+
+	/// Update the cache with the `added` values and the `accessed` values.
+	///
+	/// The `added` values are the ones that have been collected by doing operations on the trie and
+	/// now should be stored in the shared cache. The `accessed` values are only referenced by the
+	/// [`ValueCacheKeyHash`] and represent the values that were retrieved from this shared cache.
+	/// These `accessed` values are being put to the front of the internal [`LruMap`] like the
+	/// `added` ones.
+	pub fn update(
+		&mut self,
+		added: impl IntoIterator<Item = (ValueCacheKey<H>, CachedValue<H>)>,
+		accessed: impl IntoIterator<Item = ValueCacheKeyHash>,
+		config: &LocalValueCacheConfig,
+		metrics: &Option<Metrics>,
+	) {
+		let mut access_count = 0;
+		let mut add_count = 0;
+
+		for hash in accessed {
+			// Access every node in the map to put it to the front.
+			//
+			// Since we are only comparing the hashes here it may lead us to promoting the wrong
+			// values as the most recently accessed ones. However this is harmless as the only
+			// consequence is that we may accidentally prune a recently used value too early.
+			self.lru.get_by_hash(hash.raw(), |existing_key, _| existing_key.hash == hash);
+			access_count += 1;
+		}
+
+		// Insert all of the new items which were *not* found in the shared cache.
+		//
+		// Limit how many items we'll replace in the shared cache in one go so that
+		// we don't evict the whole shared cache nor we keep spinning our wheels
+		// evicting items which we've added ourselves in previous iterations of this loop.
+
+		self.lru.limiter_mut().items_evicted = 0;
+		self.lru.limiter_mut().max_items_evicted =
+			self.lru.len() * 100 / config.shared_value_cache_max_replace_percent;
+
+		for (key, value) in added {
+			self.lru.insert(key, value);
+			add_count += 1;
+
+			if self.lru.limiter().items_evicted > self.lru.limiter().max_items_evicted {
+				// Stop when we've evicted a big enough chunk of the shared cache.
+				break;
+			}
+		}
+
+		metrics.as_ref().map(|m| {
+			m.observe_value_cache_inline_size(self.lru.memory_usage());
+			m.observe_value_cache_heap_size(self.lru.limiter().heap_size);
+		});
+
+		tracing::debug!(
+			target: super::LOG_TARGET,
+			"Updated the shared value cache: {} accesses, {} new values, {}/{} evicted (length = {}, known_storage_keys = {}, inline size={}/{}, heap size={}/{})",
+			access_count,
+			add_count,
+			self.lru.limiter().items_evicted,
+			self.lru.limiter().max_items_evicted,
+			self.lru.len(),
+			self.lru.limiter().known_storage_keys.len(),
+			self.lru.memory_usage(),
+			self.lru.limiter().max_inline_size,
+			self.lru.limiter().heap_size,
+			self.lru.limiter().max_heap_size
+		);
+	}
+
+	/// Reset the cache.
+	fn reset(&mut self) {
+		self.lru.clear();
+	}
+}
+
+/// The inner of [`SharedTrieCache`].
+pub(super) struct SharedTrieCacheInner<H: Hasher> {
+	node_cache: SharedNodeCache<H::Out>,
+	value_cache: SharedValueCache<H::Out>,
+	stats: TrieHitStats,
+	metrics: Option<Metrics>,
+	previous_stats_dump: Instant,
+}
+
+impl<H: Hasher> SharedTrieCacheInner<H> {
+	/// Returns a reference to the [`SharedValueCache`].
+	#[cfg(test)]
+	pub(super) fn value_cache(&self) -> &SharedValueCache<H::Out> {
+		&self.value_cache
+	}
+
+	/// Returns a mutable reference to the [`SharedValueCache`].
+	pub(super) fn value_cache_mut(&mut self) -> &mut SharedValueCache<H::Out> {
+		&mut self.value_cache
+	}
+
+	/// Returns a reference to the [`SharedNodeCache`].
+	#[cfg(test)]
+	pub(super) fn node_cache(&self) -> &SharedNodeCache<H::Out> {
+		&self.node_cache
+	}
+
+	/// Returns a mutable reference to the [`SharedNodeCache`].
+	pub(super) fn node_cache_mut(&mut self) -> &mut SharedNodeCache<H::Out> {
+		&mut self.node_cache
+	}
+
+	pub(super) fn metrics(&self) -> Option<&Metrics> {
+		self.metrics.as_ref()
+	}
+
+	/// Returns a mutable reference to the [`TrieHitStats`].
+	pub(super) fn stats_add_snapshot(&mut self, snapshot: &TrieHitStatsSnapshot) {
+		self.stats.add_snapshot(&snapshot);
+		// Print trie cache stats every 60 seconds.
+		if self.previous_stats_dump.elapsed() > Duration::from_secs(60) {
+			self.previous_stats_dump = Instant::now();
+			let snapshot = self.stats.snapshot();
+			tracing::trace!(target: LOG_TARGET, node_cache = %snapshot.node_cache, value_cache = %snapshot.value_cache, "Shared trie cache stats");
+		}
+	}
+}
+
+/// The shared trie cache.
+///
+/// It should be instantiated once per node. It will hold the trie nodes and values of all
+/// operations to the state. To not use all available memory it will ensure to stay in the
+/// bounds given via the [`CacheSize`] at startup.
+///
+/// The instance of this object can be shared between multiple threads.
+pub struct SharedTrieCache<H: Hasher> {
+	inner: Arc<RwLock<SharedTrieCacheInner<H>>>,
+	trusted_node_cache_config: LocalNodeCacheConfig,
+	trusted_value_cache_config: LocalValueCacheConfig,
+}
+
+impl<H: Hasher> Clone for SharedTrieCache<H> {
+	fn clone(&self) -> Self {
+		Self {
+			inner: self.inner.clone(),
+			trusted_node_cache_config: self.trusted_node_cache_config,
+			trusted_value_cache_config: self.trusted_value_cache_config,
+		}
+	}
+}
+
+impl<H: Hasher> SharedTrieCache<H> {
+	/// Create a new [`SharedTrieCache`].
+	pub fn new(cache_size: CacheSize, metrics_registry: Option<&Registry>) -> Self {
+		let total_budget = cache_size.0;
+
+		// Split our memory budget between the two types of caches.
+		let value_cache_budget = (total_budget as f32 * 0.20) as usize; // 20% for the value cache
+		let node_cache_budget = total_budget - value_cache_budget; // 80% for the node cache
+
+		// Split our memory budget between what we'll be holding inline in the map,
+		// and what we'll be holding on the heap.
+		let value_cache_inline_budget = (value_cache_budget as f32 * 0.70) as usize;
+		let node_cache_inline_budget = (node_cache_budget as f32 * 0.70) as usize;
+
+		// Calculate how much memory the maps will be allowed to hold inline given our budget.
+		let value_cache_max_inline_size =
+			SharedValueCacheMap::<H::Out>::memory_usage_for_memory_budget(
+				value_cache_inline_budget,
+			);
+
+		let node_cache_max_inline_size =
+			SharedNodeCacheMap::<H::Out>::memory_usage_for_memory_budget(node_cache_inline_budget);
+
+		// And this is how much data we'll at most keep on the heap for each cache.
+		let value_cache_max_heap_size = value_cache_budget - value_cache_max_inline_size;
+		let node_cache_max_heap_size = node_cache_budget - node_cache_max_inline_size;
+
+		tracing::debug!(
+			target: super::LOG_TARGET,
+			"Configured a shared trie cache with a budget of ~{} bytes (node_cache_max_inline_size = {}, node_cache_max_heap_size = {}, value_cache_max_inline_size = {}, value_cache_max_heap_size = {})",
+			total_budget,
+			node_cache_max_inline_size,
+			node_cache_max_heap_size,
+			value_cache_max_inline_size,
+			value_cache_max_heap_size,
+		);
+
+		Self {
+			inner: Arc::new(RwLock::new(SharedTrieCacheInner {
+				node_cache: SharedNodeCache::new(
+					node_cache_max_inline_size,
+					node_cache_max_heap_size,
+				),
+				value_cache: SharedValueCache::new(
+					value_cache_max_inline_size,
+					value_cache_max_heap_size,
+				),
+				stats: Default::default(),
+				previous_stats_dump: Instant::now(),
+				metrics: metrics_registry.and_then(|registry| Metrics::register(registry).ok()),
+			})),
+			trusted_node_cache_config: LocalNodeCacheConfig::trusted(
+				node_cache_max_heap_size,
+				node_cache_max_inline_size,
+			),
+			trusted_value_cache_config: LocalValueCacheConfig::trusted(
+				value_cache_max_heap_size,
+				value_cache_max_inline_size,
+			),
+		}
+	}
+
+	/// Create a new [`LocalTrieCache`](super::LocalTrieCache) instance from this shared cache.
+	pub fn local_cache_untrusted(&self) -> super::LocalTrieCache<H> {
+		let local_value_cache_config = LocalValueCacheConfig::untrusted();
+		let local_node_cache_config = LocalNodeCacheConfig::untrusted();
+		tracing::debug!(
+			target: super::LOG_TARGET,
+			"Configuring a local un-trusted cache"
+		);
+
+		super::LocalTrieCache {
+			shared: self.clone(),
+			node_cache: Mutex::new(LruMap::new(LocalNodeCacheLimiter::new(
+				local_node_cache_config,
+			))),
+			value_cache: Mutex::new(LruMap::with_hasher(
+				LocalValueCacheLimiter::new(local_value_cache_config),
+				Default::default(),
+			)),
+			shared_value_cache_access: Mutex::new(super::ValueAccessSet::with_hasher(
+				schnellru::ByLength::new(
+					local_value_cache_config.shared_value_cache_max_promoted_keys,
+				),
+				Default::default(),
+			)),
+			value_cache_config: local_value_cache_config,
+			node_cache_config: local_node_cache_config,
+			stats: Default::default(),
+			trusted: false,
+		}
+	}
+
+	/// Creates a TrieCache that allows the local_caches to grow to indefinitely.
+	///
+	/// This is safe to be used only for trusted paths because it removes all limits on cache
+	/// growth and promotion, which could lead to excessive memory usage if used in untrusted or
+	/// uncontrolled environments. It is intended for scenarios like block authoring or importing,
+	/// where the operations are bounded already and there are no risks of unbounded memory usage.
+	pub fn local_cache_trusted(&self) -> super::LocalTrieCache<H> {
+		tracing::debug!(
+			target: super::LOG_TARGET,
+			"Configuring a local trusted cache"
+		);
+		super::LocalTrieCache {
+			shared: self.clone(),
+			node_cache: Mutex::new(LruMap::new(LocalNodeCacheLimiter::new(
+				self.trusted_node_cache_config,
+			))),
+			value_cache: Mutex::new(LruMap::with_hasher(
+				LocalValueCacheLimiter::new(self.trusted_value_cache_config),
+				Default::default(),
+			)),
+			shared_value_cache_access: Mutex::new(super::ValueAccessSet::with_hasher(
+				schnellru::ByLength::new(
+					self.trusted_value_cache_config.shared_value_cache_max_promoted_keys,
+				),
+				Default::default(),
+			)),
+			value_cache_config: self.trusted_value_cache_config,
+			node_cache_config: self.trusted_node_cache_config,
+			stats: Default::default(),
+			trusted: true,
+		}
+	}
+
+	/// Get a copy of the node for `key`.
+	///
+	/// This will temporarily lock the shared cache for reading.
+	///
+	/// This doesn't change the least recently order in the internal [`LruMap`].
+	#[inline]
+	pub fn peek_node(&self, key: &H::Out) -> Option<NodeOwned<H::Out>> {
+		self.inner.read().node_cache.lru.peek(key).cloned()
+	}
+
+	/// Get a copy of the [`CachedValue`] for `key`.
+	///
+	/// This will temporarily lock the shared cache for reading.
+	///
+	/// This doesn't reorder any of the elements in the internal [`LruMap`].
+	pub fn peek_value_by_hash(
+		&self,
+		hash: ValueCacheKeyHash,
+		storage_root: &H::Out,
+		storage_key: &[u8],
+	) -> Option<CachedValue<H::Out>> {
+		self.inner
+			.read()
+			.value_cache
+			.lru
+			.peek_by_hash(hash.0, |existing_key, _| existing_key.is_eq(storage_root, storage_key))
+			.cloned()
+	}
+
+	/// Returns the used memory size of this cache in bytes.
+	pub fn used_memory_size(&self) -> usize {
+		let inner = self.inner.read();
+		let value_cache_size =
+			inner.value_cache.lru.memory_usage() + inner.value_cache.lru.limiter().heap_size;
+		let node_cache_size =
+			inner.node_cache.lru.memory_usage() + inner.node_cache.lru.limiter().heap_size;
+
+		node_cache_size + value_cache_size
+	}
+
+	/// Reset the node cache.
+	pub fn reset_node_cache(&self) {
+		self.inner.write().node_cache.reset();
+	}
+
+	/// Reset the value cache.
+	pub fn reset_value_cache(&self) {
+		self.inner.write().value_cache.reset();
+	}
+
+	/// Reset the entire cache.
+	pub fn reset(&self) {
+		self.reset_node_cache();
+		self.reset_value_cache();
+	}
+
+	/// Returns the read locked inner.
+	#[cfg(test)]
+	pub(super) fn read_lock_inner(
+		&self,
+	) -> parking_lot::RwLockReadGuard<'_, SharedTrieCacheInner<H>> {
+		self.inner.read()
+	}
+
+	/// Returns the write locked inner.
+	pub(super) fn write_lock_inner(&self) -> Option<RwLockWriteGuard<'_, SharedTrieCacheInner<H>>> {
+		// This should never happen, but we *really* don't want to deadlock. So let's have it
+		// timeout, just in case. At worst it'll do nothing, and at best it'll avert a catastrophe
+		// and notify us that there's a problem.
+		self.inner.try_write_for(super::SHARED_CACHE_WRITE_LOCK_TIMEOUT)
+	}
+}
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+	use pezsp_core::H256 as Hash;
+
+	#[test]
+	fn shared_value_cache_works() {
+		let mut cache = SharedValueCache::<pezsp_core::H256>::new(usize::MAX, 10 * 10);
+
+		let key = vec![0; 10];
+
+		let root0 = Hash::repeat_byte(1);
+		let root1 = Hash::repeat_byte(2);
+
+		cache.update(
+			vec![
+				(ValueCacheKey::new_value(&key[..], root0), CachedValue::NonExisting),
+				(ValueCacheKey::new_value(&key[..], root1), CachedValue::NonExisting),
+			],
+			vec![],
+			&LocalValueCacheConfig::untrusted(),
+			&None,
+		);
+
+		// Ensure that the basics are working
+		assert_eq!(1, cache.lru.limiter_mut().known_storage_keys.len());
+		assert_eq!(
+			3, // Two instances inside the cache + one extra in `known_storage_keys`.
+			Arc::strong_count(
+				cache.lru.limiter_mut().known_storage_keys.get_key_value(&key[..]).unwrap().0
+			)
+		);
+		assert_eq!(key.len(), cache.lru.limiter().heap_size);
+		assert_eq!(cache.lru.len(), 2);
+		assert_eq!(cache.lru.peek_newest().unwrap().0.storage_root, root1);
+		assert_eq!(cache.lru.peek_oldest().unwrap().0.storage_root, root0);
+		assert!(cache.lru.limiter().heap_size <= cache.lru.limiter().max_heap_size);
+		assert_eq!(cache.lru.limiter().heap_size, 10);
+
+		// Just accessing a key should not change anything on the size and number of entries.
+		cache.update(
+			vec![],
+			vec![ValueCacheKey::hash_data(&key[..], &root0)],
+			&LocalValueCacheConfig::untrusted(),
+			&None,
+		);
+		assert_eq!(1, cache.lru.limiter_mut().known_storage_keys.len());
+		assert_eq!(
+			3,
+			Arc::strong_count(
+				cache.lru.limiter_mut().known_storage_keys.get_key_value(&key[..]).unwrap().0
+			)
+		);
+		assert_eq!(key.len(), cache.lru.limiter().heap_size);
+		assert_eq!(cache.lru.len(), 2);
+		assert_eq!(cache.lru.peek_newest().unwrap().0.storage_root, root0);
+		assert_eq!(cache.lru.peek_oldest().unwrap().0.storage_root, root1);
+		assert!(cache.lru.limiter().heap_size <= cache.lru.limiter().max_heap_size);
+		assert_eq!(cache.lru.limiter().heap_size, 10);
+
+		// Updating the cache again with exactly the same data should not change anything.
+		cache.update(
+			vec![
+				(ValueCacheKey::new_value(&key[..], root1), CachedValue::NonExisting),
+				(ValueCacheKey::new_value(&key[..], root0), CachedValue::NonExisting),
+			],
+			vec![],
+			&LocalValueCacheConfig::untrusted(),
+			&None,
+		);
+		assert_eq!(1, cache.lru.limiter_mut().known_storage_keys.len());
+		assert_eq!(
+			3,
+			Arc::strong_count(
+				cache.lru.limiter_mut().known_storage_keys.get_key_value(&key[..]).unwrap().0
+			)
+		);
+		assert_eq!(key.len(), cache.lru.limiter().heap_size);
+		assert_eq!(cache.lru.len(), 2);
+		assert_eq!(cache.lru.peek_newest().unwrap().0.storage_root, root0);
+		assert_eq!(cache.lru.peek_oldest().unwrap().0.storage_root, root1);
+		assert!(cache.lru.limiter().heap_size <= cache.lru.limiter().max_heap_size);
+		assert_eq!(cache.lru.limiter().items_evicted, 0);
+		assert_eq!(cache.lru.limiter().heap_size, 10);
+
+		// Add 10 other entries and this should move out two of the initial entries.
+		cache.update(
+			(1..11)
+				.map(|i| vec![i; 10])
+				.map(|key| (ValueCacheKey::new_value(&key[..], root0), CachedValue::NonExisting)),
+			vec![],
+			&LocalValueCacheConfig::untrusted(),
+			&None,
+		);
+
+		assert_eq!(cache.lru.limiter().items_evicted, 2);
+		assert_eq!(10, cache.lru.len());
+		assert_eq!(10, cache.lru.limiter_mut().known_storage_keys.len());
+		assert!(cache.lru.limiter_mut().known_storage_keys.get_key_value(&key[..]).is_none());
+		assert_eq!(key.len() * 10, cache.lru.limiter().heap_size);
+		assert_eq!(cache.lru.len(), 10);
+		assert!(cache.lru.limiter().heap_size <= cache.lru.limiter().max_heap_size);
+		assert_eq!(cache.lru.limiter().heap_size, 100);
+
+		assert!(matches!(
+			cache.lru.peek(&ValueCacheKey::new_value(&[1; 10][..], root0)).unwrap(),
+			CachedValue::<Hash>::NonExisting
+		));
+
+		assert!(cache.lru.peek(&ValueCacheKey::new_value(&[1; 10][..], root1)).is_none(),);
+
+		assert!(cache.lru.peek(&ValueCacheKey::new_value(&key[..], root0)).is_none());
+		assert!(cache.lru.peek(&ValueCacheKey::new_value(&key[..], root1)).is_none());
+
+		cache.update(
+			vec![(ValueCacheKey::new_value(vec![10; 10], root0), CachedValue::NonExisting)],
+			vec![],
+			&LocalValueCacheConfig::untrusted(),
+			&None,
+		);
+
+		assert!(cache.lru.limiter_mut().known_storage_keys.get_key_value(&key[..]).is_none());
+	}
+}

bizinikiwi/primitives/trie/src/error.rs

@@ -0,0 +1,47 @@
+// This file is part of Bizinikiwi.
+
+// Copyright (C) Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use alloc::{boxed::Box, vec::Vec};
+
+/// Error type used for trie related errors.
+#[derive(Debug, PartialEq, Eq, Clone)]
+#[cfg_attr(feature = "std", derive(thiserror::Error))]
+pub enum Error<H> {
+	#[cfg_attr(feature = "std", error("Bad format"))]
+	BadFormat,
+	#[cfg_attr(feature = "std", error("Decoding failed: {0}"))]
+	Decode(#[cfg_attr(feature = "std", source)] codec::Error),
+	#[cfg_attr(
+		feature = "std",
+		error("Recorded key ({0:x?}) access with value as found={1}, but could not confirm with trie.")
+	)]
+	InvalidRecording(Vec<u8>, bool),
+	#[cfg_attr(feature = "std", error("Trie error: {0:?}"))]
+	TrieError(Box<trie_db::TrieError<H, Self>>),
+}
+
+impl<H> From<codec::Error> for Error<H> {
+	fn from(x: codec::Error) -> Self {
+		Error::Decode(x)
+	}
+}
+
+impl<H> From<Box<trie_db::TrieError<H, Self>>> for Error<H> {
+	fn from(x: Box<trie_db::TrieError<H, Self>>) -> Self {
+		Error::TrieError(x)
+	}
+}

bizinikiwi/primitives/trie/src/hasher_random_state.rs

@@ -0,0 +1,171 @@
+// This file is part of Bizinikiwi.
+
+// Copyright (C) Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Utility module to use a custom random state for HashMap and friends
+//! in a no_std environment.
+
+use core::{
+	cell::UnsafeCell,
+	hash::Hasher as CoreHasher,
+	sync::atomic::{AtomicU8, Ordering},
+};
+
+use core::hash::BuildHasher;
+use foldhash::quality::RandomState as FoldHashBuilder;
+
+// Constants to represent the state of the global extra randomness.
+// UNINITIALIZED: The extra randomness has not been set yet.
+const UNINITIALIZED: u8 = 0;
+// LOCKED: The extra randomness is being set.
+const LOCKED: u8 = 1;
+// INITIALIZED: The extra randomness has been set and is ready to use.
+const INITIALIZED: u8 = 2;
+
+// SAFETY: we only mutate the UnsafeCells when state is in the thread-exclusive
+// LOCKED state, and only read when state is in the INITIALIZED state.
+unsafe impl Sync for GlobalExtraRandomnesss {}
+struct GlobalExtraRandomnesss {
+	initialized: AtomicU8,
+	randomness: UnsafeCell<[u8; 16]>,
+}
+
+// Extra randomness to be used besides the one provided by the `FoldHashBuilder`.
+static EXTRA_RANDOMNESS: GlobalExtraRandomnesss = GlobalExtraRandomnesss {
+	initialized: AtomicU8::new(UNINITIALIZED),
+	randomness: UnsafeCell::new([0u8; 16]),
+};
+
+/// Adds extra randomness to be used by all new instances of RandomState.
+pub fn add_extra_randomness(extra_randomness: [u8; 16]) {
+	match EXTRA_RANDOMNESS.initialized.compare_exchange(
+		UNINITIALIZED,
+		LOCKED,
+		Ordering::Acquire,
+		Ordering::Acquire,
+	) {
+		Ok(_) => {
+			// SAFETY: We are the only ones writing exclusively to this memory.
+			unsafe { *EXTRA_RANDOMNESS.randomness.get() = extra_randomness };
+			EXTRA_RANDOMNESS.initialized.store(INITIALIZED, Ordering::Release);
+		},
+		Err(_) => {
+			panic!("Extra randomness has already been set, cannot set it again.");
+		},
+	}
+}
+
+// Returns the extra randomness if it has been set, otherwise returns None.
+fn extra_randomness() -> Option<&'static [u8; 16]> {
+	// SAFETY: We are reading from a static memory location that is initialized
+	// only once, so it is safe to read from it.
+	if EXTRA_RANDOMNESS.initialized.load(Ordering::Acquire) == INITIALIZED {
+		Some(unsafe { &*EXTRA_RANDOMNESS.randomness.get() })
+	} else {
+		None
+	}
+}
+
+/// A wrapper around `FoldHashBuilder` that adds extra randomness to the hashers it creates.
+#[derive(Copy, Clone, Debug)]
+pub struct RandomState {
+	default: FoldHashBuilder,
+	extra_randomness: Option<&'static [u8; 16]>,
+}
+
+impl Default for RandomState {
+	#[inline(always)]
+	fn default() -> Self {
+		RandomState {
+			// FoldHashBuilder already uses a random seed, so we use that as the base.
+			default: FoldHashBuilder::default(),
+			extra_randomness: extra_randomness(),
+		}
+	}
+}
+
+impl BuildHasher for RandomState {
+	type Hasher = <FoldHashBuilder as BuildHasher>::Hasher;
+
+	#[inline(always)]
+	fn build_hasher(&self) -> Self::Hasher {
+		let mut hasher = self.default.build_hasher();
+		if let Some(extra) = self.extra_randomness {
+			// If extra randomness is set, we write it into the hasher.
+			hasher.write(extra);
+		}
+
+		hasher
+	}
+}
+
+#[cfg(test)]
+mod tests {
+	use core::hash::{BuildHasher, Hasher};
+
+	#[test]
+	fn hashbuilder_produces_same_result() {
+		let haser_builder = super::RandomState::default();
+		let mut hasher_1 = haser_builder.build_hasher();
+		let mut hasher_2 = haser_builder.build_hasher();
+
+		hasher_1.write_u32(8128);
+		hasher_2.write_u32(8128);
+
+		assert_eq!(hasher_1.finish(), hasher_2.finish());
+	}
+
+	#[test]
+	fn adding_randomness_does_not_affect_already_instantiated_builders() {
+		let hasher_builder = super::RandomState::default();
+		let mut hasher_1 = hasher_builder.build_hasher();
+
+		let randomness = [0xde; 16];
+		super::add_extra_randomness(randomness);
+		let builder_after_randomness_added = super::RandomState::default();
+		assert_eq!(builder_after_randomness_added.extra_randomness, Some(&randomness));
+
+		let mut hasher_2 = hasher_builder.build_hasher();
+
+		hasher_1.write_u32(8128);
+		hasher_2.write_u32(8128);
+
+		assert_eq!(hasher_1.finish(), hasher_2.finish());
+	}
+
+	#[test]
+	fn sanity_check() {
+		let haser_builder = super::RandomState::default();
+		let mut hasher_create_manually =
+			hashbrown::HashMap::<u32, u32, _>::with_hasher(haser_builder);
+		let mut default_built = hashbrown::HashMap::<u32, u32, super::RandomState>::default();
+
+		for x in 0..100 {
+			default_built.insert(x, x * 2);
+			hasher_create_manually.insert(x, x * 2);
+		}
+
+		for x in 0..100 {
+			assert_eq!(default_built.get(&x), Some(&(x * 2)));
+			assert_eq!(hasher_create_manually.get(&x), Some(&(x * 2)));
+		}
+
+		for x in 100..200 {
+			assert_eq!(default_built.get(&x), None);
+			assert_eq!(hasher_create_manually.get(&x), None);
+		}
+	}
+}

bizinikiwi/primitives/trie/src/node_codec.rs

@@ -0,0 +1,333 @@
+// This file is part of Bizinikiwi.
+
+// Copyright (C) Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! `NodeCodec` implementation for Bizinikiwi's trie format.
+
+use super::node_header::{NodeHeader, NodeKind};
+use crate::{error::Error, trie_constants};
+use alloc::{borrow::Borrow, vec::Vec};
+use codec::{Compact, Decode, Encode, Input};
+use core::{marker::PhantomData, ops::Range};
+use hash_db::Hasher;
+use trie_db::{
+	nibble_ops,
+	node::{NibbleSlicePlan, NodeHandlePlan, NodePlan, Value, ValuePlan},
+	ChildReference, NodeCodec as NodeCodecT,
+};
+
+/// Helper struct for trie node decoder. This implements `codec::Input` on a byte slice, while
+/// tracking the absolute position. This is similar to `std::io::Cursor` but does not implement
+/// `Read` and `io` are not in `core` or `alloc`.
+struct ByteSliceInput<'a> {
+	data: &'a [u8],
+	offset: usize,
+}
+
+impl<'a> ByteSliceInput<'a> {
+	fn new(data: &'a [u8]) -> Self {
+		ByteSliceInput { data, offset: 0 }
+	}
+
+	fn take(&mut self, count: usize) -> Result<Range<usize>, codec::Error> {
+		if self.offset + count > self.data.len() {
+			return Err("out of data".into());
+		}
+
+		let range = self.offset..(self.offset + count);
+		self.offset += count;
+		Ok(range)
+	}
+}
+
+impl<'a> Input for ByteSliceInput<'a> {
+	fn remaining_len(&mut self) -> Result<Option<usize>, codec::Error> {
+		Ok(Some(self.data.len().saturating_sub(self.offset)))
+	}
+
+	fn read(&mut self, into: &mut [u8]) -> Result<(), codec::Error> {
+		let range = self.take(into.len())?;
+		into.copy_from_slice(&self.data[range]);
+		Ok(())
+	}
+
+	fn read_byte(&mut self) -> Result<u8, codec::Error> {
+		if self.offset + 1 > self.data.len() {
+			return Err("out of data".into());
+		}
+
+		let byte = self.data[self.offset];
+		self.offset += 1;
+		Ok(byte)
+	}
+}
+
+/// Concrete implementation of a [`NodeCodecT`] with SCALE encoding.
+///
+/// It is generic over `H` the [`Hasher`].
+#[derive(Default, Clone)]
+pub struct NodeCodec<H>(PhantomData<H>);
+
+impl<H> NodeCodecT for NodeCodec<H>
+where
+	H: Hasher,
+{
+	const ESCAPE_HEADER: Option<u8> = Some(trie_constants::ESCAPE_COMPACT_HEADER);
+	type Error = Error<H::Out>;
+	type HashOut = H::Out;
+
+	fn hashed_null_node() -> <H as Hasher>::Out {
+		H::hash(<Self as NodeCodecT>::empty_node())
+	}
+
+	fn decode_plan(data: &[u8]) -> Result<NodePlan, Self::Error> {
+		let mut input = ByteSliceInput::new(data);
+
+		let header = NodeHeader::decode(&mut input)?;
+		let contains_hash = header.contains_hash_of_value();
+
+		let branch_has_value = if let NodeHeader::Branch(has_value, _) = &header {
+			*has_value
+		} else {
+			// hashed_value_branch
+			true
+		};
+
+		match header {
+			NodeHeader::Null => Ok(NodePlan::Empty),
+			NodeHeader::HashedValueBranch(nibble_count) | NodeHeader::Branch(_, nibble_count) => {
+				let padding = nibble_count % nibble_ops::NIBBLE_PER_BYTE != 0;
+				// data should be at least of size offset + 1
+				if data.len() < input.offset + 1 {
+					return Err(Error::BadFormat);
+				}
+				// check that the padding is valid (if any)
+				if padding && nibble_ops::pad_left(data[input.offset]) != 0 {
+					return Err(Error::BadFormat);
+				}
+				let partial = input.take(nibble_count.div_ceil(nibble_ops::NIBBLE_PER_BYTE))?;
+				let partial_padding = nibble_ops::number_padding(nibble_count);
+				let bitmap_range = input.take(BITMAP_LENGTH)?;
+				let bitmap = Bitmap::decode(&data[bitmap_range])?;
+				let value = if branch_has_value {
+					Some(if contains_hash {
+						ValuePlan::Node(input.take(H::LENGTH)?)
+					} else {
+						let count = <Compact<u32>>::decode(&mut input)?.0 as usize;
+						ValuePlan::Inline(input.take(count)?)
+					})
+				} else {
+					None
+				};
+				let mut children = [
+					None, None, None, None, None, None, None, None, None, None, None, None, None,
+					None, None, None,
+				];
+				for i in 0..nibble_ops::NIBBLE_LENGTH {
+					if bitmap.value_at(i) {
+						let count = <Compact<u32>>::decode(&mut input)?.0 as usize;
+						let range = input.take(count)?;
+						children[i] = Some(if count == H::LENGTH {
+							NodeHandlePlan::Hash(range)
+						} else {
+							NodeHandlePlan::Inline(range)
+						});
+					}
+				}
+				Ok(NodePlan::NibbledBranch {
+					partial: NibbleSlicePlan::new(partial, partial_padding),
+					value,
+					children,
+				})
+			},
+			NodeHeader::HashedValueLeaf(nibble_count) | NodeHeader::Leaf(nibble_count) => {
+				let padding = nibble_count % nibble_ops::NIBBLE_PER_BYTE != 0;
+				// data should be at least of size offset + 1
+				if data.len() < input.offset + 1 {
+					return Err(Error::BadFormat);
+				}
+				// check that the padding is valid (if any)
+				if padding && nibble_ops::pad_left(data[input.offset]) != 0 {
+					return Err(Error::BadFormat);
+				}
+				let partial = input.take(nibble_count.div_ceil(nibble_ops::NIBBLE_PER_BYTE))?;
+				let partial_padding = nibble_ops::number_padding(nibble_count);
+				let value = if contains_hash {
+					ValuePlan::Node(input.take(H::LENGTH)?)
+				} else {
+					let count = <Compact<u32>>::decode(&mut input)?.0 as usize;
+					ValuePlan::Inline(input.take(count)?)
+				};
+
+				Ok(NodePlan::Leaf {
+					partial: NibbleSlicePlan::new(partial, partial_padding),
+					value,
+				})
+			},
+		}
+	}
+
+	fn is_empty_node(data: &[u8]) -> bool {
+		data == <Self as NodeCodecT>::empty_node()
+	}
+
+	fn empty_node() -> &'static [u8] {
+		&[trie_constants::EMPTY_TRIE]
+	}
+
+	fn leaf_node(partial: impl Iterator<Item = u8>, number_nibble: usize, value: Value) -> Vec<u8> {
+		let contains_hash = matches!(&value, Value::Node(..));
+		let mut output = if contains_hash {
+			partial_from_iterator_encode(partial, number_nibble, NodeKind::HashedValueLeaf)
+		} else {
+			partial_from_iterator_encode(partial, number_nibble, NodeKind::Leaf)
+		};
+		match value {
+			Value::Inline(value) => {
+				Compact(value.len() as u32).encode_to(&mut output);
+				output.extend_from_slice(value);
+			},
+			Value::Node(hash) => {
+				debug_assert!(hash.len() == H::LENGTH);
+				output.extend_from_slice(hash);
+			},
+		}
+		output
+	}
+
+	fn extension_node(
+		_partial: impl Iterator<Item = u8>,
+		_nbnibble: usize,
+		_child: ChildReference<<H as Hasher>::Out>,
+	) -> Vec<u8> {
+		unreachable!("No extension codec.")
+	}
+
+	fn branch_node(
+		_children: impl Iterator<Item = impl Borrow<Option<ChildReference<<H as Hasher>::Out>>>>,
+		_maybe_value: Option<Value>,
+	) -> Vec<u8> {
+		unreachable!("No extension codec.")
+	}
+
+	fn branch_node_nibbled(
+		partial: impl Iterator<Item = u8>,
+		number_nibble: usize,
+		children: impl Iterator<Item = impl Borrow<Option<ChildReference<<H as Hasher>::Out>>>>,
+		value: Option<Value>,
+	) -> Vec<u8> {
+		let contains_hash = matches!(&value, Some(Value::Node(..)));
+		let mut output = match (&value, contains_hash) {
+			(&None, _) =>
+				partial_from_iterator_encode(partial, number_nibble, NodeKind::BranchNoValue),
+			(_, false) =>
+				partial_from_iterator_encode(partial, number_nibble, NodeKind::BranchWithValue),
+			(_, true) =>
+				partial_from_iterator_encode(partial, number_nibble, NodeKind::HashedValueBranch),
+		};
+
+		let bitmap_index = output.len();
+		let mut bitmap: [u8; BITMAP_LENGTH] = [0; BITMAP_LENGTH];
+		(0..BITMAP_LENGTH).for_each(|_| output.push(0));
+		match value {
+			Some(Value::Inline(value)) => {
+				Compact(value.len() as u32).encode_to(&mut output);
+				output.extend_from_slice(value);
+			},
+			Some(Value::Node(hash)) => {
+				debug_assert!(hash.len() == H::LENGTH);
+				output.extend_from_slice(hash);
+			},
+			None => (),
+		}
+		Bitmap::encode(
+			children.map(|maybe_child| match maybe_child.borrow() {
+				Some(ChildReference::Hash(h)) => {
+					h.as_ref().encode_to(&mut output);
+					true
+				},
+				&Some(ChildReference::Inline(inline_data, len)) => {
+					inline_data.as_ref()[..len].encode_to(&mut output);
+					true
+				},
+				None => false,
+			}),
+			bitmap.as_mut(),
+		);
+		output[bitmap_index..bitmap_index + BITMAP_LENGTH]
+			.copy_from_slice(&bitmap[..BITMAP_LENGTH]);
+		output
+	}
+}
+
+// utils
+
+/// Encode and allocate node type header (type and size), and partial value.
+/// It uses an iterator over encoded partial bytes as input.
+fn partial_from_iterator_encode<I: Iterator<Item = u8>>(
+	partial: I,
+	nibble_count: usize,
+	node_kind: NodeKind,
+) -> Vec<u8> {
+	let mut output = Vec::with_capacity(4 + (nibble_count / nibble_ops::NIBBLE_PER_BYTE));
+	match node_kind {
+		NodeKind::Leaf => NodeHeader::Leaf(nibble_count).encode_to(&mut output),
+		NodeKind::BranchWithValue => NodeHeader::Branch(true, nibble_count).encode_to(&mut output),
+		NodeKind::BranchNoValue => NodeHeader::Branch(false, nibble_count).encode_to(&mut output),
+		NodeKind::HashedValueLeaf =>
+			NodeHeader::HashedValueLeaf(nibble_count).encode_to(&mut output),
+		NodeKind::HashedValueBranch =>
+			NodeHeader::HashedValueBranch(nibble_count).encode_to(&mut output),
+	};
+	output.extend(partial);
+	output
+}
+
+const BITMAP_LENGTH: usize = 2;
+
+/// Radix 16 trie, bitmap encoding implementation,
+/// it contains children mapping information for a branch
+/// (children presence only), it encodes into
+/// a compact bitmap encoding representation.
+pub(crate) struct Bitmap(u16);
+
+impl Bitmap {
+	pub fn decode(data: &[u8]) -> Result<Self, codec::Error> {
+		let value = u16::decode(&mut &data[..])?;
+		if value == 0 {
+			Err("Bitmap without a child.".into())
+		} else {
+			Ok(Bitmap(value))
+		}
+	}
+
+	pub fn value_at(&self, i: usize) -> bool {
+		self.0 & (1u16 << i) != 0
+	}
+
+	pub fn encode<I: Iterator<Item = bool>>(has_children: I, dest: &mut [u8]) {
+		let mut bitmap: u16 = 0;
+		let mut cursor: u16 = 1;
+		for v in has_children {
+			if v {
+				bitmap |= cursor
+			}
+			cursor <<= 1;
+		}
+		dest[0] = (bitmap % 256) as u8;
+		dest[1] = (bitmap / 256) as u8;
+	}
+}

bizinikiwi/primitives/trie/src/node_header.rs

@@ -0,0 +1,173 @@
+// This file is part of Bizinikiwi.
+
+// Copyright (C) Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! The node header.
+
+use crate::trie_constants;
+use codec::{Decode, Encode, Input, Output};
+use core::iter::once;
+
+/// A node header
+#[derive(Copy, Clone, PartialEq, Eq, pezsp_core::RuntimeDebug)]
+pub(crate) enum NodeHeader {
+	Null,
+	// contains wether there is a value and nibble count
+	Branch(bool, usize),
+	// contains nibble count
+	Leaf(usize),
+	// contains nibble count.
+	HashedValueBranch(usize),
+	// contains nibble count.
+	HashedValueLeaf(usize),
+}
+
+impl NodeHeader {
+	pub(crate) fn contains_hash_of_value(&self) -> bool {
+		matches!(self, NodeHeader::HashedValueBranch(_) | NodeHeader::HashedValueLeaf(_))
+	}
+}
+
+/// NodeHeader without content
+pub(crate) enum NodeKind {
+	Leaf,
+	BranchNoValue,
+	BranchWithValue,
+	HashedValueLeaf,
+	HashedValueBranch,
+}
+
+impl Encode for NodeHeader {
+	fn encode_to<T: Output + ?Sized>(&self, output: &mut T) {
+		match self {
+			NodeHeader::Null => output.push_byte(trie_constants::EMPTY_TRIE),
+			NodeHeader::Branch(true, nibble_count) =>
+				encode_size_and_prefix(*nibble_count, trie_constants::BRANCH_WITH_MASK, 2, output),
+			NodeHeader::Branch(false, nibble_count) => encode_size_and_prefix(
+				*nibble_count,
+				trie_constants::BRANCH_WITHOUT_MASK,
+				2,
+				output,
+			),
+			NodeHeader::Leaf(nibble_count) =>
+				encode_size_and_prefix(*nibble_count, trie_constants::LEAF_PREFIX_MASK, 2, output),
+			NodeHeader::HashedValueBranch(nibble_count) => encode_size_and_prefix(
+				*nibble_count,
+				trie_constants::ALT_HASHING_BRANCH_WITH_MASK,
+				4,
+				output,
+			),
+			NodeHeader::HashedValueLeaf(nibble_count) => encode_size_and_prefix(
+				*nibble_count,
+				trie_constants::ALT_HASHING_LEAF_PREFIX_MASK,
+				3,
+				output,
+			),
+		}
+	}
+}
+
+impl codec::EncodeLike for NodeHeader {}
+
+impl Decode for NodeHeader {
+	fn decode<I: Input>(input: &mut I) -> Result<Self, codec::Error> {
+		let i = input.read_byte()?;
+		if i == trie_constants::EMPTY_TRIE {
+			return Ok(NodeHeader::Null);
+		}
+		match i & (0b11 << 6) {
+			trie_constants::LEAF_PREFIX_MASK => Ok(NodeHeader::Leaf(decode_size(i, input, 2)?)),
+			trie_constants::BRANCH_WITH_MASK =>
+				Ok(NodeHeader::Branch(true, decode_size(i, input, 2)?)),
+			trie_constants::BRANCH_WITHOUT_MASK =>
+				Ok(NodeHeader::Branch(false, decode_size(i, input, 2)?)),
+			trie_constants::EMPTY_TRIE => {
+				if i & (0b111 << 5) == trie_constants::ALT_HASHING_LEAF_PREFIX_MASK {
+					Ok(NodeHeader::HashedValueLeaf(decode_size(i, input, 3)?))
+				} else if i & (0b1111 << 4) == trie_constants::ALT_HASHING_BRANCH_WITH_MASK {
+					Ok(NodeHeader::HashedValueBranch(decode_size(i, input, 4)?))
+				} else {
+					// do not allow any special encoding
+					Err("Unallowed encoding".into())
+				}
+			},
+			_ => unreachable!(),
+		}
+	}
+}
+
+/// Returns an iterator over encoded bytes for node header and size.
+/// Size encoding allows unlimited, length inefficient, representation, but
+/// is bounded to 16 bit maximum value to avoid possible DOS.
+pub(crate) fn size_and_prefix_iterator(
+	size: usize,
+	prefix: u8,
+	prefix_mask: usize,
+) -> impl Iterator<Item = u8> {
+	let max_value = 255u8 >> prefix_mask;
+	let l1 = core::cmp::min((max_value as usize).saturating_sub(1), size);
+	let (first_byte, mut rem) = if size == l1 {
+		(once(prefix + l1 as u8), 0)
+	} else {
+		(once(prefix + max_value as u8), size - l1)
+	};
+	let next_bytes = move || {
+		if rem > 0 {
+			if rem < 256 {
+				let result = rem - 1;
+				rem = 0;
+				Some(result as u8)
+			} else {
+				rem = rem.saturating_sub(255);
+				Some(255)
+			}
+		} else {
+			None
+		}
+	};
+	first_byte.chain(core::iter::from_fn(next_bytes))
+}
+
+/// Encodes size and prefix to a stream output.
+fn encode_size_and_prefix<W>(size: usize, prefix: u8, prefix_mask: usize, out: &mut W)
+where
+	W: Output + ?Sized,
+{
+	for b in size_and_prefix_iterator(size, prefix, prefix_mask) {
+		out.push_byte(b)
+	}
+}
+
+/// Decode size only from stream input and header byte.
+fn decode_size(
+	first: u8,
+	input: &mut impl Input,
+	prefix_mask: usize,
+) -> Result<usize, codec::Error> {
+	let max_value = 255u8 >> prefix_mask;
+	let mut result = (first & max_value) as usize;
+	if result < max_value as usize {
+		return Ok(result);
+	}
+	result -= 1;
+	loop {
+		let n = input.read_byte()? as usize;
+		if n < 255 {
+			return Ok(result + n + 1);
+		}
+		result += 255;
+	}
+}

bizinikiwi/primitives/trie/src/proof_size_extension.rs

@@ -0,0 +1,497 @@
+// This file is part of Bizinikiwi.
+
+// Copyright (C) Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Externalities extension that provides access to the current proof size
+//! of the underlying recorder.
+
+use parking_lot::Mutex;
+
+use crate::ProofSizeProvider;
+use std::{collections::VecDeque, sync::Arc};
+
+pezsp_externalities::decl_extension! {
+	/// The proof size extension to fetch the current storage proof size
+	/// in externalities.
+	pub struct ProofSizeExt(Box<dyn ProofSizeProvider + 'static + Sync + Send>);
+
+	impl ProofSizeExt {
+		fn start_transaction(&mut self, ty: pezsp_externalities::TransactionType) {
+			self.0.start_transaction(ty.is_host());
+		}
+
+		fn rollback_transaction(&mut self, ty: pezsp_externalities::TransactionType) {
+			self.0.rollback_transaction(ty.is_host());
+		}
+
+		fn commit_transaction(&mut self, ty: pezsp_externalities::TransactionType) {
+			self.0.commit_transaction(ty.is_host());
+		}
+	}
+}
+
+impl ProofSizeExt {
+	/// Creates a new instance of [`ProofSizeExt`].
+	pub fn new<T: ProofSizeProvider + Sync + Send + 'static>(recorder: T) -> Self {
+		ProofSizeExt(Box::new(recorder))
+	}
+
+	/// Returns the storage proof size.
+	pub fn storage_proof_size(&self) -> u64 {
+		self.0.estimate_encoded_size() as _
+	}
+}
+
+/// Proof size estimations as recorded by [`RecordingProofSizeProvider`].
+///
+/// Each item is the estimated proof size as observed when calling
+/// [`ProofSizeProvider::estimate_encoded_size`]. The items are ordered by their observation and
+/// need to be replayed in the exact same order.
+pub struct RecordedProofSizeEstimations(pub VecDeque<usize>);
+
+/// Inner structure of [`RecordingProofSizeProvider`].
+struct RecordingProofSizeProviderInner {
+	inner: Box<dyn ProofSizeProvider + Send + Sync>,
+	/// Stores the observed proof estimations (in order of observation) per transaction.
+	///
+	/// Last element of the outer vector is the active transaction.
+	proof_size_estimations: Vec<Vec<usize>>,
+}
+
+/// An implementation of [`ProofSizeProvider`] that records the return value of the calls to
+/// [`ProofSizeProvider::estimate_encoded_size`].
+///
+/// Wraps an inner [`ProofSizeProvider`] that is used to get the actual encoded size estimations.
+/// Each estimation is recorded in the order it was observed.
+#[derive(Clone)]
+pub struct RecordingProofSizeProvider {
+	inner: Arc<Mutex<RecordingProofSizeProviderInner>>,
+}
+
+impl RecordingProofSizeProvider {
+	/// Creates a new instance of [`RecordingProofSizeProvider`].
+	pub fn new<T: ProofSizeProvider + Sync + Send + 'static>(recorder: T) -> Self {
+		Self {
+			inner: Arc::new(Mutex::new(RecordingProofSizeProviderInner {
+				inner: Box::new(recorder),
+				// Init the always existing transaction.
+				proof_size_estimations: vec![Vec::new()],
+			})),
+		}
+	}
+
+	/// Returns the recorded estimations returned by each call to
+	/// [`Self::estimate_encoded_size`].
+	pub fn recorded_estimations(&self) -> Vec<usize> {
+		self.inner.lock().proof_size_estimations.iter().flatten().copied().collect()
+	}
+}
+
+impl ProofSizeProvider for RecordingProofSizeProvider {
+	fn estimate_encoded_size(&self) -> usize {
+		let mut inner = self.inner.lock();
+
+		let estimation = inner.inner.estimate_encoded_size();
+
+		inner
+			.proof_size_estimations
+			.last_mut()
+			.expect("There is always at least one transaction open; qed")
+			.push(estimation);
+
+		estimation
+	}
+
+	fn start_transaction(&mut self, is_host: bool) {
+		// We don't care about runtime transactions, because they are part of the consensus critical
+		// path, that will always deterministically call this code.
+		//
+		// For example a runtime execution is creating 10 runtime transaction and calling in every
+		// transaction the proof size estimation host function and 8 of these transactions are
+		// rolled back. We need to keep all the 10 estimations. When the runtime execution is
+		// replayed (by e.g. importing a block), we will deterministically again create 10 runtime
+		// executions and roll back 8. However, in between we require all 10 estimations as
+		// otherwise the execution would not be deterministically anymore.
+		//
+		// A host transaction is only rolled back while for example building a block and an
+		// extrinsic failed in the early checks in the runtime. In this case, the extrinsic will
+		// also never appear in a block and thus, will not need to be replayed later on.
+		if is_host {
+			self.inner.lock().proof_size_estimations.push(Default::default());
+		}
+	}
+
+	fn rollback_transaction(&mut self, is_host: bool) {
+		let mut inner = self.inner.lock();
+
+		// The host side transaction needs to be reverted, because this is only done when an
+		// entire execution is rolled back. So, the execution will never be part of the consensus
+		// critical path.
+		if is_host && inner.proof_size_estimations.len() > 1 {
+			inner.proof_size_estimations.pop();
+		}
+	}
+
+	fn commit_transaction(&mut self, is_host: bool) {
+		let mut inner = self.inner.lock();
+
+		if is_host && inner.proof_size_estimations.len() > 1 {
+			let last = inner
+				.proof_size_estimations
+				.pop()
+				.expect("There are more than one element in the vector; qed");
+
+			inner
+				.proof_size_estimations
+				.last_mut()
+				.expect("There are more than one element in the vector; qed")
+				.extend(last);
+		}
+	}
+}
+
+/// An implementation of [`ProofSizeProvider`] that replays estimations recorded by
+/// [`RecordingProofSizeProvider`].
+///
+/// The recorded estimations are removed as they are required by calls to
+/// [`Self::estimate_encoded_size`]. Will return `0` when all estimations are consumed.
+pub struct ReplayProofSizeProvider(Arc<Mutex<RecordedProofSizeEstimations>>);
+
+impl ReplayProofSizeProvider {
+	/// Creates a new instance from the given [`RecordedProofSizeEstimations`].
+	pub fn from_recorded(recorded: RecordedProofSizeEstimations) -> Self {
+		Self(Arc::new(Mutex::new(recorded)))
+	}
+}
+
+impl From<RecordedProofSizeEstimations> for ReplayProofSizeProvider {
+	fn from(value: RecordedProofSizeEstimations) -> Self {
+		Self::from_recorded(value)
+	}
+}
+
+impl ProofSizeProvider for ReplayProofSizeProvider {
+	fn estimate_encoded_size(&self) -> usize {
+		self.0.lock().0.pop_front().unwrap_or_default()
+	}
+}
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+	use std::sync::atomic::{AtomicUsize, Ordering};
+
+	// Mock ProofSizeProvider for testing
+	#[derive(Clone)]
+	struct MockProofSizeProvider {
+		size: Arc<AtomicUsize>,
+	}
+
+	impl MockProofSizeProvider {
+		fn new(initial_size: usize) -> Self {
+			Self { size: Arc::new(AtomicUsize::new(initial_size)) }
+		}
+
+		fn set_size(&self, new_size: usize) {
+			self.size.store(new_size, Ordering::Relaxed);
+		}
+	}
+
+	impl ProofSizeProvider for MockProofSizeProvider {
+		fn estimate_encoded_size(&self) -> usize {
+			self.size.load(Ordering::Relaxed)
+		}
+
+		fn start_transaction(&mut self, _is_host: bool) {}
+		fn rollback_transaction(&mut self, _is_host: bool) {}
+		fn commit_transaction(&mut self, _is_host: bool) {}
+	}
+
+	#[test]
+	fn recording_proof_size_provider_basic_functionality() {
+		let mock = MockProofSizeProvider::new(100);
+		let tracker = RecordingProofSizeProvider::new(mock.clone());
+
+		// Initial state - no estimations recorded yet
+		assert_eq!(tracker.recorded_estimations(), Vec::<usize>::new());
+
+		// Call estimate_encoded_size and verify it's recorded
+		let size = tracker.estimate_encoded_size();
+		assert_eq!(size, 100);
+		assert_eq!(tracker.recorded_estimations(), vec![100]);
+
+		// Change the mock size and call again
+		mock.set_size(200);
+		let size = tracker.estimate_encoded_size();
+		assert_eq!(size, 200);
+		assert_eq!(tracker.recorded_estimations(), vec![100, 200]);
+
+		// Multiple calls with same size
+		let size = tracker.estimate_encoded_size();
+		assert_eq!(size, 200);
+		assert_eq!(tracker.recorded_estimations(), vec![100, 200, 200]);
+	}
+
+	#[test]
+	fn recording_proof_size_provider_host_transactions() {
+		let mock = MockProofSizeProvider::new(100);
+		let mut tracker = RecordingProofSizeProvider::new(mock.clone());
+
+		// Record some estimations in the initial transaction
+		tracker.estimate_encoded_size();
+		tracker.estimate_encoded_size();
+		assert_eq!(tracker.recorded_estimations(), vec![100, 100]);
+
+		// Start a host transaction
+		tracker.start_transaction(true);
+		mock.set_size(200);
+		tracker.estimate_encoded_size();
+
+		// Should have 3 estimations total
+		assert_eq!(tracker.recorded_estimations(), vec![100, 100, 200]);
+
+		// Commit the host transaction
+		tracker.commit_transaction(true);
+
+		// All estimations should still be there
+		assert_eq!(tracker.recorded_estimations(), vec![100, 100, 200]);
+
+		// Add more estimations
+		mock.set_size(300);
+		tracker.estimate_encoded_size();
+		assert_eq!(tracker.recorded_estimations(), vec![100, 100, 200, 300]);
+	}
+
+	#[test]
+	fn recording_proof_size_provider_host_transaction_rollback() {
+		let mock = MockProofSizeProvider::new(100);
+		let mut tracker = RecordingProofSizeProvider::new(mock.clone());
+
+		// Record some estimations in the initial transaction
+		tracker.estimate_encoded_size();
+		assert_eq!(tracker.recorded_estimations(), vec![100]);
+
+		// Start a host transaction
+		tracker.start_transaction(true);
+		mock.set_size(200);
+		tracker.estimate_encoded_size();
+		tracker.estimate_encoded_size();
+
+		// Should have 3 estimations total
+		assert_eq!(tracker.recorded_estimations(), vec![100, 200, 200]);
+
+		// Rollback the host transaction
+		tracker.rollback_transaction(true);
+
+		// Should only have the original estimation
+		assert_eq!(tracker.recorded_estimations(), vec![100]);
+	}
+
+	#[test]
+	fn recording_proof_size_provider_runtime_transactions_ignored() {
+		let mock = MockProofSizeProvider::new(100);
+		let mut tracker = RecordingProofSizeProvider::new(mock.clone());
+
+		// Record initial estimation
+		tracker.estimate_encoded_size();
+		assert_eq!(tracker.recorded_estimations(), vec![100]);
+
+		// Start a runtime transaction (is_host = false)
+		tracker.start_transaction(false);
+		mock.set_size(200);
+		tracker.estimate_encoded_size();
+
+		// Should have both estimations
+		assert_eq!(tracker.recorded_estimations(), vec![100, 200]);
+
+		// Commit runtime transaction - should not affect recording
+		tracker.commit_transaction(false);
+		assert_eq!(tracker.recorded_estimations(), vec![100, 200]);
+
+		// Rollback runtime transaction - should not affect recording
+		tracker.rollback_transaction(false);
+		assert_eq!(tracker.recorded_estimations(), vec![100, 200]);
+	}
+
+	#[test]
+	fn recording_proof_size_provider_nested_host_transactions() {
+		let mock = MockProofSizeProvider::new(100);
+		let mut tracker = RecordingProofSizeProvider::new(mock.clone());
+
+		// Initial estimation
+		tracker.estimate_encoded_size();
+		assert_eq!(tracker.recorded_estimations(), vec![100]);
+
+		// Start first host transaction
+		tracker.start_transaction(true);
+		mock.set_size(200);
+		tracker.estimate_encoded_size();
+
+		// Start nested host transaction
+		tracker.start_transaction(true);
+		mock.set_size(300);
+		tracker.estimate_encoded_size();
+
+		assert_eq!(tracker.recorded_estimations(), vec![100, 200, 300]);
+
+		// Commit nested transaction
+		tracker.commit_transaction(true);
+		assert_eq!(tracker.recorded_estimations(), vec![100, 200, 300]);
+
+		// Commit outer transaction
+		tracker.commit_transaction(true);
+		assert_eq!(tracker.recorded_estimations(), vec![100, 200, 300]);
+	}
+
+	#[test]
+	fn recording_proof_size_provider_nested_host_transaction_rollback() {
+		let mock = MockProofSizeProvider::new(100);
+		let mut tracker = RecordingProofSizeProvider::new(mock.clone());
+
+		// Initial estimation
+		tracker.estimate_encoded_size();
+
+		// Start first host transaction
+		tracker.start_transaction(true);
+		mock.set_size(200);
+		tracker.estimate_encoded_size();
+
+		// Start nested host transaction
+		tracker.start_transaction(true);
+		mock.set_size(300);
+		tracker.estimate_encoded_size();
+
+		assert_eq!(tracker.recorded_estimations(), vec![100, 200, 300]);
+
+		// Rollback nested transaction
+		tracker.rollback_transaction(true);
+		assert_eq!(tracker.recorded_estimations(), vec![100, 200]);
+
+		// Rollback outer transaction
+		tracker.rollback_transaction(true);
+		assert_eq!(tracker.recorded_estimations(), vec![100]);
+	}
+
+	#[test]
+	fn recording_proof_size_provider_rollback_on_base_transaction_does_nothing() {
+		let mock = MockProofSizeProvider::new(100);
+		let mut tracker = RecordingProofSizeProvider::new(mock.clone());
+
+		// Record some estimations
+		tracker.estimate_encoded_size();
+		tracker.estimate_encoded_size();
+		assert_eq!(tracker.recorded_estimations(), vec![100, 100]);
+
+		// Try to rollback the base transaction - should do nothing
+		tracker.rollback_transaction(true);
+		assert_eq!(tracker.recorded_estimations(), vec![100, 100]);
+	}
+
+	#[test]
+	fn recorded_proof_size_estimations_struct() {
+		let estimations = vec![100, 200, 300];
+		let recorded = RecordedProofSizeEstimations(estimations.into());
+		let expected: VecDeque<usize> = vec![100, 200, 300].into();
+		assert_eq!(recorded.0, expected);
+	}
+
+	#[test]
+	fn replay_proof_size_provider_basic_functionality() {
+		let estimations = vec![100, 200, 300, 150];
+		let recorded = RecordedProofSizeEstimations(estimations.into());
+		let replay = ReplayProofSizeProvider::from_recorded(recorded);
+
+		// Should replay estimations in order
+		assert_eq!(replay.estimate_encoded_size(), 100);
+		assert_eq!(replay.estimate_encoded_size(), 200);
+		assert_eq!(replay.estimate_encoded_size(), 300);
+		assert_eq!(replay.estimate_encoded_size(), 150);
+	}
+
+	#[test]
+	fn replay_proof_size_provider_exhausted_returns_zero() {
+		let estimations = vec![100, 200];
+		let recorded = RecordedProofSizeEstimations(estimations.into());
+		let replay = ReplayProofSizeProvider::from_recorded(recorded);
+
+		// Consume all estimations
+		assert_eq!(replay.estimate_encoded_size(), 100);
+		assert_eq!(replay.estimate_encoded_size(), 200);
+
+		// Should return 0 when exhausted
+		assert_eq!(replay.estimate_encoded_size(), 0);
+		assert_eq!(replay.estimate_encoded_size(), 0);
+	}
+
+	#[test]
+	fn replay_proof_size_provider_empty_returns_zero() {
+		let recorded = RecordedProofSizeEstimations(VecDeque::new());
+		let replay = ReplayProofSizeProvider::from_recorded(recorded);
+
+		// Should return 0 for empty estimations
+		assert_eq!(replay.estimate_encoded_size(), 0);
+		assert_eq!(replay.estimate_encoded_size(), 0);
+	}
+
+	#[test]
+	fn replay_proof_size_provider_from_trait() {
+		let estimations = vec![42, 84];
+		let recorded = RecordedProofSizeEstimations(estimations.into());
+		let replay: ReplayProofSizeProvider = recorded.into();
+
+		assert_eq!(replay.estimate_encoded_size(), 42);
+		assert_eq!(replay.estimate_encoded_size(), 84);
+		assert_eq!(replay.estimate_encoded_size(), 0);
+	}
+
+	#[test]
+	fn record_and_replay_integration() {
+		let mock = MockProofSizeProvider::new(100);
+		let recorder = RecordingProofSizeProvider::new(mock.clone());
+
+		// Record some estimations
+		recorder.estimate_encoded_size();
+		mock.set_size(200);
+		recorder.estimate_encoded_size();
+		mock.set_size(300);
+		recorder.estimate_encoded_size();
+
+		// Get recorded estimations
+		let recorded_estimations = recorder.recorded_estimations();
+		assert_eq!(recorded_estimations, vec![100, 200, 300]);
+
+		// Create replay provider from recorded estimations
+		let recorded = RecordedProofSizeEstimations(recorded_estimations.into());
+		let replay = ReplayProofSizeProvider::from_recorded(recorded);
+
+		// Replay should return the same sequence
+		assert_eq!(replay.estimate_encoded_size(), 100);
+		assert_eq!(replay.estimate_encoded_size(), 200);
+		assert_eq!(replay.estimate_encoded_size(), 300);
+		assert_eq!(replay.estimate_encoded_size(), 0);
+	}
+
+	#[test]
+	fn replay_proof_size_provider_single_value() {
+		let estimations = vec![42];
+		let recorded = RecordedProofSizeEstimations(estimations.into());
+		let replay = ReplayProofSizeProvider::from_recorded(recorded);
+
+		// Should return the single value then default to 0
+		assert_eq!(replay.estimate_encoded_size(), 42);
+		assert_eq!(replay.estimate_encoded_size(), 0);
+	}
+}

bizinikiwi/primitives/trie/src/recorder.rs

@@ -0,0 +1,883 @@
+// This file is part of Bizinikiwi.
+
+// Copyright (C) Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Trie recorder
+//!
+//! Provides an implementation of the [`TrieRecorder`](trie_db::TrieRecorder) trait. It can be used
+//! to record storage accesses to the state to generate a [`StorageProof`].
+
+use crate::{GenericMemoryDB, NodeCodec, StorageProof};
+use codec::Encode;
+use hash_db::Hasher;
+use memory_db::KeyFunction;
+use parking_lot::{Mutex, MutexGuard};
+use std::{
+	collections::{HashMap, HashSet},
+	marker::PhantomData,
+	mem,
+	ops::DerefMut,
+	sync::{
+		atomic::{AtomicUsize, Ordering},
+		Arc,
+	},
+};
+use trie_db::{RecordedForKey, TrieAccess};
+
+const LOG_TARGET: &str = "trie-recorder";
+
+/// A list of ignored nodes for [`Recorder`].
+///
+/// These nodes when passed to a recorder will be ignored and not recorded by the recorder.
+#[derive(Clone)]
+pub struct IgnoredNodes<H> {
+	nodes: HashSet<H>,
+}
+
+impl<H> Default for IgnoredNodes<H> {
+	fn default() -> Self {
+		Self { nodes: HashSet::default() }
+	}
+}
+
+impl<H: Eq + std::hash::Hash + Clone> IgnoredNodes<H> {
+	/// Initialize from the given storage proof.
+	///
+	/// So, all recorded nodes of the proof will be the ignored nodes.
+	pub fn from_storage_proof<Hasher: trie_db::Hasher<Out = H>>(proof: &StorageProof) -> Self {
+		Self { nodes: proof.iter_nodes().map(|n| Hasher::hash(&n)).collect() }
+	}
+
+	/// Initialize from the given memory db.
+	///
+	/// All nodes that have a reference count > 0 will be used as ignored nodes.
+	pub fn from_memory_db<Hasher: trie_db::Hasher<Out = H>, KF: KeyFunction<Hasher>>(
+		mut memory_db: GenericMemoryDB<Hasher, KF>,
+	) -> Self {
+		Self {
+			nodes: memory_db
+				.drain()
+				.into_iter()
+				// We do not want to add removed nodes.
+				.filter(|(_, (_, counter))| *counter > 0)
+				.map(|(_, (data, _))| Hasher::hash(&data))
+				.collect(),
+		}
+	}
+
+	/// Extend `self` with the other instance of ignored nodes.
+	pub fn extend(&mut self, other: Self) {
+		self.nodes.extend(other.nodes.into_iter());
+	}
+
+	/// Returns `true` if the node is ignored.
+	pub fn is_ignored(&self, node: &H) -> bool {
+		self.nodes.contains(node)
+	}
+}
+
+/// Stores all the information per transaction.
+#[derive(Default)]
+struct Transaction<H> {
+	/// Stores transaction information about [`RecorderInner::recorded_keys`].
+	///
+	/// For each transaction we only store the `storage_root` and the old states per key. `None`
+	/// state means that the key wasn't recorded before.
+	recorded_keys: HashMap<H, HashMap<Arc<[u8]>, Option<RecordedForKey>>>,
+	/// Stores transaction information about [`RecorderInner::accessed_nodes`].
+	///
+	/// For each transaction we only store the hashes of added nodes.
+	accessed_nodes: HashSet<H>,
+}
+
+/// The internals of [`Recorder`].
+struct RecorderInner<H> {
+	/// The keys for that we have recorded the trie nodes and if we have recorded up to the value.
+	///
+	/// Mapping: `StorageRoot -> (Key -> RecordedForKey)`.
+	recorded_keys: HashMap<H, HashMap<Arc<[u8]>, RecordedForKey>>,
+
+	/// Currently active transactions.
+	transactions: Vec<Transaction<H>>,
+
+	/// The encoded nodes we accessed while recording.
+	///
+	/// Mapping: `Hash(Node) -> Node`.
+	accessed_nodes: HashMap<H, Vec<u8>>,
+
+	/// Nodes that should be ignored and not recorded.
+	ignored_nodes: IgnoredNodes<H>,
+}
+
+impl<H> Default for RecorderInner<H> {
+	fn default() -> Self {
+		Self {
+			recorded_keys: Default::default(),
+			accessed_nodes: Default::default(),
+			transactions: Vec::new(),
+			ignored_nodes: Default::default(),
+		}
+	}
+}
+
+/// The trie recorder.
+///
+/// Owns the recorded data. Is used to transform data into a storage
+/// proof and to provide transaction support. The `as_trie_recorder` method provides a
+/// [`trie_db::TrieDB`] compatible recorder that implements the actual recording logic.
+pub struct Recorder<H: Hasher> {
+	inner: Arc<Mutex<RecorderInner<H::Out>>>,
+	/// The estimated encoded size of the storage proof this recorder will produce.
+	///
+	/// We store this in an atomic to be able to fetch the value while the `inner` is may locked.
+	encoded_size_estimation: Arc<AtomicUsize>,
+}
+
+impl<H: Hasher> Default for Recorder<H> {
+	fn default() -> Self {
+		Self { inner: Default::default(), encoded_size_estimation: Arc::new(0.into()) }
+	}
+}
+
+impl<H: Hasher> Clone for Recorder<H> {
+	fn clone(&self) -> Self {
+		Self {
+			inner: self.inner.clone(),
+			encoded_size_estimation: self.encoded_size_estimation.clone(),
+		}
+	}
+}
+
+impl<H: Hasher> Recorder<H> {
+	/// Create a new instance with the given `ingored_nodes`.
+	///
+	/// These ignored nodes are not recorded when accessed.
+	pub fn with_ignored_nodes(ignored_nodes: IgnoredNodes<H::Out>) -> Self {
+		Self {
+			inner: Arc::new(Mutex::new(RecorderInner { ignored_nodes, ..Default::default() })),
+			..Default::default()
+		}
+	}
+
+	/// Returns [`RecordedForKey`] per recorded key per trie.
+	///
+	/// There are multiple tries when working with e.g. child tries.
+	pub fn recorded_keys(&self) -> HashMap<H::Out, HashMap<Arc<[u8]>, RecordedForKey>> {
+		self.inner.lock().recorded_keys.clone()
+	}
+
+	/// Returns the recorder as [`TrieRecorder`](trie_db::TrieRecorder) compatible type.
+	///
+	/// - `storage_root`: The storage root of the trie for which accesses are recorded. This is
+	///   important when recording access to different tries at once (like top and child tries).
+	///
+	/// NOTE: This locks a mutex that stays locked until the return value is dropped.
+	#[inline]
+	pub fn as_trie_recorder(&self, storage_root: H::Out) -> TrieRecorder<'_, H> {
+		TrieRecorder::<H> {
+			inner: self.inner.lock(),
+			storage_root,
+			encoded_size_estimation: self.encoded_size_estimation.clone(),
+			_phantom: PhantomData,
+		}
+	}
+
+	/// Drain the recording into a [`StorageProof`].
+	///
+	/// While a recorder can be cloned, all share the same internal state. After calling this
+	/// function, all other instances will have their internal state reset as well.
+	///
+	/// If you don't want to drain the recorded state, use [`Self::to_storage_proof`].
+	///
+	/// Returns the [`StorageProof`].
+	pub fn drain_storage_proof(self) -> StorageProof {
+		let mut recorder = mem::take(&mut *self.inner.lock());
+		StorageProof::new(recorder.accessed_nodes.drain().map(|(_, v)| v))
+	}
+
+	/// Convert the recording to a [`StorageProof`].
+	///
+	/// In contrast to [`Self::drain_storage_proof`] this doesn't consume and doesn't clear the
+	/// recordings.
+	///
+	/// Returns the [`StorageProof`].
+	pub fn to_storage_proof(&self) -> StorageProof {
+		let recorder = self.inner.lock();
+		StorageProof::new(recorder.accessed_nodes.values().cloned())
+	}
+
+	/// Returns the estimated encoded size of the proof.
+	///
+	/// The estimation is based on all the nodes that were accessed until now while
+	/// accessing the trie.
+	pub fn estimate_encoded_size(&self) -> usize {
+		self.encoded_size_estimation.load(Ordering::Relaxed)
+	}
+
+	/// Reset the state.
+	///
+	/// This discards all recorded data.
+	pub fn reset(&self) {
+		mem::take(&mut *self.inner.lock());
+		self.encoded_size_estimation.store(0, Ordering::Relaxed);
+	}
+
+	/// Start a new transaction.
+	pub fn start_transaction(&self) {
+		let mut inner = self.inner.lock();
+		inner.transactions.push(Default::default());
+	}
+
+	/// Rollback the latest transaction.
+	///
+	/// Returns an error if there wasn't any active transaction.
+	pub fn rollback_transaction(&self) -> Result<(), ()> {
+		let mut inner = self.inner.lock();
+
+		// We locked `inner` and can just update the encoded size locally and then store it back to
+		// the atomic.
+		let mut new_encoded_size_estimation = self.encoded_size_estimation.load(Ordering::Relaxed);
+		let transaction = inner.transactions.pop().ok_or(())?;
+
+		transaction.accessed_nodes.into_iter().for_each(|n| {
+			if let Some(old) = inner.accessed_nodes.remove(&n) {
+				new_encoded_size_estimation =
+					new_encoded_size_estimation.saturating_sub(old.encoded_size());
+			}
+		});
+
+		transaction.recorded_keys.into_iter().for_each(|(storage_root, keys)| {
+			keys.into_iter().for_each(|(k, old_state)| {
+				if let Some(state) = old_state {
+					inner.recorded_keys.entry(storage_root).or_default().insert(k, state);
+				} else {
+					inner.recorded_keys.entry(storage_root).or_default().remove(&k);
+				}
+			});
+		});
+
+		self.encoded_size_estimation
+			.store(new_encoded_size_estimation, Ordering::Relaxed);
+
+		Ok(())
+	}
+
+	/// Commit the latest transaction.
+	///
+	/// Returns an error if there wasn't any active transaction.
+	pub fn commit_transaction(&self) -> Result<(), ()> {
+		let mut inner = self.inner.lock();
+
+		let transaction = inner.transactions.pop().ok_or(())?;
+
+		if let Some(parent_transaction) = inner.transactions.last_mut() {
+			parent_transaction.accessed_nodes.extend(transaction.accessed_nodes);
+
+			transaction.recorded_keys.into_iter().for_each(|(storage_root, keys)| {
+				keys.into_iter().for_each(|(k, old_state)| {
+					parent_transaction
+						.recorded_keys
+						.entry(storage_root)
+						.or_default()
+						.entry(k)
+						.or_insert(old_state);
+				})
+			});
+		}
+
+		Ok(())
+	}
+}
+
+impl<H: Hasher> crate::ProofSizeProvider for Recorder<H> {
+	fn estimate_encoded_size(&self) -> usize {
+		Recorder::estimate_encoded_size(self)
+	}
+}
+
+/// The [`TrieRecorder`](trie_db::TrieRecorder) implementation.
+pub struct TrieRecorder<'a, H: Hasher> {
+	inner: MutexGuard<'a, RecorderInner<H::Out>>,
+	storage_root: H::Out,
+	encoded_size_estimation: Arc<AtomicUsize>,
+	_phantom: PhantomData<H>,
+}
+
+impl<H: Hasher> crate::TrieRecorderProvider<H> for Recorder<H> {
+	type Recorder<'a>
+		= TrieRecorder<'a, H>
+	where
+		H: 'a;
+
+	fn drain_storage_proof(self) -> Option<StorageProof> {
+		Some(Recorder::drain_storage_proof(self))
+	}
+
+	fn as_trie_recorder(&self, storage_root: H::Out) -> Self::Recorder<'_> {
+		Recorder::as_trie_recorder(&self, storage_root)
+	}
+}
+
+impl<'a, H: Hasher> TrieRecorder<'a, H> {
+	/// Update the recorded keys entry for the given `full_key`.
+	fn update_recorded_keys(&mut self, full_key: &[u8], access: RecordedForKey) {
+		let inner = self.inner.deref_mut();
+
+		let entry =
+			inner.recorded_keys.entry(self.storage_root).or_default().entry(full_key.into());
+
+		let key = entry.key().clone();
+
+		// We don't need to update the record if we only accessed the `Hash` for the given
+		// `full_key`. Only `Value` access can be an upgrade from `Hash`.
+		let entry = if matches!(access, RecordedForKey::Value) {
+			entry.and_modify(|e| {
+				if let Some(tx) = inner.transactions.last_mut() {
+					// Store the previous state only once per transaction.
+					tx.recorded_keys
+						.entry(self.storage_root)
+						.or_default()
+						.entry(key.clone())
+						.or_insert(Some(*e));
+				}
+
+				*e = access;
+			})
+		} else {
+			entry
+		};
+
+		entry.or_insert_with(|| {
+			if let Some(tx) = inner.transactions.last_mut() {
+				// The key wasn't yet recorded, so there isn't any old state.
+				tx.recorded_keys
+					.entry(self.storage_root)
+					.or_default()
+					.entry(key)
+					.or_insert(None);
+			}
+
+			access
+		});
+	}
+}
+
+impl<'a, H: Hasher> trie_db::TrieRecorder<H::Out> for TrieRecorder<'a, H> {
+	fn record(&mut self, access: TrieAccess<H::Out>) {
+		let mut encoded_size_update = 0;
+
+		match access {
+			TrieAccess::NodeOwned { hash, node_owned } => {
+				let inner = self.inner.deref_mut();
+
+				if inner.ignored_nodes.is_ignored(&hash) {
+					tracing::trace!(
+						target: LOG_TARGET,
+						?hash,
+						"Ignoring node",
+					);
+					return;
+				}
+
+				tracing::trace!(
+					target: LOG_TARGET,
+					?hash,
+					"Recording node",
+				);
+
+				inner.accessed_nodes.entry(hash).or_insert_with(|| {
+					let node = node_owned.to_encoded::<NodeCodec<H>>();
+
+					encoded_size_update += node.encoded_size();
+
+					if let Some(tx) = inner.transactions.last_mut() {
+						tx.accessed_nodes.insert(hash);
+					}
+
+					node
+				});
+			},
+			TrieAccess::EncodedNode { hash, encoded_node } => {
+				let inner = self.inner.deref_mut();
+
+				if inner.ignored_nodes.is_ignored(&hash) {
+					tracing::trace!(
+						target: LOG_TARGET,
+						?hash,
+						"Ignoring node",
+					);
+					return;
+				}
+
+				tracing::trace!(
+					target: LOG_TARGET,
+					hash = ?hash,
+					"Recording node",
+				);
+
+				inner.accessed_nodes.entry(hash).or_insert_with(|| {
+					let node = encoded_node.into_owned();
+
+					encoded_size_update += node.encoded_size();
+
+					if let Some(tx) = inner.transactions.last_mut() {
+						tx.accessed_nodes.insert(hash);
+					}
+
+					node
+				});
+			},
+			TrieAccess::Value { hash, value, full_key } => {
+				let inner = self.inner.deref_mut();
+
+				// A value is also just a node.
+				if inner.ignored_nodes.is_ignored(&hash) {
+					tracing::trace!(
+						target: LOG_TARGET,
+						?hash,
+						"Ignoring value",
+					);
+					return;
+				}
+
+				tracing::trace!(
+					target: LOG_TARGET,
+					hash = ?hash,
+					key = ?pezsp_core::hexdisplay::HexDisplay::from(&full_key),
+					"Recording value",
+				);
+
+				inner.accessed_nodes.entry(hash).or_insert_with(|| {
+					let value = value.into_owned();
+
+					encoded_size_update += value.encoded_size();
+
+					if let Some(tx) = inner.transactions.last_mut() {
+						tx.accessed_nodes.insert(hash);
+					}
+
+					value
+				});
+
+				self.update_recorded_keys(full_key, RecordedForKey::Value);
+			},
+			TrieAccess::Hash { full_key } => {
+				tracing::trace!(
+					target: LOG_TARGET,
+					key = ?pezsp_core::hexdisplay::HexDisplay::from(&full_key),
+					"Recorded hash access for key",
+				);
+
+				// We don't need to update the `encoded_size_update` as the hash was already
+				// accounted for by the recorded node that holds the hash.
+				self.update_recorded_keys(full_key, RecordedForKey::Hash);
+			},
+			TrieAccess::NonExisting { full_key } => {
+				tracing::trace!(
+					target: LOG_TARGET,
+					key = ?pezsp_core::hexdisplay::HexDisplay::from(&full_key),
+					"Recorded non-existing value access for key",
+				);
+
+				// Non-existing access means we recorded all trie nodes up to the value.
+				// Not the actual value, as it doesn't exist, but all trie nodes to know
+				// that the value doesn't exist in the trie.
+				self.update_recorded_keys(full_key, RecordedForKey::Value);
+			},
+			TrieAccess::InlineValue { full_key } => {
+				tracing::trace!(
+					target: LOG_TARGET,
+					key = ?pezsp_core::hexdisplay::HexDisplay::from(&full_key),
+					"Recorded inline value access for key",
+				);
+
+				// A value was accessed that is stored inline a node and we recorded all trie nodes
+				// to access this value.
+				self.update_recorded_keys(full_key, RecordedForKey::Value);
+			},
+		};
+
+		self.encoded_size_estimation.fetch_add(encoded_size_update, Ordering::Relaxed);
+	}
+
+	fn trie_nodes_recorded_for_key(&self, key: &[u8]) -> RecordedForKey {
+		self.inner
+			.recorded_keys
+			.get(&self.storage_root)
+			.and_then(|k| k.get(key).copied())
+			.unwrap_or(RecordedForKey::None)
+	}
+}
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+	use crate::tests::create_trie;
+	use trie_db::{Trie, TrieDBBuilder, TrieRecorder};
+
+	type MemoryDB = crate::MemoryDB<pezsp_core::Blake2Hasher>;
+	type Layout = crate::LayoutV1<pezsp_core::Blake2Hasher>;
+	type Recorder = super::Recorder<pezsp_core::Blake2Hasher>;
+
+	const TEST_DATA: &[(&[u8], &[u8])] =
+		&[(b"key1", &[1; 64]), (b"key2", &[2; 64]), (b"key3", &[3; 64]), (b"key4", &[4; 64])];
+
+	#[test]
+	fn recorder_works() {
+		let (db, root) = create_trie::<Layout>(TEST_DATA);
+
+		let recorder = Recorder::default();
+
+		{
+			let mut trie_recorder = recorder.as_trie_recorder(root);
+			let trie = TrieDBBuilder::<Layout>::new(&db, &root)
+				.with_recorder(&mut trie_recorder)
+				.build();
+			assert_eq!(TEST_DATA[0].1.to_vec(), trie.get(TEST_DATA[0].0).unwrap().unwrap());
+		}
+
+		let storage_proof = recorder.drain_storage_proof();
+		let memory_db: MemoryDB = storage_proof.into_memory_db();
+
+		// Check that we recorded the required data
+		let trie = TrieDBBuilder::<Layout>::new(&memory_db, &root).build();
+		assert_eq!(TEST_DATA[0].1.to_vec(), trie.get(TEST_DATA[0].0).unwrap().unwrap());
+	}
+
+	#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
+	struct RecorderStats {
+		accessed_nodes: usize,
+		recorded_keys: usize,
+		estimated_size: usize,
+	}
+
+	impl RecorderStats {
+		fn extract(recorder: &Recorder) -> Self {
+			let inner = recorder.inner.lock();
+
+			let recorded_keys =
+				inner.recorded_keys.iter().flat_map(|(_, keys)| keys.keys()).count();
+
+			Self {
+				recorded_keys,
+				accessed_nodes: inner.accessed_nodes.len(),
+				estimated_size: recorder.estimate_encoded_size(),
+			}
+		}
+	}
+
+	#[test]
+	fn recorder_transactions_rollback_work() {
+		let (db, root) = create_trie::<Layout>(TEST_DATA);
+
+		let recorder = Recorder::default();
+		let mut stats = vec![RecorderStats::default()];
+
+		for i in 0..4 {
+			recorder.start_transaction();
+			{
+				let mut trie_recorder = recorder.as_trie_recorder(root);
+				let trie = TrieDBBuilder::<Layout>::new(&db, &root)
+					.with_recorder(&mut trie_recorder)
+					.build();
+
+				assert_eq!(TEST_DATA[i].1.to_vec(), trie.get(TEST_DATA[i].0).unwrap().unwrap());
+			}
+			stats.push(RecorderStats::extract(&recorder));
+		}
+
+		assert_eq!(4, recorder.inner.lock().transactions.len());
+
+		for i in 0..5 {
+			assert_eq!(stats[4 - i], RecorderStats::extract(&recorder));
+
+			let storage_proof = recorder.to_storage_proof();
+			let memory_db: MemoryDB = storage_proof.into_memory_db();
+
+			// Check that we recorded the required data
+			let trie = TrieDBBuilder::<Layout>::new(&memory_db, &root).build();
+
+			// Check that the required data is still present.
+			for a in 0..4 {
+				if a < 4 - i {
+					assert_eq!(TEST_DATA[a].1.to_vec(), trie.get(TEST_DATA[a].0).unwrap().unwrap());
+				} else {
+					// All the data that we already rolled back, should be gone!
+					assert!(trie.get(TEST_DATA[a].0).is_err());
+				}
+			}
+
+			if i < 4 {
+				recorder.rollback_transaction().unwrap();
+			}
+		}
+
+		assert_eq!(0, recorder.inner.lock().transactions.len());
+	}
+
+	#[test]
+	fn recorder_transactions_commit_work() {
+		let (db, root) = create_trie::<Layout>(TEST_DATA);
+
+		let recorder = Recorder::default();
+
+		for i in 0..4 {
+			recorder.start_transaction();
+			{
+				let mut trie_recorder = recorder.as_trie_recorder(root);
+				let trie = TrieDBBuilder::<Layout>::new(&db, &root)
+					.with_recorder(&mut trie_recorder)
+					.build();
+
+				assert_eq!(TEST_DATA[i].1.to_vec(), trie.get(TEST_DATA[i].0).unwrap().unwrap());
+			}
+		}
+
+		let stats = RecorderStats::extract(&recorder);
+		assert_eq!(4, recorder.inner.lock().transactions.len());
+
+		for _ in 0..4 {
+			recorder.commit_transaction().unwrap();
+		}
+		assert_eq!(0, recorder.inner.lock().transactions.len());
+		assert_eq!(stats, RecorderStats::extract(&recorder));
+
+		let storage_proof = recorder.to_storage_proof();
+		let memory_db: MemoryDB = storage_proof.into_memory_db();
+
+		// Check that we recorded the required data
+		let trie = TrieDBBuilder::<Layout>::new(&memory_db, &root).build();
+
+		// Check that the required data is still present.
+		for i in 0..4 {
+			assert_eq!(TEST_DATA[i].1.to_vec(), trie.get(TEST_DATA[i].0).unwrap().unwrap());
+		}
+	}
+
+	#[test]
+	fn recorder_transactions_commit_and_rollback_work() {
+		let (db, root) = create_trie::<Layout>(TEST_DATA);
+
+		let recorder = Recorder::default();
+
+		for i in 0..2 {
+			recorder.start_transaction();
+			{
+				let mut trie_recorder = recorder.as_trie_recorder(root);
+				let trie = TrieDBBuilder::<Layout>::new(&db, &root)
+					.with_recorder(&mut trie_recorder)
+					.build();
+
+				assert_eq!(TEST_DATA[i].1.to_vec(), trie.get(TEST_DATA[i].0).unwrap().unwrap());
+			}
+		}
+
+		recorder.rollback_transaction().unwrap();
+
+		for i in 2..4 {
+			recorder.start_transaction();
+			{
+				let mut trie_recorder = recorder.as_trie_recorder(root);
+				let trie = TrieDBBuilder::<Layout>::new(&db, &root)
+					.with_recorder(&mut trie_recorder)
+					.build();
+
+				assert_eq!(TEST_DATA[i].1.to_vec(), trie.get(TEST_DATA[i].0).unwrap().unwrap());
+			}
+		}
+
+		recorder.rollback_transaction().unwrap();
+
+		assert_eq!(2, recorder.inner.lock().transactions.len());
+
+		for _ in 0..2 {
+			recorder.commit_transaction().unwrap();
+		}
+
+		assert_eq!(0, recorder.inner.lock().transactions.len());
+
+		let storage_proof = recorder.to_storage_proof();
+		let memory_db: MemoryDB = storage_proof.into_memory_db();
+
+		// Check that we recorded the required data
+		let trie = TrieDBBuilder::<Layout>::new(&memory_db, &root).build();
+
+		// Check that the required data is still present.
+		for i in 0..4 {
+			if i % 2 == 0 {
+				assert_eq!(TEST_DATA[i].1.to_vec(), trie.get(TEST_DATA[i].0).unwrap().unwrap());
+			} else {
+				assert!(trie.get(TEST_DATA[i].0).is_err());
+			}
+		}
+	}
+
+	#[test]
+	fn recorder_transaction_accessed_keys_works() {
+		let key = TEST_DATA[0].0;
+		let (db, root) = create_trie::<Layout>(TEST_DATA);
+
+		let recorder = Recorder::default();
+
+		{
+			let trie_recorder = recorder.as_trie_recorder(root);
+			assert!(matches!(trie_recorder.trie_nodes_recorded_for_key(key), RecordedForKey::None));
+		}
+
+		recorder.start_transaction();
+		{
+			let mut trie_recorder = recorder.as_trie_recorder(root);
+			let trie = TrieDBBuilder::<Layout>::new(&db, &root)
+				.with_recorder(&mut trie_recorder)
+				.build();
+
+			assert_eq!(
+				pezsp_core::Blake2Hasher::hash(TEST_DATA[0].1),
+				trie.get_hash(TEST_DATA[0].0).unwrap().unwrap()
+			);
+			assert!(matches!(trie_recorder.trie_nodes_recorded_for_key(key), RecordedForKey::Hash));
+		}
+
+		recorder.start_transaction();
+		{
+			let mut trie_recorder = recorder.as_trie_recorder(root);
+			let trie = TrieDBBuilder::<Layout>::new(&db, &root)
+				.with_recorder(&mut trie_recorder)
+				.build();
+
+			assert_eq!(TEST_DATA[0].1.to_vec(), trie.get(TEST_DATA[0].0).unwrap().unwrap());
+			assert!(matches!(
+				trie_recorder.trie_nodes_recorded_for_key(key),
+				RecordedForKey::Value,
+			));
+		}
+
+		recorder.rollback_transaction().unwrap();
+		{
+			let trie_recorder = recorder.as_trie_recorder(root);
+			assert!(matches!(trie_recorder.trie_nodes_recorded_for_key(key), RecordedForKey::Hash));
+		}
+
+		recorder.rollback_transaction().unwrap();
+		{
+			let trie_recorder = recorder.as_trie_recorder(root);
+			assert!(matches!(trie_recorder.trie_nodes_recorded_for_key(key), RecordedForKey::None));
+		}
+
+		recorder.start_transaction();
+		{
+			let mut trie_recorder = recorder.as_trie_recorder(root);
+			let trie = TrieDBBuilder::<Layout>::new(&db, &root)
+				.with_recorder(&mut trie_recorder)
+				.build();
+
+			assert_eq!(TEST_DATA[0].1.to_vec(), trie.get(TEST_DATA[0].0).unwrap().unwrap());
+			assert!(matches!(
+				trie_recorder.trie_nodes_recorded_for_key(key),
+				RecordedForKey::Value,
+			));
+		}
+
+		recorder.start_transaction();
+		{
+			let mut trie_recorder = recorder.as_trie_recorder(root);
+			let trie = TrieDBBuilder::<Layout>::new(&db, &root)
+				.with_recorder(&mut trie_recorder)
+				.build();
+
+			assert_eq!(
+				pezsp_core::Blake2Hasher::hash(TEST_DATA[0].1),
+				trie.get_hash(TEST_DATA[0].0).unwrap().unwrap()
+			);
+			assert!(matches!(
+				trie_recorder.trie_nodes_recorded_for_key(key),
+				RecordedForKey::Value
+			));
+		}
+
+		recorder.rollback_transaction().unwrap();
+		{
+			let trie_recorder = recorder.as_trie_recorder(root);
+			assert!(matches!(
+				trie_recorder.trie_nodes_recorded_for_key(key),
+				RecordedForKey::Value
+			));
+		}
+
+		recorder.rollback_transaction().unwrap();
+		{
+			let trie_recorder = recorder.as_trie_recorder(root);
+			assert!(matches!(trie_recorder.trie_nodes_recorded_for_key(key), RecordedForKey::None));
+		}
+	}
+
+	#[test]
+	fn recorder_ignoring_nodes_works() {
+		let (db, root) = create_trie::<Layout>(TEST_DATA);
+
+		let recorder = Recorder::default();
+
+		{
+			let mut trie_recorder = recorder.as_trie_recorder(root);
+			let trie = TrieDBBuilder::<Layout>::new(&db, &root)
+				.with_recorder(&mut trie_recorder)
+				.build();
+
+			for (key, data) in TEST_DATA.iter().take(3) {
+				assert_eq!(data.to_vec(), trie.get(&key).unwrap().unwrap());
+			}
+		}
+
+		assert!(recorder.estimate_encoded_size() > 10);
+		let mut ignored_nodes = IgnoredNodes::from_storage_proof::<pezsp_core::Blake2Hasher>(
+			&recorder.drain_storage_proof(),
+		);
+
+		let recorder = Recorder::with_ignored_nodes(ignored_nodes.clone());
+
+		{
+			let mut trie_recorder = recorder.as_trie_recorder(root);
+			let trie = TrieDBBuilder::<Layout>::new(&db, &root)
+				.with_recorder(&mut trie_recorder)
+				.build();
+
+			for (key, data) in TEST_DATA {
+				assert_eq!(data.to_vec(), trie.get(&key).unwrap().unwrap());
+			}
+		}
+
+		assert!(recorder.estimate_encoded_size() > TEST_DATA[3].1.len());
+		let ignored_nodes2 = IgnoredNodes::from_storage_proof::<pezsp_core::Blake2Hasher>(
+			&recorder.drain_storage_proof(),
+		);
+
+		ignored_nodes.extend(ignored_nodes2);
+
+		let recorder = Recorder::with_ignored_nodes(ignored_nodes);
+
+		{
+			let mut trie_recorder = recorder.as_trie_recorder(root);
+			let trie = TrieDBBuilder::<Layout>::new(&db, &root)
+				.with_recorder(&mut trie_recorder)
+				.build();
+
+			for (key, data) in TEST_DATA {
+				assert_eq!(data.to_vec(), trie.get(&key).unwrap().unwrap());
+			}
+		}
+		assert_eq!(0, recorder.estimate_encoded_size());
+	}
+}

bizinikiwi/primitives/trie/src/recorder_ext.rs

@@ -0,0 +1,47 @@
+// This file is part of Bizinikiwi.
+
+// Copyright (C) Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Extension for the default recorder.
+
+use crate::RawStorageProof;
+use alloc::{collections::BTreeSet, vec::Vec};
+use trie_db::{Recorder, TrieLayout};
+
+/// Convenience extension for the `Recorder` struct.
+///
+/// Used to deduplicate some logic.
+pub trait RecorderExt<L: TrieLayout>
+where
+	Self: Sized,
+{
+	/// Convert the recorder into a `BTreeSet`.
+	fn into_set(self) -> BTreeSet<Vec<u8>>;
+
+	/// Convert the recorder into a `RawStorageProof`, avoiding duplicate nodes.
+	fn into_raw_storage_proof(self) -> RawStorageProof {
+		// The recorder may record the same trie node multiple times,
+		// and we don't want duplicate nodes in our proofs
+		// => let's deduplicate it by collecting to a BTreeSet first
+		self.into_set().into_iter().collect()
+	}
+}
+
+impl<L: TrieLayout> RecorderExt<L> for Recorder<L> {
+	fn into_set(mut self) -> BTreeSet<Vec<u8>> {
+		self.drain().into_iter().map(|record| record.data).collect::<BTreeSet<_>>()
+	}
+}

bizinikiwi/primitives/trie/src/storage_proof.rs

@@ -0,0 +1,256 @@
+// This file is part of Bizinikiwi.
+
+// Copyright (C) Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use alloc::{collections::btree_set::BTreeSet, vec::Vec};
+use codec::{Decode, DecodeWithMemTracking, Encode};
+use core::iter::{DoubleEndedIterator, IntoIterator};
+use hash_db::{HashDB, Hasher};
+use scale_info::TypeInfo;
+
+// Note that `LayoutV1` usage here (proof compaction) is compatible
+// with `LayoutV0`.
+use crate::LayoutV1 as Layout;
+
+/// Error associated with the `storage_proof` module.
+#[derive(Encode, Decode, Clone, Eq, PartialEq, Debug, TypeInfo)]
+pub enum StorageProofError {
+	/// The proof contains duplicate nodes.
+	DuplicateNodes,
+}
+
+/// A proof that some set of key-value pairs are included in the storage trie. The proof contains
+/// the storage values so that the partial storage backend can be reconstructed by a verifier that
+/// does not already have access to the key-value pairs.
+///
+/// The proof consists of the set of serialized nodes in the storage trie accessed when looking up
+/// the keys covered by the proof. Verifying the proof requires constructing the partial trie from
+/// the serialized nodes and performing the key lookups.
+#[derive(Debug, PartialEq, Eq, Clone, Encode, Decode, DecodeWithMemTracking, TypeInfo)]
+pub struct StorageProof {
+	trie_nodes: BTreeSet<Vec<u8>>,
+}
+
+impl StorageProof {
+	/// Constructs a storage proof from a subset of encoded trie nodes in a storage backend.
+	pub fn new(trie_nodes: impl IntoIterator<Item = Vec<u8>>) -> Self {
+		StorageProof { trie_nodes: BTreeSet::from_iter(trie_nodes) }
+	}
+
+	/// Constructs a storage proof from a subset of encoded trie nodes in a storage backend.
+	///
+	/// Returns an error if the provided subset of encoded trie nodes contains duplicates.
+	pub fn new_with_duplicate_nodes_check(
+		trie_nodes: impl IntoIterator<Item = Vec<u8>>,
+	) -> Result<Self, StorageProofError> {
+		let mut trie_nodes_set = BTreeSet::new();
+		for node in trie_nodes {
+			if !trie_nodes_set.insert(node) {
+				return Err(StorageProofError::DuplicateNodes);
+			}
+		}
+
+		Ok(StorageProof { trie_nodes: trie_nodes_set })
+	}
+
+	/// Returns a new empty proof.
+	///
+	/// An empty proof is capable of only proving trivial statements (ie. that an empty set of
+	/// key-value pairs exist in storage).
+	pub fn empty() -> Self {
+		StorageProof { trie_nodes: BTreeSet::new() }
+	}
+
+	/// Returns whether this is an empty proof.
+	pub fn is_empty(&self) -> bool {
+		self.trie_nodes.is_empty()
+	}
+
+	/// Returns the number of nodes in the proof.
+	pub fn len(&self) -> usize {
+		self.trie_nodes.len()
+	}
+
+	/// Convert into an iterator over encoded trie nodes in lexicographical order constructed
+	/// from the proof.
+	pub fn into_iter_nodes(self) -> impl Sized + DoubleEndedIterator<Item = Vec<u8>> {
+		self.trie_nodes.into_iter()
+	}
+
+	/// Create an iterator over encoded trie nodes in lexicographical order constructed
+	/// from the proof.
+	pub fn iter_nodes(&self) -> impl Sized + DoubleEndedIterator<Item = &Vec<u8>> {
+		self.trie_nodes.iter()
+	}
+
+	/// Convert into plain node vector.
+	pub fn into_nodes(self) -> BTreeSet<Vec<u8>> {
+		self.trie_nodes
+	}
+
+	/// Creates a [`MemoryDB`](crate::MemoryDB) from `Self`.
+	pub fn into_memory_db<H: Hasher>(self) -> crate::MemoryDB<H> {
+		self.into()
+	}
+
+	/// Creates a [`MemoryDB`](crate::MemoryDB) from `Self` reference.
+	pub fn to_memory_db<H: Hasher>(&self) -> crate::MemoryDB<H> {
+		self.into()
+	}
+
+	/// Merges multiple storage proofs covering potentially different sets of keys into one proof
+	/// covering all keys. The merged proof output may be smaller than the aggregate size of the
+	/// input proofs due to deduplication of trie nodes.
+	pub fn merge(proofs: impl IntoIterator<Item = Self>) -> Self {
+		let trie_nodes = proofs
+			.into_iter()
+			.flat_map(|proof| proof.into_iter_nodes())
+			.collect::<BTreeSet<_>>()
+			.into_iter()
+			.collect();
+
+		Self { trie_nodes }
+	}
+
+	/// Encode as a compact proof with default trie layout.
+	pub fn into_compact_proof<H: Hasher>(
+		self,
+		root: H::Out,
+	) -> Result<CompactProof, crate::CompactProofError<H::Out, crate::Error<H::Out>>> {
+		let db = self.into_memory_db();
+		crate::encode_compact::<Layout<H>, crate::MemoryDB<H>>(&db, &root)
+	}
+
+	/// Encode as a compact proof with default trie layout.
+	pub fn to_compact_proof<H: Hasher>(
+		&self,
+		root: H::Out,
+	) -> Result<CompactProof, crate::CompactProofError<H::Out, crate::Error<H::Out>>> {
+		let db = self.to_memory_db();
+		crate::encode_compact::<Layout<H>, crate::MemoryDB<H>>(&db, &root)
+	}
+
+	/// Returns the estimated encoded size of the compact proof.
+	///
+	/// Running this operation is a slow operation (build the whole compact proof) and should only
+	/// be in non sensitive path.
+	///
+	/// Return `None` on error.
+	pub fn encoded_compact_size<H: Hasher>(self, root: H::Out) -> Option<usize> {
+		let compact_proof = self.into_compact_proof::<H>(root);
+		compact_proof.ok().map(|p| p.encoded_size())
+	}
+}
+
+impl<H: Hasher> From<StorageProof> for crate::MemoryDB<H> {
+	fn from(proof: StorageProof) -> Self {
+		From::from(&proof)
+	}
+}
+
+impl<H: Hasher> From<&StorageProof> for crate::MemoryDB<H> {
+	fn from(proof: &StorageProof) -> Self {
+		let mut db = crate::MemoryDB::with_hasher(crate::RandomState::default());
+		proof.iter_nodes().for_each(|n| {
+			db.insert(crate::EMPTY_PREFIX, &n);
+		});
+		db
+	}
+}
+
+/// Storage proof in compact form.
+#[derive(Debug, PartialEq, Eq, Clone, Encode, Decode, TypeInfo)]
+pub struct CompactProof {
+	pub encoded_nodes: Vec<Vec<u8>>,
+}
+
+impl CompactProof {
+	/// Return an iterator on the compact encoded nodes.
+	pub fn iter_compact_encoded_nodes(&self) -> impl Iterator<Item = &[u8]> {
+		self.encoded_nodes.iter().map(Vec::as_slice)
+	}
+
+	/// Decode to a full storage_proof.
+	pub fn to_storage_proof<H: Hasher>(
+		&self,
+		expected_root: Option<&H::Out>,
+	) -> Result<(StorageProof, H::Out), crate::CompactProofError<H::Out, crate::Error<H::Out>>> {
+		let mut db = crate::MemoryDB::<H>::new(&[]);
+		let root = crate::decode_compact::<Layout<H>, _, _>(
+			&mut db,
+			self.iter_compact_encoded_nodes(),
+			expected_root,
+		)?;
+		Ok((
+			StorageProof::new(db.drain().into_iter().filter_map(|kv| {
+				if (kv.1).1 > 0 {
+					Some((kv.1).0)
+				} else {
+					None
+				}
+			})),
+			root,
+		))
+	}
+
+	/// Convert self into a [`MemoryDB`](crate::MemoryDB).
+	///
+	/// `expected_root` is the expected root of this compact proof.
+	///
+	/// Returns the memory db and the root of the trie.
+	pub fn to_memory_db<H: Hasher>(
+		&self,
+		expected_root: Option<&H::Out>,
+	) -> Result<(crate::MemoryDB<H>, H::Out), crate::CompactProofError<H::Out, crate::Error<H::Out>>>
+	{
+		let mut db = crate::MemoryDB::<H>::new(&[]);
+		let root = crate::decode_compact::<Layout<H>, _, _>(
+			&mut db,
+			self.iter_compact_encoded_nodes(),
+			expected_root,
+		)?;
+
+		Ok((db, root))
+	}
+}
+
+#[cfg(test)]
+pub mod tests {
+	use super::*;
+	use crate::{tests::create_storage_proof, StorageProof};
+
+	type Hasher = pezsp_core::Blake2Hasher;
+	type Layout = crate::LayoutV1<Hasher>;
+
+	const TEST_DATA: &[(&[u8], &[u8])] =
+		&[(b"key1", &[1; 64]), (b"key2", &[2; 64]), (b"key3", &[3; 64]), (b"key11", &[4; 64])];
+
+	#[test]
+	fn proof_with_duplicate_nodes_is_rejected() {
+		let (raw_proof, _root) = create_storage_proof::<Layout>(TEST_DATA);
+		assert!(matches!(
+			StorageProof::new_with_duplicate_nodes_check(raw_proof),
+			Err(StorageProofError::DuplicateNodes)
+		));
+	}
+
+	#[test]
+	fn invalid_compact_proof_does_not_panic_when_decoding() {
+		let invalid_proof = CompactProof { encoded_nodes: vec![vec![135]] };
+		let result = invalid_proof.to_memory_db::<Hasher>(None);
+		assert!(result.is_err());
+	}
+}

bizinikiwi/primitives/trie/src/trie_codec.rs

@@ -0,0 +1,439 @@
+// This file is part of Bizinikiwi.
+
+// Copyright (C) Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Compact proof support.
+//!
+//! This uses compact proof from trie crate and extends
+//! it to bizinikiwi specific layout and child trie system.
+
+use crate::{CompactProof, HashDBT, TrieConfiguration, TrieHash, EMPTY_PREFIX};
+use alloc::{boxed::Box, vec::Vec};
+use trie_db::{CError, Trie};
+
+/// Error for trie node decoding.
+#[derive(Debug)]
+#[cfg_attr(feature = "std", derive(thiserror::Error))]
+pub enum Error<H, CodecError> {
+	#[cfg_attr(feature = "std", error("Invalid root {0:x?}, expected {1:x?}"))]
+	RootMismatch(H, H),
+	#[cfg_attr(feature = "std", error("Missing nodes in the proof"))]
+	IncompleteProof,
+	#[cfg_attr(feature = "std", error("Child node content with no root in proof"))]
+	ExtraneousChildNode,
+	#[cfg_attr(feature = "std", error("Proof of child trie {0:x?} not in parent proof"))]
+	ExtraneousChildProof(H),
+	#[cfg_attr(feature = "std", error("Invalid root {0:x?}, expected {1:x?}"))]
+	InvalidChildRoot(Vec<u8>, Vec<u8>),
+	#[cfg_attr(feature = "std", error("Trie error: {0:?}"))]
+	TrieError(Box<trie_db::TrieError<H, CodecError>>),
+}
+
+impl<H, CodecError> From<Box<trie_db::TrieError<H, CodecError>>> for Error<H, CodecError> {
+	fn from(error: Box<trie_db::TrieError<H, CodecError>>) -> Self {
+		Error::TrieError(error)
+	}
+}
+
+/// Decode a compact proof.
+///
+/// Takes as input a destination `db` for decoded node and `encoded`
+/// an iterator of compact encoded nodes.
+///
+/// Child trie are decoded in order of child trie root present
+/// in the top trie.
+pub fn decode_compact<'a, L, DB, I>(
+	db: &mut DB,
+	encoded: I,
+	expected_root: Option<&TrieHash<L>>,
+) -> Result<TrieHash<L>, Error<TrieHash<L>, CError<L>>>
+where
+	L: TrieConfiguration,
+	DB: HashDBT<L::Hash, trie_db::DBValue> + hash_db::HashDBRef<L::Hash, trie_db::DBValue>,
+	I: IntoIterator<Item = &'a [u8]>,
+{
+	let mut nodes_iter = encoded.into_iter();
+	let (top_root, _nb_used) = trie_db::decode_compact_from_iter::<L, _, _>(db, &mut nodes_iter)?;
+
+	// Only check root if expected root is passed as argument.
+	if let Some(expected_root) = expected_root.filter(|expected| *expected != &top_root) {
+		return Err(Error::RootMismatch(top_root, *expected_root));
+	}
+
+	let mut child_tries = Vec::new();
+	{
+		// fetch child trie roots
+		let trie = crate::TrieDBBuilder::<L>::new(db, &top_root).build();
+
+		let mut iter = trie.iter()?;
+
+		let childtrie_roots = pezsp_core::storage::well_known_keys::DEFAULT_CHILD_STORAGE_KEY_PREFIX;
+		if iter.seek(childtrie_roots).is_ok() {
+			loop {
+				match iter.next() {
+					Some(Ok((key, value))) if key.starts_with(childtrie_roots) => {
+						// we expect all default child trie root to be correctly encoded.
+						// see other child trie functions.
+						let mut root = TrieHash::<L>::default();
+						// still in a proof so prevent panic
+						if root.as_mut().len() != value.as_slice().len() {
+							return Err(Error::InvalidChildRoot(key, value));
+						}
+						root.as_mut().copy_from_slice(value.as_ref());
+						child_tries.push(root);
+					},
+					// allow incomplete database error: we only
+					// require access to data in the proof.
+					Some(Err(error)) => match *error {
+						trie_db::TrieError::IncompleteDatabase(..) => (),
+						e => return Err(Box::new(e).into()),
+					},
+					_ => break,
+				}
+			}
+		}
+	}
+
+	if !HashDBT::<L::Hash, _>::contains(db, &top_root, EMPTY_PREFIX) {
+		return Err(Error::IncompleteProof);
+	}
+
+	let mut previous_extracted_child_trie = None;
+	let mut nodes_iter = nodes_iter.peekable();
+	for child_root in child_tries.into_iter() {
+		if previous_extracted_child_trie.is_none() && nodes_iter.peek().is_some() {
+			let (top_root, _) = trie_db::decode_compact_from_iter::<L, _, _>(db, &mut nodes_iter)?;
+			previous_extracted_child_trie = Some(top_root);
+		}
+
+		// we do not early exit on root mismatch but try the
+		// other read from proof (some child root may be
+		// in proof without actual child content).
+		if Some(child_root) == previous_extracted_child_trie {
+			previous_extracted_child_trie = None;
+		}
+	}
+
+	if let Some(child_root) = previous_extracted_child_trie {
+		// A child root was read from proof but is not present
+		// in top trie.
+		return Err(Error::ExtraneousChildProof(child_root));
+	}
+
+	if nodes_iter.next().is_some() {
+		return Err(Error::ExtraneousChildNode);
+	}
+
+	Ok(top_root)
+}
+
+/// Encode a compact proof.
+///
+/// Takes as input all full encoded node from the proof, and
+/// the root.
+/// Then parse all child trie root and compress main trie content first
+/// then all child trie contents.
+/// Child trie are ordered by the order of their roots in the top trie.
+pub fn encode_compact<L, DB>(
+	partial_db: &DB,
+	root: &TrieHash<L>,
+) -> Result<CompactProof, Error<TrieHash<L>, CError<L>>>
+where
+	L: TrieConfiguration,
+	DB: HashDBT<L::Hash, trie_db::DBValue> + hash_db::HashDBRef<L::Hash, trie_db::DBValue>,
+{
+	let mut child_tries = Vec::new();
+	let mut compact_proof = {
+		let trie = crate::TrieDBBuilder::<L>::new(partial_db, root).build();
+
+		let mut iter = trie.iter()?;
+
+		let childtrie_roots = pezsp_core::storage::well_known_keys::DEFAULT_CHILD_STORAGE_KEY_PREFIX;
+		if iter.seek(childtrie_roots).is_ok() {
+			loop {
+				match iter.next() {
+					Some(Ok((key, value))) if key.starts_with(childtrie_roots) => {
+						let mut root = TrieHash::<L>::default();
+						if root.as_mut().len() != value.as_slice().len() {
+							// some child trie root in top trie are not an encoded hash.
+							return Err(Error::InvalidChildRoot(key.to_vec(), value.to_vec()));
+						}
+						root.as_mut().copy_from_slice(value.as_ref());
+						child_tries.push(root);
+					},
+					// allow incomplete database error: we only
+					// require access to data in the proof.
+					Some(Err(error)) => match *error {
+						trie_db::TrieError::IncompleteDatabase(..) => (),
+						e => return Err(Box::new(e).into()),
+					},
+					_ => break,
+				}
+			}
+		}
+
+		trie_db::encode_compact::<L>(&trie)?
+	};
+
+	for child_root in child_tries {
+		if !HashDBT::<L::Hash, _>::contains(partial_db, &child_root, EMPTY_PREFIX) {
+			// child proof are allowed to be missing (unused root can be included
+			// due to trie structure modification).
+			continue;
+		}
+
+		let trie = crate::TrieDBBuilder::<L>::new(partial_db, &child_root).build();
+		let child_proof = trie_db::encode_compact::<L>(&trie)?;
+
+		compact_proof.extend(child_proof);
+	}
+
+	Ok(CompactProof { encoded_nodes: compact_proof })
+}
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+	use crate::{delta_trie_root, recorder::IgnoredNodes, HashDB, StorageProof};
+	use codec::Encode;
+	use hash_db::AsHashDB;
+	use pezsp_core::{Blake2Hasher, H256};
+	use std::collections::HashSet;
+	use trie_db::{DBValue, Trie, TrieDBBuilder, TrieDBMutBuilder, TrieHash, TrieMut};
+
+	type MemoryDB = crate::MemoryDB<pezsp_core::Blake2Hasher>;
+	type Layout = crate::LayoutV1<pezsp_core::Blake2Hasher>;
+	type Recorder = crate::recorder::Recorder<pezsp_core::Blake2Hasher>;
+
+	fn create_trie(num_keys: u32) -> (MemoryDB, TrieHash<Layout>) {
+		let mut db = MemoryDB::default();
+		let mut root = Default::default();
+
+		{
+			let mut trie = TrieDBMutBuilder::<Layout>::new(&mut db, &mut root).build();
+			for i in 0..num_keys {
+				trie.insert(
+					&i.encode(),
+					&vec![1u8; 64].into_iter().chain(i.encode()).collect::<Vec<_>>(),
+				)
+				.expect("Inserts data");
+			}
+		}
+
+		(db, root)
+	}
+
+	struct Overlay<'a> {
+		db: &'a MemoryDB,
+		write: MemoryDB,
+	}
+
+	impl hash_db::HashDB<pezsp_core::Blake2Hasher, DBValue> for Overlay<'_> {
+		fn get(
+			&self,
+			key: &<pezsp_core::Blake2Hasher as hash_db::Hasher>::Out,
+			prefix: hash_db::Prefix,
+		) -> Option<DBValue> {
+			HashDB::get(self.db, key, prefix)
+		}
+
+		fn contains(
+			&self,
+			key: &<pezsp_core::Blake2Hasher as hash_db::Hasher>::Out,
+			prefix: hash_db::Prefix,
+		) -> bool {
+			HashDB::contains(self.db, key, prefix)
+		}
+
+		fn insert(
+			&mut self,
+			prefix: hash_db::Prefix,
+			value: &[u8],
+		) -> <pezsp_core::Blake2Hasher as hash_db::Hasher>::Out {
+			self.write.insert(prefix, value)
+		}
+
+		fn emplace(
+			&mut self,
+			key: <pezsp_core::Blake2Hasher as hash_db::Hasher>::Out,
+			prefix: hash_db::Prefix,
+			value: DBValue,
+		) {
+			self.write.emplace(key, prefix, value);
+		}
+
+		fn remove(
+			&mut self,
+			key: &<pezsp_core::Blake2Hasher as hash_db::Hasher>::Out,
+			prefix: hash_db::Prefix,
+		) {
+			self.write.remove(key, prefix);
+		}
+	}
+
+	impl AsHashDB<Blake2Hasher, DBValue> for Overlay<'_> {
+		fn as_hash_db(&self) -> &dyn HashDBT<Blake2Hasher, DBValue> {
+			self
+		}
+
+		fn as_hash_db_mut<'a>(&'a mut self) -> &'a mut (dyn HashDBT<Blake2Hasher, DBValue> + 'a) {
+			self
+		}
+	}
+
+	fn emulate_block_building(
+		state: &MemoryDB,
+		root: H256,
+		read_keys: &[u32],
+		write_keys: &[u32],
+		nodes_to_ignore: IgnoredNodes<H256>,
+	) -> (Recorder, MemoryDB, H256) {
+		let recorder = Recorder::with_ignored_nodes(nodes_to_ignore);
+
+		{
+			let mut trie_recorder = recorder.as_trie_recorder(root);
+			let trie = TrieDBBuilder::<Layout>::new(state, &root)
+				.with_recorder(&mut trie_recorder)
+				.build();
+
+			for key in read_keys {
+				trie.get(&key.encode()).unwrap().unwrap();
+			}
+		}
+
+		let mut overlay = Overlay { db: state, write: Default::default() };
+
+		let new_root = {
+			let mut trie_recorder = recorder.as_trie_recorder(root);
+			delta_trie_root::<Layout, _, _, _, _, _>(
+				&mut overlay,
+				root,
+				write_keys.iter().map(|k| {
+					(
+						k.encode(),
+						Some(vec![2u8; 64].into_iter().chain(k.encode()).collect::<Vec<_>>()),
+					)
+				}),
+				Some(&mut trie_recorder),
+				None,
+			)
+			.unwrap()
+		};
+
+		(recorder, overlay.write, new_root)
+	}
+
+	fn build_known_nodes_list(recorder: &Recorder, transaction: &MemoryDB) -> IgnoredNodes<H256> {
+		let mut ignored_nodes =
+			IgnoredNodes::from_storage_proof::<Blake2Hasher>(&recorder.to_storage_proof());
+
+		ignored_nodes.extend(IgnoredNodes::from_memory_db::<Blake2Hasher, _>(transaction.clone()));
+
+		ignored_nodes
+	}
+
+	#[test]
+	fn ensure_multiple_tries_encode_compact_works() {
+		let (mut db, root) = create_trie(100);
+
+		let mut nodes_to_ignore = IgnoredNodes::default();
+		let (recorder, transaction, root1) = emulate_block_building(
+			&db,
+			root,
+			&[2, 4, 5, 6, 7, 8],
+			&[9, 10, 11, 12, 13, 14],
+			nodes_to_ignore.clone(),
+		);
+
+		db.consolidate(transaction.clone());
+		nodes_to_ignore.extend(build_known_nodes_list(&recorder, &transaction));
+
+		let (recorder2, transaction, root2) = emulate_block_building(
+			&db,
+			root1,
+			&[9, 10, 11, 12, 13, 14],
+			&[15, 16, 17, 18, 19, 20],
+			nodes_to_ignore.clone(),
+		);
+
+		db.consolidate(transaction.clone());
+		nodes_to_ignore.extend(build_known_nodes_list(&recorder2, &transaction));
+
+		let (recorder3, _, root3) = emulate_block_building(
+			&db,
+			root2,
+			&[20, 30, 40, 41, 42],
+			&[80, 90, 91, 92, 93],
+			nodes_to_ignore,
+		);
+
+		let proof = recorder.to_storage_proof();
+		let proof2 = recorder2.to_storage_proof();
+		let proof3 = recorder3.to_storage_proof();
+
+		let mut combined = HashSet::<Vec<u8>>::from_iter(proof.into_iter_nodes());
+		proof2.iter_nodes().for_each(|n| assert!(combined.insert(n.clone())));
+		proof3.iter_nodes().for_each(|n| assert!(combined.insert(n.clone())));
+
+		let proof = StorageProof::new(combined.into_iter());
+
+		let compact_proof = encode_compact::<Layout, _>(&proof.to_memory_db(), &root).unwrap();
+
+		assert!(proof.encoded_size() > compact_proof.encoded_size());
+
+		let mut res_db = crate::MemoryDB::<Blake2Hasher>::new(&[]);
+		decode_compact::<Layout, _, _>(
+			&mut res_db,
+			compact_proof.iter_compact_encoded_nodes(),
+			Some(&root),
+		)
+		.unwrap();
+
+		let (_, transaction, root1_proof) = emulate_block_building(
+			&res_db,
+			root,
+			&[2, 4, 5, 6, 7, 8],
+			&[9, 10, 11, 12, 13, 14],
+			Default::default(),
+		);
+
+		assert_eq!(root1, root1_proof);
+
+		res_db.consolidate(transaction);
+
+		let (_, transaction2, root2_proof) = emulate_block_building(
+			&res_db,
+			root1,
+			&[9, 10, 11, 12, 13, 14],
+			&[15, 16, 17, 18, 19, 20],
+			Default::default(),
+		);
+
+		assert_eq!(root2, root2_proof);
+
+		res_db.consolidate(transaction2);
+
+		let (_, _, root3_proof) = emulate_block_building(
+			&res_db,
+			root2,
+			&[20, 30, 40, 41, 42],
+			&[80, 90, 91, 92, 93],
+			Default::default(),
+		);
+
+		assert_eq!(root3, root3_proof);
+	}
+}

bizinikiwi/primitives/trie/src/trie_stream.rs

@@ -0,0 +1,147 @@
+// This file is part of Bizinikiwi.
+
+// Copyright (C) Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! `TrieStream` implementation for Bizinikiwi's trie format.
+
+use crate::{
+	node_header::{size_and_prefix_iterator, NodeKind},
+	trie_constants,
+};
+use alloc::vec::Vec;
+use codec::{Compact, Encode};
+use hash_db::Hasher;
+use trie_root;
+
+/// Codec-flavored TrieStream.
+#[derive(Default, Clone)]
+pub struct TrieStream {
+	/// Current node buffer.
+	buffer: Vec<u8>,
+}
+
+impl TrieStream {
+	// useful for debugging but not used otherwise
+	pub fn as_raw(&self) -> &[u8] {
+		&self.buffer
+	}
+}
+
+fn branch_node_bit_mask(has_children: impl Iterator<Item = bool>) -> (u8, u8) {
+	let mut bitmap: u16 = 0;
+	let mut cursor: u16 = 1;
+	for v in has_children {
+		if v {
+			bitmap |= cursor
+		}
+		cursor <<= 1;
+	}
+	((bitmap % 256) as u8, (bitmap / 256) as u8)
+}
+
+/// Create a leaf/branch node, encoding a number of nibbles.
+fn fuse_nibbles_node(nibbles: &[u8], kind: NodeKind) -> impl Iterator<Item = u8> + '_ {
+	let size = nibbles.len();
+	let iter_start = match kind {
+		NodeKind::Leaf => size_and_prefix_iterator(size, trie_constants::LEAF_PREFIX_MASK, 2),
+		NodeKind::BranchNoValue =>
+			size_and_prefix_iterator(size, trie_constants::BRANCH_WITHOUT_MASK, 2),
+		NodeKind::BranchWithValue =>
+			size_and_prefix_iterator(size, trie_constants::BRANCH_WITH_MASK, 2),
+		NodeKind::HashedValueLeaf =>
+			size_and_prefix_iterator(size, trie_constants::ALT_HASHING_LEAF_PREFIX_MASK, 3),
+		NodeKind::HashedValueBranch =>
+			size_and_prefix_iterator(size, trie_constants::ALT_HASHING_BRANCH_WITH_MASK, 4),
+	};
+	iter_start
+		.chain(if nibbles.len() % 2 == 1 { Some(nibbles[0]) } else { None })
+		.chain(nibbles[nibbles.len() % 2..].chunks(2).map(|ch| (ch[0] << 4) | ch[1]))
+}
+
+use trie_root::Value as TrieStreamValue;
+impl trie_root::TrieStream for TrieStream {
+	fn new() -> Self {
+		Self { buffer: Vec::new() }
+	}
+
+	fn append_empty_data(&mut self) {
+		self.buffer.push(trie_constants::EMPTY_TRIE);
+	}
+
+	fn append_leaf(&mut self, key: &[u8], value: TrieStreamValue) {
+		let kind = match &value {
+			TrieStreamValue::Inline(..) => NodeKind::Leaf,
+			TrieStreamValue::Node(..) => NodeKind::HashedValueLeaf,
+		};
+		self.buffer.extend(fuse_nibbles_node(key, kind));
+		match &value {
+			TrieStreamValue::Inline(value) => {
+				Compact(value.len() as u32).encode_to(&mut self.buffer);
+				self.buffer.extend_from_slice(value);
+			},
+			TrieStreamValue::Node(hash) => {
+				self.buffer.extend_from_slice(hash.as_slice());
+			},
+		};
+	}
+
+	fn begin_branch(
+		&mut self,
+		maybe_partial: Option<&[u8]>,
+		maybe_value: Option<TrieStreamValue>,
+		has_children: impl Iterator<Item = bool>,
+	) {
+		if let Some(partial) = maybe_partial {
+			let kind = match &maybe_value {
+				None => NodeKind::BranchNoValue,
+				Some(TrieStreamValue::Inline(..)) => NodeKind::BranchWithValue,
+				Some(TrieStreamValue::Node(..)) => NodeKind::HashedValueBranch,
+			};
+
+			self.buffer.extend(fuse_nibbles_node(partial, kind));
+			let bm = branch_node_bit_mask(has_children);
+			self.buffer.extend([bm.0, bm.1].iter());
+		} else {
+			unreachable!("trie stream codec only for no extension trie");
+		}
+		match maybe_value {
+			None => (),
+			Some(TrieStreamValue::Inline(value)) => {
+				Compact(value.len() as u32).encode_to(&mut self.buffer);
+				self.buffer.extend_from_slice(value);
+			},
+			Some(TrieStreamValue::Node(hash)) => {
+				self.buffer.extend_from_slice(hash.as_slice());
+			},
+		}
+	}
+
+	fn append_extension(&mut self, _key: &[u8]) {
+		debug_assert!(false, "trie stream codec only for no extension trie");
+	}
+
+	fn append_substream<H: Hasher>(&mut self, other: Self) {
+		let data = other.out();
+		match data.len() {
+			0..=31 => data.encode_to(&mut self.buffer),
+			_ => H::hash(&data).as_ref().encode_to(&mut self.buffer),
+		}
+	}
+
+	fn out(self) -> Vec<u8> {
+		self.buffer
+	}
+}

bizinikiwi/primitives/trie/test-res/storage_root

@@ -0,0 +1 @@
+‡=Τ[Α42%αJPΆΑh¦Kwι)R� 0¤Τ­uΏΓ