avail_rust_core/
decoded_storage.rs

1use crate::{Error, rpc};
2use codec::{Decode, Encode};
3use primitive_types::H256;
4use std::marker::PhantomData;
5use subxt_rpcs::RpcClient;
6
7#[derive(Debug, Clone, Copy)]
8pub enum StorageHasher {
9	/// 128-bit Blake2 hash.
10	Blake2_128,
11	/// 256-bit Blake2 hash.
12	Blake2_256,
13	/// Multiple 128-bit Blake2 hashes concatenated.
14	Blake2_128Concat,
15	/// 128-bit XX hash.
16	Twox128,
17	/// 256-bit XX hash.
18	Twox256,
19	/// Multiple 64-bit XX hashes concatenated.
20	Twox64Concat,
21	/// Identity hashing (no hashing).
22	Identity,
23}
24
25impl StorageHasher {
26	pub fn hash(&self, data: &[u8]) -> Vec<u8> {
27		match self {
28			StorageHasher::Blake2_128 => sp_crypto_hashing::blake2_128(data).into(),
29			StorageHasher::Blake2_256 => sp_crypto_hashing::blake2_256(data).into(),
30			StorageHasher::Blake2_128Concat => {
31				let mut hash = sp_crypto_hashing::blake2_128(data).to_vec();
32				hash.extend_from_slice(data);
33				hash
34			},
35			StorageHasher::Twox128 => sp_crypto_hashing::twox_128(data).into(),
36			StorageHasher::Twox256 => sp_crypto_hashing::twox_256(data).into(),
37			StorageHasher::Twox64Concat => {
38				let mut hash = sp_crypto_hashing::twox_64(data).to_vec();
39				hash.extend_from_slice(data);
40				hash
41			},
42			StorageHasher::Identity => data.to_vec(),
43		}
44	}
45
46	pub fn from_hash<Key: codec::Decode>(&self, data: &mut &[u8]) -> Result<Key, codec::Error> {
47		match self {
48			StorageHasher::Blake2_128Concat => {
49				if data.len() < 17 {
50					return Err(codec::Error::from("Not enough data to compute Blake2_128Concat"));
51				}
52				Key::decode(&mut &data[16..])
53			},
54			StorageHasher::Twox64Concat => {
55				if data.len() < 9 {
56					return Err(codec::Error::from("Not enough data to compute Twox64Concat"));
57				}
58				Key::decode(&mut &data[8..])
59			},
60			StorageHasher::Identity => Key::decode(data),
61			_ => unimplemented!(),
62		}
63	}
64}
65
66pub trait StorageValue {
67	const PALLET_NAME: &str;
68	const STORAGE_NAME: &str;
69	type VALUE: codec::Decode;
70
71	fn encode_storage_key() -> [u8; 32] {
72		use sp_crypto_hashing::twox_128;
73
74		let mut encoded_storage_key = [0u8; 32];
75		encoded_storage_key[0..16].copy_from_slice(&twox_128(Self::PALLET_NAME.as_bytes()));
76		encoded_storage_key[16..].copy_from_slice(&twox_128(Self::STORAGE_NAME.as_bytes()));
77
78		encoded_storage_key
79	}
80
81	fn hex_encode_storage_key() -> String {
82		std::format!("0x{}", hex::encode(&Self::encode_storage_key()))
83	}
84
85	fn decode(value: &mut &[u8]) -> Result<Self::VALUE, codec::Error> {
86		Self::VALUE::decode(value)
87	}
88
89	fn fetch(
90		client: &RpcClient,
91		at: Option<H256>,
92	) -> impl std::future::Future<Output = Result<Option<Self::VALUE>, Error>> {
93		async move {
94			let storage_key = hex::encode(Self::encode_storage_key());
95
96			let storage_value = rpc::state::get_storage(client, &storage_key, at).await?;
97			let Some(storage_value) = storage_value else {
98				return Ok(None);
99			};
100
101			let storage_value = Self::decode(&mut storage_value.as_slice())?;
102			return Ok(Some(storage_value));
103		}
104	}
105}
106
107pub trait StorageMap {
108	const PALLET_NAME: &str;
109	const STORAGE_NAME: &str;
110	const KEY_HASHER: StorageHasher;
111	type KEY: codec::Decode + codec::Encode;
112	type VALUE: codec::Decode;
113
114	fn encode_partial_key() -> [u8; 32] {
115		use sp_crypto_hashing::twox_128;
116
117		let mut encoded_storage_key = [0u8; 32];
118		encoded_storage_key[0..16].copy_from_slice(&twox_128(Self::PALLET_NAME.as_bytes()));
119		encoded_storage_key[16..].copy_from_slice(&twox_128(Self::STORAGE_NAME.as_bytes()));
120
121		encoded_storage_key
122	}
123
124	fn hex_encode_partial_key() -> String {
125		std::format!("0x{}", hex::encode(&Self::encode_partial_key()))
126	}
127
128	fn encode_storage_key(key: &Self::KEY) -> Vec<u8> {
129		let mut storage_key: Vec<u8> = Vec::new();
130		storage_key.extend_from_slice(&Self::encode_partial_key());
131
132		let encoded_key = key.encode();
133		storage_key.extend_from_slice(&Self::KEY_HASHER.hash(&encoded_key));
134
135		storage_key
136	}
137
138	fn hex_encode_storage_key(key: &Self::KEY) -> String {
139		std::format!("0x{}", hex::encode(&Self::encode_storage_key(key)))
140	}
141
142	fn decode_storage_key(value: &mut &[u8]) -> Result<Self::KEY, codec::Error> {
143		// Skip pallet/variant
144		*value = &value[32..];
145
146		Self::KEY_HASHER.from_hash::<Self::KEY>(value)
147	}
148
149	fn decode_storage_value(value: &mut &[u8]) -> Result<Self::VALUE, codec::Error> {
150		Self::VALUE::decode(value)
151	}
152
153	fn fetch(
154		client: &RpcClient,
155		key: &Self::KEY,
156		at: Option<H256>,
157	) -> impl std::future::Future<Output = Result<Option<Self::VALUE>, Error>> {
158		async move {
159			let storage_key = hex::encode(Self::encode_storage_key(key));
160			let storage_value = rpc::state::get_storage(client, &storage_key, at).await?;
161			let Some(storage_value) = storage_value else {
162				return Ok(None);
163			};
164
165			let storage_value = Self::decode_storage_value(&mut storage_value.as_slice())?;
166			return Ok(Some(storage_value));
167		}
168	}
169
170	fn iter(client: RpcClient, block_hash: H256) -> StorageMapIterator<Self>
171	where
172		Self: Sized,
173	{
174		StorageMapIterator::new(client, block_hash)
175	}
176}
177
178pub trait StorageDoubleMap {
179	const PALLET_NAME: &str;
180	const STORAGE_NAME: &str;
181	const KEY1_HASHER: StorageHasher;
182	const KEY2_HASHER: StorageHasher;
183	type KEY1: codec::Decode + codec::Encode;
184	type KEY2: codec::Decode + codec::Encode;
185	type VALUE: codec::Decode;
186
187	fn encode_partial_key(key1: &Self::KEY1) -> Vec<u8> {
188		use sp_crypto_hashing::twox_128;
189
190		let mut encoded_storage_key = Vec::new();
191		encoded_storage_key.extend_from_slice(&twox_128(Self::PALLET_NAME.as_bytes()));
192		encoded_storage_key.extend_from_slice(&twox_128(Self::STORAGE_NAME.as_bytes()));
193		encoded_storage_key.extend_from_slice(&Self::KEY1_HASHER.hash(&key1.encode()));
194
195		encoded_storage_key
196	}
197
198	fn hex_encode_partial_key(key1: &Self::KEY1) -> String {
199		std::format!("0x{}", hex::encode(&Self::encode_partial_key(key1)))
200	}
201
202	fn encode_storage_key(key1: &Self::KEY1, key2: &Self::KEY2) -> Vec<u8> {
203		let mut storage_key: Vec<u8> = Vec::new();
204		storage_key.extend_from_slice(&Self::encode_partial_key(key1));
205		storage_key.extend_from_slice(&Self::KEY2_HASHER.hash(&key2.encode()));
206
207		storage_key
208	}
209
210	fn hex_encode_storage_key(key1: &Self::KEY1, key2: &Self::KEY2) -> String {
211		std::format!("0x{}", hex::encode(&Self::encode_storage_key(key1, key2)))
212	}
213
214	fn decode_partial_key(value: &mut &[u8]) -> Result<Self::KEY1, codec::Error> {
215		// Skip pallet/variant
216		*value = &value[32..];
217
218		Self::KEY1_HASHER.from_hash::<Self::KEY1>(value)
219	}
220
221	fn decode_storage_key(value: &mut &[u8]) -> Result<(Self::KEY1, Self::KEY2), codec::Error> {
222		// Skip pallet/variant
223		*value = &value[32..];
224
225		let key1 = Self::KEY1_HASHER.from_hash::<Self::KEY1>(value)?;
226		let key2 = Self::KEY2_HASHER.from_hash::<Self::KEY2>(value)?;
227		Ok((key1, key2))
228	}
229
230	fn decode_storage_value(value: &mut &[u8]) -> Result<Self::VALUE, codec::Error> {
231		Self::VALUE::decode(value)
232	}
233
234	fn fetch(
235		client: &RpcClient,
236		key_1: &Self::KEY1,
237		key_2: &Self::KEY2,
238		at: Option<H256>,
239	) -> impl std::future::Future<Output = Result<Option<Self::VALUE>, Error>> {
240		async move {
241			let storage_key = hex::encode(Self::encode_storage_key(key_1, key_2));
242			let storage_value = rpc::state::get_storage(client, &storage_key, at).await?;
243			let Some(storage_value) = storage_value else {
244				return Ok(None);
245			};
246
247			let storage_value = Self::decode_storage_value(&mut storage_value.as_slice())?;
248			return Ok(Some(storage_value));
249		}
250	}
251
252	fn iter(client: RpcClient, key_1: &Self::KEY1, block_hash: H256) -> StorageDoubleMapIterator<Self>
253	where
254		Self: Sized,
255	{
256		StorageDoubleMapIterator::new(client, key_1, block_hash)
257	}
258}
259
260#[derive(Clone)]
261pub struct StorageMapIterator<T: StorageMap> {
262	client: RpcClient,
263	phantom: PhantomData<T>,
264	block_hash: H256,
265	fetched_keys: Vec<String>,
266	last_key: Option<String>,
267	is_done: bool,
268	prefix: String,
269}
270
271impl<T: StorageMap> StorageMapIterator<T> {
272	pub fn new(client: RpcClient, block_hash: H256) -> Self {
273		Self {
274			client,
275			phantom: PhantomData::<T>,
276			block_hash,
277			fetched_keys: Vec::new(),
278			last_key: None,
279			is_done: false,
280			prefix: hex::encode(T::encode_partial_key()),
281		}
282	}
283
284	pub async fn next_key_value(&mut self) -> Result<Option<(T::KEY, T::VALUE)>, Error> {
285		if self.is_done {
286			return Ok(None);
287		}
288
289		// Fetch new keys
290		if self.fetched_keys.is_empty() {
291			self.fetch_new_keys().await?;
292		}
293
294		let Some(storage_key) = self.fetched_keys.last() else {
295			return Ok(None);
296		};
297
298		let Some(storage_value) = self.fetch_storage_value(storage_key).await? else {
299			return Ok(None);
300		};
301
302		let key = hex::decode(storage_key.trim_start_matches("0x")).map_err(|x| x.to_string())?;
303		let key = T::decode_storage_key(&mut key.as_slice())?;
304
305		self.last_key = Some(storage_key.clone());
306		self.fetched_keys.pop();
307
308		Ok(Some((key, storage_value)))
309	}
310
311	pub async fn next(&mut self) -> Result<Option<T::VALUE>, Error> {
312		if self.is_done {
313			return Ok(None);
314		}
315
316		// Fetch new keys
317		if self.fetched_keys.is_empty() {
318			self.fetch_new_keys().await?;
319		}
320
321		let Some(storage_key) = self.fetched_keys.last() else {
322			return Ok(None);
323		};
324
325		let Some(storage_value) = self.fetch_storage_value(storage_key).await? else {
326			return Ok(None);
327		};
328
329		self.last_key = Some(storage_key.clone());
330		self.fetched_keys.pop();
331
332		Ok(Some(storage_value))
333	}
334
335	async fn fetch_new_keys(&mut self) -> Result<(), Error> {
336		self.fetched_keys = rpc::state::get_keys_paged(
337			&self.client,
338			Some(self.prefix.clone()),
339			100,
340			self.last_key.clone(),
341			Some(self.block_hash),
342		)
343		.await?;
344
345		self.fetched_keys.reverse();
346		if self.fetched_keys.is_empty() {
347			self.is_done = true
348		}
349
350		Ok(())
351	}
352
353	async fn fetch_storage_value(&self, key: &str) -> Result<Option<T::VALUE>, Error> {
354		let storage_value = rpc::state::get_storage(&self.client, key, Some(self.block_hash)).await?;
355		let Some(storage_value) = storage_value else {
356			return Ok(None);
357		};
358		let storage_value = T::decode_storage_value(&mut storage_value.as_slice())?;
359
360		Ok(Some(storage_value))
361	}
362}
363
364#[derive(Clone)]
365pub struct StorageDoubleMapIterator<T: StorageDoubleMap> {
366	client: RpcClient,
367	phantom: PhantomData<T>,
368	block_hash: H256,
369	fetched_keys: Vec<String>,
370	last_key: Option<String>,
371	is_done: bool,
372	prefix: String,
373}
374
375impl<T: StorageDoubleMap> StorageDoubleMapIterator<T> {
376	pub fn new(client: RpcClient, key_1: &T::KEY1, block_hash: H256) -> Self {
377		Self {
378			client,
379			phantom: PhantomData::<T>,
380			block_hash,
381			fetched_keys: Vec::new(),
382			last_key: None,
383			is_done: false,
384
385			prefix: hex::encode(T::encode_partial_key(key_1)),
386		}
387	}
388
389	pub async fn next_key_value(&mut self) -> Result<Option<(T::KEY1, T::KEY2, T::VALUE)>, Error> {
390		if self.is_done {
391			return Ok(None);
392		}
393
394		// Fetch new keys
395		if self.fetched_keys.is_empty() {
396			self.fetch_new_keys().await?;
397		}
398
399		let Some(storage_key) = self.fetched_keys.last() else {
400			return Ok(None);
401		};
402
403		let Some(storage_value) = self.fetch_storage_value(storage_key).await? else {
404			return Ok(None);
405		};
406
407		let key = hex::decode(storage_key.trim_start_matches("0x")).map_err(|x| x.to_string())?;
408		let (key1, key2) = T::decode_storage_key(&mut key.as_slice())?;
409
410		self.last_key = Some(storage_key.clone());
411		self.fetched_keys.pop();
412
413		Ok(Some((key1, key2, storage_value)))
414	}
415
416	pub async fn next(&mut self) -> Result<Option<T::VALUE>, Error> {
417		if self.is_done {
418			return Ok(None);
419		}
420
421		// Fetch new keys
422		if self.fetched_keys.is_empty() {
423			self.fetch_new_keys().await?;
424		}
425
426		let Some(storage_key) = self.fetched_keys.last() else {
427			return Ok(None);
428		};
429
430		let Some(storage_value) = self.fetch_storage_value(storage_key).await? else {
431			return Ok(None);
432		};
433
434		self.last_key = Some(storage_key.clone());
435		self.fetched_keys.pop();
436
437		Ok(Some(storage_value))
438	}
439
440	async fn fetch_new_keys(&mut self) -> Result<(), Error> {
441		self.fetched_keys = rpc::state::get_keys_paged(
442			&self.client,
443			Some(self.prefix.clone()),
444			100,
445			self.last_key.clone(),
446			Some(self.block_hash),
447		)
448		.await?;
449
450		self.fetched_keys.reverse();
451		if self.fetched_keys.is_empty() {
452			self.is_done = true
453		}
454
455		Ok(())
456	}
457
458	async fn fetch_storage_value(&self, key: &str) -> Result<Option<T::VALUE>, Error> {
459		let storage_value = rpc::state::get_storage(&self.client, key, Some(self.block_hash)).await?;
460		let Some(storage_value) = storage_value else {
461			return Ok(None);
462		};
463		let storage_value = T::decode_storage_value(&mut storage_value.as_slice())?;
464
465		Ok(Some(storage_value))
466	}
467}
avail_rust_core/decoded_storage.rs

avail_rust_core/
decoded_storage.rs
