Add ChildTriePrefixIterator and methods (#8478)

* Make use of PrefixIterator underneath Storage[Key]Iterator * Add ChildTriePrefixIterator and methods * Add documentation on ChilTriePrefixIterator fields * Deprecate Storage[Key]Iterator API instead of removing them * Allow fetching for the prefix as an option for ChildTriePrefixIterator * Rename prefix_fetch to fetch_previous_key * fix implementation + test * make gitdiff better * Add test for storage_iter and storage_key_iter Co-authored-by: thiolliere <gui.thiolliere@gmail.com>
2026-06-24 09:01:08 +00:00 · 2021-04-01 07:20:43 -07:00
parent 643d2b669f
commit 1fdc8fa1c1
3 changed files with 307 additions and 5 deletions
@@ -17,9 +17,10 @@

 //! Stuff to do with the runtime's storage.

+use sp_core::storage::ChildInfo;
 use sp_std::prelude::*;
 use codec::{FullCodec, FullEncode, Encode, EncodeLike, Decode};
-use crate::hash::{Twox128, StorageHasher};
+use crate::hash::{Twox128, StorageHasher, ReversibleStorageHasher};
 use sp_runtime::generic::{Digest, DigestItem};
 pub use sp_runtime::TransactionOutcome;

@@ -519,6 +520,14 @@ pub struct PrefixIterator<T> {
 	closure: fn(&[u8], &[u8]) -> Result<T, codec::Error>,
 }

+impl<T> PrefixIterator<T> {
+	/// Mutate this iterator into a draining iterator; items iterated are removed from storage.
+	pub fn drain(mut self) -> Self {
+		self.drain = true;
+		self
+	}
+}
+
 impl<T> Iterator for PrefixIterator<T> {
 	type Item = T;

@@ -563,6 +572,133 @@ impl<T> Iterator for PrefixIterator<T> {
 	}
 }

+/// Iterate over a prefix of a child trie and decode raw_key and raw_value into `T`.
+///
+/// If any decoding fails it skips the key and continues to the next one.
+pub struct ChildTriePrefixIterator<T> {
+	/// The prefix iterated on
+	prefix: Vec<u8>,
+	/// child info for child trie
+	child_info: ChildInfo,
+	/// The last key iterated on
+	previous_key: Vec<u8>,
+	/// If true then values are removed while iterating
+	drain: bool,
+	/// Whether or not we should fetch the previous key
+	fetch_previous_key: bool,
+	/// Function that takes `(raw_key_without_prefix, raw_value)` and decode `T`.
+	/// `raw_key_without_prefix` is the raw storage key without the prefix iterated on.
+	closure: fn(&[u8], &[u8]) -> Result<T, codec::Error>,
+}
+
+impl<T> ChildTriePrefixIterator<T> {
+	/// Mutate this iterator into a draining iterator; items iterated are removed from storage.
+	pub fn drain(mut self) -> Self {
+		self.drain = true;
+		self
+	}
+}
+
+impl<T: Decode + Sized> ChildTriePrefixIterator<(Vec<u8>, T)> {
+	/// Construct iterator to iterate over child trie items in `child_info` with the prefix `prefix`.
+	///
+	/// NOTE: Iterator with [`Self::drain`] will remove any value who failed to decode
+	pub fn with_prefix(child_info: &ChildInfo, prefix: &[u8]) -> Self {
+		let prefix = prefix.to_vec();
+		let previous_key = prefix.clone();
+		let closure = |raw_key_without_prefix: &[u8], raw_value: &[u8]| {
+			let value = T::decode(&mut &raw_value[..])?;
+			Ok((raw_key_without_prefix.to_vec(), value))
+		};
+
+		Self {
+			prefix,
+			child_info: child_info.clone(),
+			previous_key,
+			drain: false,
+			fetch_previous_key: true,
+			closure,
+		}
+	}
+}
+
+impl<K: Decode + Sized, T: Decode + Sized> ChildTriePrefixIterator<(K, T)> {
+	/// Construct iterator to iterate over child trie items in `child_info` with the prefix `prefix`.
+	///
+	/// NOTE: Iterator with [`Self::drain`] will remove any key or value who failed to decode
+	pub fn with_prefix_over_key<H: ReversibleStorageHasher>(child_info: &ChildInfo, prefix: &[u8]) -> Self {
+		let prefix = prefix.to_vec();
+		let previous_key = prefix.clone();
+		let closure = |raw_key_without_prefix: &[u8], raw_value: &[u8]| {
+			let mut key_material = H::reverse(raw_key_without_prefix);
+			let key = K::decode(&mut key_material)?;
+			let value = T::decode(&mut &raw_value[..])?;
+			Ok((key, value))
+		};
+
+		Self {
+			prefix,
+			child_info: child_info.clone(),
+			previous_key,
+			drain: false,
+			fetch_previous_key: true,
+			closure,
+		 }
+	}
+}
+
+impl<T> Iterator for ChildTriePrefixIterator<T> {
+	type Item = T;
+
+	fn next(&mut self) -> Option<Self::Item> {
+		loop {
+			let maybe_next = if self.fetch_previous_key {
+				self.fetch_previous_key = false;
+				Some(self.previous_key.clone())
+			} else {
+				sp_io::default_child_storage::next_key(
+					&self.child_info.storage_key(),
+					&self.previous_key,
+				)
+					.filter(|n| n.starts_with(&self.prefix))
+			};
+			break match maybe_next {
+				Some(next) => {
+					self.previous_key = next;
+					let raw_value = match child::get_raw(&self.child_info, &self.previous_key) {
+						Some(raw_value) => raw_value,
+						None => {
+							log::error!(
+								"next_key returned a key with no value at {:?}",
+								self.previous_key,
+							);
+							continue
+						}
+					};
+					if self.drain {
+						child::kill(&self.child_info, &self.previous_key)
+					}
+					let raw_key_without_prefix = &self.previous_key[self.prefix.len()..];
+					let item = match (self.closure)(raw_key_without_prefix, &raw_value[..]) {
+						Ok(item) => item,
+						Err(e) => {
+							log::error!(
+								"(key, value) failed to decode at {:?}: {:?}",
+								self.previous_key,
+								e,
+							);
+							continue
+						}
+					};
+
+					Some(item)
+				}
+				None => None,
+			}
+		}
+	}
+}
+
 /// Trait for maps that store all its value after a unique prefix.
 ///
 /// By default the final prefix is:
@@ -689,6 +825,7 @@ impl<Hash: Encode> StorageAppend<DigestItem<Hash>> for Digest<Hash> {}
 mod test {
 	use super::*;
 	use sp_core::hashing::twox_128;
+	use crate::hash::Identity;
 	use sp_io::TestExternalities;
 	use generator::StorageValue as _;

@@ -825,4 +962,78 @@ mod test {
 			});
 		});
 	}
+
+	#[test]
+	fn child_trie_prefixed_map_works() {
+		TestExternalities::default().execute_with(|| {
+			let child_info_a = child::ChildInfo::new_default(b"a");
+			child::put(&child_info_a, &[1, 2, 3], &8u16);
+			child::put(&child_info_a, &[2], &8u16);
+			child::put(&child_info_a, &[2, 1, 3], &8u8);
+			child::put(&child_info_a, &[2, 2, 3], &8u16);
+			child::put(&child_info_a, &[3], &8u16);
+
+			assert_eq!(
+				ChildTriePrefixIterator::with_prefix(&child_info_a, &[2])
+					.collect::<Vec<(Vec<u8>, u16)>>(),
+				vec![
+					(vec![], 8),
+					(vec![2, 3], 8),
+				],
+			);
+
+			assert_eq!(
+				ChildTriePrefixIterator::with_prefix(&child_info_a, &[2])
+					.drain()
+					.collect::<Vec<(Vec<u8>, u16)>>(),
+				vec![
+					(vec![], 8),
+					(vec![2, 3], 8),
+				],
+			);
+
+			// The only remaining is the ones outside prefix
+			assert_eq!(
+				ChildTriePrefixIterator::with_prefix(&child_info_a, &[])
+					.collect::<Vec<(Vec<u8>, u8)>>(),
+				vec![
+					(vec![1, 2, 3], 8),
+					(vec![3], 8),
+				],
+			);
+
+			child::put(&child_info_a, &[1, 2, 3], &8u16);
+			child::put(&child_info_a, &[2], &8u16);
+			child::put(&child_info_a, &[2, 1, 3], &8u8);
+			child::put(&child_info_a, &[2, 2, 3], &8u16);
+			child::put(&child_info_a, &[3], &8u16);
+
+			assert_eq!(
+				ChildTriePrefixIterator::with_prefix_over_key::<Identity>(&child_info_a, &[2])
+					.collect::<Vec<(u16, u16)>>(),
+				vec![
+					(u16::decode(&mut &[2, 3][..]).unwrap(), 8),
+				],
+			);
+
+			assert_eq!(
+				ChildTriePrefixIterator::with_prefix_over_key::<Identity>(&child_info_a, &[2])
+					.drain()
+					.collect::<Vec<(u16, u16)>>(),
+				vec![
+					(u16::decode(&mut &[2, 3][..]).unwrap(), 8),
+				],
+			);
+
+			// The only remaining is the ones outside prefix
+			assert_eq!(
+				ChildTriePrefixIterator::with_prefix(&child_info_a, &[])
+					.collect::<Vec<(Vec<u8>, u8)>>(),
+				vec![
+					(vec![1, 2, 3], 8),
+					(vec![3], 8),
+				],
+			);
+		});
+	}
 }