pallet_stateful_storage/
stateful_child_tree.rs

1use core::marker::PhantomData;
2
3use common_primitives::msa::MessageSourceId;
4use core::fmt::Debug;
5use frame_support::{
6	storage::{child, child::ChildInfo, ChildTriePrefixIterator},
7	Blake2_128, Blake2_256, Hashable, StorageHasher, Twox128, Twox256,
8};
9use parity_scale_codec::{Decode, DecodeAll, Encode};
10use sp_core::{ConstU8, Get};
11use sp_io::hashing::twox_64;
12extern crate alloc;
13use alloc::{boxed::Box, vec::Vec};
14
15/// Hasher to use to hash keys to insert to storage.
16pub trait MultipartKeyStorageHasher: StorageHasher {
17	type HashSize: Get<u8>;
18	/// Split the hash part out of the input.
19	///
20	/// I.e. for input `&[hash ++ hash ++ key]` returns `&[key]`
21	fn reverse(x: &[u8], num_key_parts: u8) -> &[u8] {
22		let hash_len: usize = (Self::HashSize::get() * num_key_parts).into();
23		if x.len() < hash_len {
24			log::error!("Invalid reverse: hash length too short");
25			return &[]
26		}
27		&x[hash_len..]
28	}
29}
30/// 128-bit Blake2 hash.
31impl MultipartKeyStorageHasher for Blake2_128 {
32	type HashSize = ConstU8<16>;
33}
34/// 256-bit Blake2 hash.
35impl MultipartKeyStorageHasher for Blake2_256 {
36	type HashSize = ConstU8<32>;
37}
38/// 128-bit XX hash.
39impl MultipartKeyStorageHasher for Twox128 {
40	type HashSize = ConstU8<16>;
41}
42/// 256-bit XX hash.
43impl MultipartKeyStorageHasher for Twox256 {
44	type HashSize = ConstU8<32>;
45}
46
47pub trait MultipartStorageKeyPart:
48	Clone + Debug + Default + Decode + Encode + Eq + PartialEq + Hashable
49{
50}
51impl<T> MultipartStorageKeyPart for T where
52	T: Clone + Debug + Default + Decode + Encode + Eq + PartialEq + Hashable
53{
54}
55
56/// A trait that defines operations for a multi-part key in child storage that allows to retrieve
57/// all different prefix combinations
58pub trait MultipartKey<H: MultipartKeyStorageHasher>: MultipartStorageKeyPart {
59	type Arity: Get<u8>;
60
61	fn hash(&self) -> Vec<u8>;
62	fn hash_prefix_only(&self) -> Vec<u8>;
63
64	#[allow(dead_code)]
65	fn decode(hash: &[u8]) -> Result<Self, parity_scale_codec::Error> {
66		let mut key_material = H::reverse(hash, Self::Arity::get());
67		<Self as Decode>::decode(&mut key_material)
68	}
69
70	fn decode_without_prefix(
71		hash: &[u8],
72		prefix_len: u8,
73	) -> Result<Self, parity_scale_codec::Error> {
74		if prefix_len > Self::Arity::get() {
75			return Err("Prefix longer than total key length".into())
76		}
77
78		let mut key_material = H::reverse(hash, Self::Arity::get() - prefix_len);
79		<Self as Decode>::decode(&mut key_material)
80	}
81}
82
83impl<H: MultipartKeyStorageHasher> MultipartKey<H> for () {
84	type Arity = ConstU8<0>;
85
86	fn hash(&self) -> Vec<u8> {
87		Vec::new()
88	}
89
90	fn hash_prefix_only(&self) -> Vec<u8> {
91		Vec::new()
92	}
93
94	fn decode(_hash: &[u8]) -> Result<Self, parity_scale_codec::Error> {
95		Ok(())
96	}
97}
98
99impl<H: MultipartKeyStorageHasher, T1: MultipartStorageKeyPart> MultipartKey<H> for (T1,) {
100	type Arity = ConstU8<1>;
101
102	fn hash(&self) -> Vec<u8> {
103		let encoded_1 = self.0.encode();
104		<H as StorageHasher>::hash(&encoded_1)
105			.as_ref()
106			.iter()
107			.chain(self.encode().iter())
108			.cloned()
109			.collect::<Vec<_>>()
110	}
111
112	fn hash_prefix_only(&self) -> Vec<u8> {
113		let encoded_1 = self.0.encode();
114		<H as StorageHasher>::hash(&encoded_1).as_ref().to_vec()
115	}
116}
117
118impl<H: MultipartKeyStorageHasher, T1: MultipartStorageKeyPart, T2: MultipartStorageKeyPart>
119	MultipartKey<H> for (T1, T2)
120{
121	type Arity = ConstU8<2>;
122
123	fn hash(&self) -> Vec<u8> {
124		let encoded_1 = self.0.encode();
125		let encoded_2 = self.1.encode();
126		H::hash(&encoded_1)
127			.as_ref()
128			.iter()
129			.chain(H::hash(&encoded_2).as_ref().iter())
130			.chain(self.encode().iter())
131			.cloned()
132			.collect::<Vec<_>>()
133	}
134
135	fn hash_prefix_only(&self) -> Vec<u8> {
136		let encoded_1 = self.0.encode();
137		let encoded_2 = self.1.encode();
138		<H as StorageHasher>::hash(&encoded_1)
139			.as_ref()
140			.iter()
141			.chain(H::hash(&encoded_2).as_ref().iter())
142			.cloned()
143			.collect::<Vec<_>>()
144	}
145}
146
147impl<
148		H: MultipartKeyStorageHasher,
149		T1: MultipartStorageKeyPart,
150		T2: MultipartStorageKeyPart,
151		T3: MultipartStorageKeyPart,
152	> MultipartKey<H> for (T1, T2, T3)
153{
154	type Arity = ConstU8<3>;
155
156	fn hash(&self) -> Vec<u8> {
157		let encoded_1 = self.0.encode();
158		let encoded_2 = self.1.encode();
159		let encoded_3 = self.2.encode();
160		H::hash(&encoded_1)
161			.as_ref()
162			.iter()
163			.chain(H::hash(&encoded_2).as_ref().iter())
164			.chain(H::hash(&encoded_3).as_ref().iter())
165			.chain(self.encode().iter())
166			.cloned()
167			.collect::<Vec<_>>()
168	}
169
170	fn hash_prefix_only(&self) -> Vec<u8> {
171		let encoded_1 = self.0.encode();
172		let encoded_2 = self.1.encode();
173		let encoded_3 = self.2.encode();
174		<H as StorageHasher>::hash(&encoded_1)
175			.as_ref()
176			.iter()
177			.chain(H::hash(&encoded_2).as_ref().iter())
178			.chain(H::hash(&encoded_3).as_ref().iter())
179			.cloned()
180			.collect::<Vec<_>>()
181	}
182}
183
184impl<
185		H: MultipartKeyStorageHasher,
186		T1: MultipartStorageKeyPart,
187		T2: MultipartStorageKeyPart,
188		T3: MultipartStorageKeyPart,
189		T4: MultipartStorageKeyPart,
190	> MultipartKey<H> for (T1, T2, T3, T4)
191{
192	type Arity = ConstU8<4>;
193
194	fn hash(&self) -> Vec<u8> {
195		let encoded_1 = self.0.encode();
196		let encoded_2 = self.1.encode();
197		let encoded_3 = self.2.encode();
198		let encoded_4 = self.3.encode();
199		H::hash(&encoded_1)
200			.as_ref()
201			.iter()
202			.chain(H::hash(&encoded_2).as_ref().iter())
203			.chain(H::hash(&encoded_3).as_ref().iter())
204			.chain(H::hash(&encoded_4).as_ref().iter())
205			.chain(self.encode().iter())
206			.cloned()
207			.collect::<Vec<_>>()
208	}
209
210	fn hash_prefix_only(&self) -> Vec<u8> {
211		let encoded_1 = self.0.encode();
212		let encoded_2 = self.1.encode();
213		let encoded_3 = self.2.encode();
214		let encoded_4 = self.3.encode();
215		<H as StorageHasher>::hash(&encoded_1)
216			.as_ref()
217			.iter()
218			.chain(H::hash(&encoded_2).as_ref().iter())
219			.chain(H::hash(&encoded_3).as_ref().iter())
220			.chain(H::hash(&encoded_4).as_ref().iter())
221			.cloned()
222			.collect::<Vec<_>>()
223	}
224}
225
226/// A trait that allows to iterate on the prefix of a multi-part key
227pub trait IsTuplePrefix<H: MultipartKeyStorageHasher, T: MultipartStorageKeyPart>:
228	MultipartKey<H>
229{
230}
231impl<H: MultipartKeyStorageHasher, T1: MultipartStorageKeyPart> IsTuplePrefix<H, (T1,)> for () {}
232impl<H: MultipartKeyStorageHasher, T1: MultipartStorageKeyPart, T2: MultipartStorageKeyPart>
233	IsTuplePrefix<H, (T1, T2)> for ()
234{
235}
236impl<H: MultipartKeyStorageHasher, T1: MultipartStorageKeyPart, T2: MultipartStorageKeyPart>
237	IsTuplePrefix<H, (T1, T2)> for (T1,)
238{
239}
240impl<
241		H: MultipartKeyStorageHasher,
242		T1: MultipartStorageKeyPart,
243		T2: MultipartStorageKeyPart,
244		T3: MultipartStorageKeyPart,
245	> IsTuplePrefix<H, (T1, T2, T3)> for ()
246{
247}
248impl<
249		H: MultipartKeyStorageHasher,
250		T1: MultipartStorageKeyPart,
251		T2: MultipartStorageKeyPart,
252		T3: MultipartStorageKeyPart,
253	> IsTuplePrefix<H, (T1, T2, T3)> for (T1,)
254{
255}
256impl<
257		H: MultipartKeyStorageHasher,
258		T1: MultipartStorageKeyPart,
259		T2: MultipartStorageKeyPart,
260		T3: MultipartStorageKeyPart,
261	> IsTuplePrefix<H, (T1, T2, T3)> for (T1, T2)
262{
263}
264impl<
265		H: MultipartKeyStorageHasher,
266		T1: MultipartStorageKeyPart,
267		T2: MultipartStorageKeyPart,
268		T3: MultipartStorageKeyPart,
269		T4: MultipartStorageKeyPart,
270	> IsTuplePrefix<H, (T1, T2, T3, T4)> for ()
271{
272}
273impl<
274		H: MultipartKeyStorageHasher,
275		T1: MultipartStorageKeyPart,
276		T2: MultipartStorageKeyPart,
277		T3: MultipartStorageKeyPart,
278		T4: MultipartStorageKeyPart,
279	> IsTuplePrefix<H, (T1, T2, T3, T4)> for (T1,)
280{
281}
282impl<
283		H: MultipartKeyStorageHasher,
284		T1: MultipartStorageKeyPart,
285		T2: MultipartStorageKeyPart,
286		T3: MultipartStorageKeyPart,
287		T4: MultipartStorageKeyPart,
288	> IsTuplePrefix<H, (T1, T2, T3, T4)> for (T1, T2)
289{
290}
291impl<
292		H: MultipartKeyStorageHasher,
293		T1: MultipartStorageKeyPart,
294		T2: MultipartStorageKeyPart,
295		T3: MultipartStorageKeyPart,
296		T4: MultipartStorageKeyPart,
297	> IsTuplePrefix<H, (T1, T2, T3, T4)> for (T1, T2, T3)
298{
299}
300
301/// Paginated Stateful data access utility
302pub struct StatefulChildTree<H: MultipartKeyStorageHasher = Twox128> {
303	hasher: PhantomData<H>,
304}
305impl<H: MultipartKeyStorageHasher> StatefulChildTree<H> {
306	/// Reads a child tree node and tries to decode it
307	///
308	/// The read is performed from the `msa_id` only. The `key` is not required. If the
309	/// data doesn't store under the given `key` `None` is returned.
310	pub fn try_read<K: MultipartKey<H>, V: DecodeAll + Sized>(
311		msa_id: &MessageSourceId,
312		pallet_name: &[u8],
313		storage_name: &[u8],
314		keys: &K,
315	) -> Result<Option<V>, ()> {
316		let child = Self::get_child_tree_for_storage(*msa_id, pallet_name, storage_name);
317		let value = child::get_raw(&child, &keys.hash());
318		match value {
319			None => Ok(None),
320			Some(v) => Ok(DecodeAll::decode_all(&mut &v[..]).map(Some).map_err(|_| ())?),
321		}
322	}
323
324	/// Read a child trie using a raw key
325	pub fn try_read_raw<V: Decode + Sized>(
326		child_info: &ChildInfo,
327		key: &[u8],
328	) -> Result<Option<V>, ()> {
329		let value = child::get_raw(child_info, key);
330		match value {
331			None => Ok(None),
332			Some(v) => Ok(Decode::decode(&mut &v[..]).map(Some).map_err(|_| ())?),
333		}
334	}
335
336	/// Prefix Iterator for a child tree
337	///
338	/// Allows getting all the keys having the same prefix
339	/// Warning: This should not be used from any on-chain transaction!
340	pub fn prefix_iterator<
341		V: Decode + Sized,
342		K: MultipartKey<H> + Sized,
343		PrefixKey: IsTuplePrefix<H, K>,
344	>(
345		msa_id: &MessageSourceId,
346		pallet_name: &[u8],
347		storage_name: &[u8],
348		prefix_keys: &PrefixKey,
349	) -> Box<impl Iterator<Item = (K, V)>> {
350		let child = Self::get_child_tree_for_storage(*msa_id, pallet_name, storage_name);
351		let result = ChildTriePrefixIterator::<(Vec<u8>, V)>::with_prefix(
352			&child,
353			&prefix_keys.hash_prefix_only(),
354		)
355		.filter_map(|(k, v)| {
356			if let Ok(key) =
357				<K as MultipartKey<H>>::decode_without_prefix(&k, PrefixKey::Arity::get())
358			{
359				Some((key, v))
360			} else {
361				None
362			}
363		});
364		Box::new(result)
365	}
366
367	/// Writes directly into child tree node
368	pub fn write<K: MultipartKey<H>, V: Encode + Sized>(
369		msa_id: &MessageSourceId,
370		pallet_name: &[u8],
371		storage_name: &[u8],
372		keys: &K,
373		new_value: V,
374	) {
375		let child_trie_info = &Self::get_child_tree_for_storage(*msa_id, pallet_name, storage_name);
376		child::put_raw(child_trie_info, &keys.hash(), new_value.encode().as_ref());
377	}
378
379	/// Writes directly into child tree node using pre-hashed keys
380	pub fn write_raw<V: Encode + Sized>(child_trie_info: &ChildInfo, key: &[u8], new_value: V) {
381		child::put_raw(child_trie_info, key, new_value.encode().as_ref());
382	}
383
384	/// Kills a child tree node
385	pub fn kill<K: MultipartKey<H>>(
386		msa_id: &MessageSourceId,
387		pallet_name: &[u8],
388		storage_name: &[u8],
389		keys: &K,
390	) {
391		let child_trie_info = &Self::get_child_tree_for_storage(*msa_id, pallet_name, storage_name);
392		child::kill(child_trie_info, &keys.hash());
393	}
394
395	/// Get the child trie at the given root
396	// These hashes should be consistent across the chain so we are hardcoding them
397	pub fn get_child_tree_for_storage(
398		msa_id: MessageSourceId,
399		pallet_name: &[u8],
400		storage_name: &[u8],
401	) -> ChildInfo {
402		let trie_root = Self::get_tree_prefix(msa_id);
403		// child tree root should be hashed by Blake128 to avoid probability of conflict
404		let hashed_keys: Vec<u8> = [
405			Blake2_128::hash(&trie_root[..]).as_ref(),
406			twox_64(pallet_name).as_ref(),
407			twox_64(storage_name).as_ref(),
408		]
409		.concat();
410		child::ChildInfo::new_default(&hashed_keys)
411	}
412
413	/// Storage Prefix for a given MSA Id
414	fn get_tree_prefix(msa_id: MessageSourceId) -> Vec<u8> {
415		let arr = [&msa_id.encode()[..], b"::"].concat();
416		arr.to_vec()
417	}
418}
pallet_stateful_storage/stateful_child_tree.rs

pallet_stateful_storage/
stateful_child_tree.rs
