frame_support/storage/generator/
map.rs

1// This file is part of Substrate.
2
3// Copyright (C) Parity Technologies (UK) Ltd.
4// SPDX-License-Identifier: Apache-2.0
5
6// Licensed under the Apache License, Version 2.0 (the "License");
7// you may not use this file except in compliance with the License.
8// You may obtain a copy of the License at
9//
10// 	http://www.apache.org/licenses/LICENSE-2.0
11//
12// Unless required by applicable law or agreed to in writing, software
13// distributed under the License is distributed on an "AS IS" BASIS,
14// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15// See the License for the specific language governing permissions and
16// limitations under the License.
17
18use crate::{
19	hash::{ReversibleStorageHasher, StorageHasher},
20	storage::{self, storage_prefix, unhashed, KeyPrefixIterator, PrefixIterator, StorageAppend},
21	Never,
22};
23use alloc::vec::Vec;
24use codec::{Decode, Encode, EncodeLike, FullCodec, FullEncode};
25
26/// Generator for `StorageMap` used by `decl_storage`.
27///
28/// By default each key value is stored at:
29/// ```nocompile
30/// Twox128(pallet_prefix) ++ Twox128(storage_prefix) ++ Hasher(encode(key))
31/// ```
32///
33/// # Warning
34///
35/// If the keys are not trusted (e.g. can be set by a user), a cryptographic `hasher` such as
36/// `blake2_256` must be used.  Otherwise, other values in storage can be compromised.
37pub trait StorageMap<K: FullEncode, V: FullCodec> {
38	/// The type that get/take returns.
39	type Query;
40
41	/// Hasher. Used for generating final key.
42	type Hasher: StorageHasher;
43
44	/// Pallet prefix. Used for generating final key.
45	fn pallet_prefix() -> &'static [u8];
46
47	/// Storage prefix. Used for generating final key.
48	fn storage_prefix() -> &'static [u8];
49
50	/// The full prefix; just the hash of `pallet_prefix` concatenated to the hash of
51	/// `storage_prefix`.
52	fn prefix_hash() -> [u8; 32];
53
54	/// Convert an optional value retrieved from storage to the type queried.
55	fn from_optional_value_to_query(v: Option<V>) -> Self::Query;
56
57	/// Convert a query to an optional value into storage.
58	fn from_query_to_optional_value(v: Self::Query) -> Option<V>;
59
60	/// Generate the full key used in top storage.
61	fn storage_map_final_key<KeyArg>(key: KeyArg) -> Vec<u8>
62	where
63		KeyArg: EncodeLike<K>,
64	{
65		let storage_prefix = storage_prefix(Self::pallet_prefix(), Self::storage_prefix());
66		let key_hashed = key.using_encoded(Self::Hasher::hash);
67
68		let mut final_key = Vec::with_capacity(storage_prefix.len() + key_hashed.as_ref().len());
69
70		final_key.extend_from_slice(&storage_prefix);
71		final_key.extend_from_slice(key_hashed.as_ref());
72
73		final_key
74	}
75}
76
77impl<K: FullCodec, V: FullCodec, G: StorageMap<K, V>> storage::IterableStorageMap<K, V> for G
78where
79	G::Hasher: ReversibleStorageHasher,
80{
81	type Iterator = PrefixIterator<(K, V)>;
82	type KeyIterator = KeyPrefixIterator<K>;
83
84	/// Enumerate all elements in the map.
85	fn iter() -> Self::Iterator {
86		let prefix = G::prefix_hash().to_vec();
87		PrefixIterator {
88			prefix: prefix.clone(),
89			previous_key: prefix,
90			drain: false,
91			closure: |raw_key_without_prefix, mut raw_value| {
92				let mut key_material = G::Hasher::reverse(raw_key_without_prefix);
93				Ok((K::decode(&mut key_material)?, V::decode(&mut raw_value)?))
94			},
95			phantom: Default::default(),
96		}
97	}
98
99	/// Enumerate all elements in the map after a given key.
100	fn iter_from(starting_raw_key: Vec<u8>) -> Self::Iterator {
101		let mut iter = Self::iter();
102		iter.set_last_raw_key(starting_raw_key);
103		iter
104	}
105
106	/// Enumerate all keys in the map.
107	fn iter_keys() -> Self::KeyIterator {
108		let prefix = G::prefix_hash().to_vec();
109		KeyPrefixIterator {
110			prefix: prefix.clone(),
111			previous_key: prefix,
112			drain: false,
113			closure: |raw_key_without_prefix| {
114				let mut key_material = G::Hasher::reverse(raw_key_without_prefix);
115				K::decode(&mut key_material)
116			},
117		}
118	}
119
120	/// Enumerate all keys in the map after a given key.
121	fn iter_keys_from(starting_raw_key: Vec<u8>) -> Self::KeyIterator {
122		let mut iter = Self::iter_keys();
123		iter.set_last_raw_key(starting_raw_key);
124		iter
125	}
126
127	/// Enumerate all elements in the map.
128	fn drain() -> Self::Iterator {
129		let mut iterator = Self::iter();
130		iterator.drain = true;
131		iterator
132	}
133
134	fn translate<O: Decode, F: FnMut(K, O) -> Option<V>>(mut f: F) {
135		let mut previous_key = None;
136		loop {
137			previous_key = Self::translate_next(previous_key, &mut f);
138			if previous_key.is_none() {
139				break;
140			}
141		}
142	}
143
144	fn translate_next<O: Decode, F: FnMut(K, O) -> Option<V>>(
145		previous_key: Option<Vec<u8>>,
146		mut f: F,
147	) -> Option<Vec<u8>> {
148		let prefix = G::prefix_hash().to_vec();
149		let previous_key = previous_key.unwrap_or_else(|| prefix.clone());
150
151		let current_key =
152			sp_io::storage::next_key(&previous_key).filter(|n| n.starts_with(&prefix))?;
153
154		let value = match unhashed::get::<O>(&current_key) {
155			Some(value) => value,
156			None => {
157				crate::defensive!(
158					"Invalid translation: failed to decode old value for key",
159					array_bytes::bytes2hex("0x", &current_key)
160				);
161				return Some(current_key);
162			},
163		};
164
165		let mut key_material = G::Hasher::reverse(&current_key[prefix.len()..]);
166		let key = match K::decode(&mut key_material) {
167			Ok(key) => key,
168			Err(_) => {
169				crate::defensive!(
170					"Invalid translation: failed to decode key",
171					array_bytes::bytes2hex("0x", &current_key)
172				);
173				return Some(current_key);
174			},
175		};
176
177		match f(key, value) {
178			Some(new) => unhashed::put::<V>(&current_key, &new),
179			None => unhashed::kill(&current_key),
180		}
181
182		Some(current_key)
183	}
184}
185
186impl<K: FullEncode, V: FullCodec, G: StorageMap<K, V>> storage::StorageMap<K, V> for G {
187	type Query = G::Query;
188
189	fn hashed_key_for<KeyArg: EncodeLike<K>>(key: KeyArg) -> Vec<u8> {
190		Self::storage_map_final_key(key)
191	}
192
193	fn swap<KeyArg1: EncodeLike<K>, KeyArg2: EncodeLike<K>>(key1: KeyArg1, key2: KeyArg2) {
194		let k1 = Self::storage_map_final_key(key1);
195		let k2 = Self::storage_map_final_key(key2);
196
197		let v1 = unhashed::get_raw(k1.as_ref());
198		if let Some(val) = unhashed::get_raw(k2.as_ref()) {
199			unhashed::put_raw(k1.as_ref(), &val);
200		} else {
201			unhashed::kill(k1.as_ref())
202		}
203		if let Some(val) = v1 {
204			unhashed::put_raw(k2.as_ref(), &val);
205		} else {
206			unhashed::kill(k2.as_ref())
207		}
208	}
209
210	fn contains_key<KeyArg: EncodeLike<K>>(key: KeyArg) -> bool {
211		unhashed::exists(Self::storage_map_final_key(key).as_ref())
212	}
213
214	fn get<KeyArg: EncodeLike<K>>(key: KeyArg) -> Self::Query {
215		G::from_optional_value_to_query(unhashed::get(Self::storage_map_final_key(key).as_ref()))
216	}
217
218	fn try_get<KeyArg: EncodeLike<K>>(key: KeyArg) -> Result<V, ()> {
219		unhashed::get(Self::storage_map_final_key(key).as_ref()).ok_or(())
220	}
221
222	fn set<KeyArg: EncodeLike<K>>(key: KeyArg, q: Self::Query) {
223		match G::from_query_to_optional_value(q) {
224			Some(v) => Self::insert(key, v),
225			None => Self::remove(key),
226		}
227	}
228
229	fn insert<KeyArg: EncodeLike<K>, ValArg: EncodeLike<V>>(key: KeyArg, val: ValArg) {
230		unhashed::put(Self::storage_map_final_key(key).as_ref(), &val)
231	}
232
233	fn remove<KeyArg: EncodeLike<K>>(key: KeyArg) {
234		unhashed::kill(Self::storage_map_final_key(key).as_ref())
235	}
236
237	fn mutate<KeyArg: EncodeLike<K>, R, F: FnOnce(&mut Self::Query) -> R>(key: KeyArg, f: F) -> R {
238		Self::try_mutate(key, |v| Ok::<R, Never>(f(v)))
239			.expect("`Never` can not be constructed; qed")
240	}
241
242	fn mutate_exists<KeyArg: EncodeLike<K>, R, F: FnOnce(&mut Option<V>) -> R>(
243		key: KeyArg,
244		f: F,
245	) -> R {
246		Self::try_mutate_exists(key, |v| Ok::<R, Never>(f(v)))
247			.expect("`Never` can not be constructed; qed")
248	}
249
250	fn try_mutate<KeyArg: EncodeLike<K>, R, E, F: FnOnce(&mut Self::Query) -> Result<R, E>>(
251		key: KeyArg,
252		f: F,
253	) -> Result<R, E> {
254		let final_key = Self::storage_map_final_key(key);
255		let mut val = G::from_optional_value_to_query(unhashed::get(final_key.as_ref()));
256
257		let ret = f(&mut val);
258		if ret.is_ok() {
259			match G::from_query_to_optional_value(val) {
260				Some(ref val) => unhashed::put(final_key.as_ref(), &val),
261				None => unhashed::kill(final_key.as_ref()),
262			}
263		}
264		ret
265	}
266
267	fn try_mutate_exists<KeyArg: EncodeLike<K>, R, E, F: FnOnce(&mut Option<V>) -> Result<R, E>>(
268		key: KeyArg,
269		f: F,
270	) -> Result<R, E> {
271		let final_key = Self::storage_map_final_key(key);
272		let mut val = unhashed::get(final_key.as_ref());
273
274		let ret = f(&mut val);
275		if ret.is_ok() {
276			match val {
277				Some(ref val) => unhashed::put(final_key.as_ref(), &val),
278				None => unhashed::kill(final_key.as_ref()),
279			}
280		}
281		ret
282	}
283
284	fn take<KeyArg: EncodeLike<K>>(key: KeyArg) -> Self::Query {
285		let key = Self::storage_map_final_key(key);
286		let value = unhashed::take(key.as_ref());
287		G::from_optional_value_to_query(value)
288	}
289
290	fn append<Item, EncodeLikeItem, EncodeLikeKey>(key: EncodeLikeKey, item: EncodeLikeItem)
291	where
292		EncodeLikeKey: EncodeLike<K>,
293		Item: Encode,
294		EncodeLikeItem: EncodeLike<Item>,
295		V: StorageAppend<Item>,
296	{
297		let key = Self::storage_map_final_key(key);
298		sp_io::storage::append(&key, item.encode());
299	}
300
301	fn migrate_key<OldHasher: StorageHasher, KeyArg: EncodeLike<K>>(key: KeyArg) -> Option<V> {
302		let old_key = {
303			let storage_prefix = storage_prefix(Self::pallet_prefix(), Self::storage_prefix());
304			let key_hashed = key.using_encoded(OldHasher::hash);
305
306			let mut final_key =
307				Vec::with_capacity(storage_prefix.len() + key_hashed.as_ref().len());
308
309			final_key.extend_from_slice(&storage_prefix);
310			final_key.extend_from_slice(key_hashed.as_ref());
311
312			final_key
313		};
314		unhashed::take(old_key.as_ref()).inspect(|value| {
315			unhashed::put(Self::storage_map_final_key(key).as_ref(), &value);
316		})
317	}
318}
319
320/// Test iterators for StorageMap
321#[cfg(test)]
322mod test_iterators {
323	use crate::{
324		hash::StorageHasher,
325		storage::{
326			generator::{tests::*, StorageMap},
327			unhashed,
328		},
329	};
330	use alloc::vec;
331	use codec::Encode;
332
333	#[test]
334	fn map_iter_from() {
335		sp_io::TestExternalities::default().execute_with(|| {
336			use crate::hash::Identity;
337			#[crate::storage_alias]
338			type MyMap = StorageMap<MyModule, Identity, u64, u64>;
339
340			MyMap::insert(1, 10);
341			MyMap::insert(2, 20);
342			MyMap::insert(3, 30);
343			MyMap::insert(4, 40);
344			MyMap::insert(5, 50);
345
346			let starting_raw_key = MyMap::storage_map_final_key(3);
347			let iter = MyMap::iter_from(starting_raw_key);
348			assert_eq!(iter.collect::<Vec<_>>(), vec![(4, 40), (5, 50)]);
349
350			let starting_raw_key = MyMap::storage_map_final_key(2);
351			let iter = MyMap::iter_keys_from(starting_raw_key);
352			assert_eq!(iter.collect::<Vec<_>>(), vec![3, 4, 5]);
353		});
354	}
355
356	#[cfg(debug_assertions)]
357	#[test]
358	#[should_panic]
359	fn map_translate_with_bad_key_in_debug_mode() {
360		sp_io::TestExternalities::default().execute_with(|| {
361			type Map = self::frame_system::Map<Runtime>;
362			let prefix = Map::prefix_hash().to_vec();
363
364			// Wrong key
365			unhashed::put(&[prefix.clone(), vec![1, 2, 3]].concat(), &3u64.encode());
366
367			// debug_assert should cause a
368			Map::translate(|_k1, v: u64| Some(v * 2));
369			assert_eq!(Map::iter().collect::<Vec<_>>(), vec![(3, 6), (0, 0), (2, 4), (1, 2)]);
370		})
371	}
372
373	#[cfg(debug_assertions)]
374	#[test]
375	#[should_panic]
376	fn map_translate_with_bad_value_in_debug_mode() {
377		sp_io::TestExternalities::default().execute_with(|| {
378			type Map = self::frame_system::Map<Runtime>;
379			let prefix = Map::prefix_hash().to_vec();
380
381			// Wrong value
382			unhashed::put(
383				&[prefix.clone(), crate::Blake2_128Concat::hash(&6u16.encode())].concat(),
384				&vec![1],
385			);
386
387			Map::translate(|_k1, v: u64| Some(v * 2));
388			assert_eq!(Map::iter().collect::<Vec<_>>(), vec![(3, 6), (0, 0), (2, 4), (1, 2)]);
389		})
390	}
391
392	#[cfg(not(debug_assertions))]
393	#[test]
394	fn map_translate_with_bad_key_in_production_mode() {
395		sp_io::TestExternalities::default().execute_with(|| {
396			type Map = self::frame_system::Map<Runtime>;
397			let prefix = Map::prefix_hash().to_vec();
398
399			// Wrong key
400			unhashed::put(&[prefix.clone(), vec![1, 2, 3]].concat(), &3u64.encode());
401
402			Map::translate(|_k1, v: u64| Some(v * 2));
403			assert_eq!(Map::iter().collect::<Vec<_>>(), vec![]);
404		})
405	}
406
407	#[cfg(not(debug_assertions))]
408	#[test]
409	fn map_translate_with_bad_value_in_production_mode() {
410		sp_io::TestExternalities::default().execute_with(|| {
411			type Map = self::frame_system::Map<Runtime>;
412			let prefix = Map::prefix_hash().to_vec();
413
414			// Wrong value
415			unhashed::put(
416				&[prefix.clone(), crate::Blake2_128Concat::hash(&6u16.encode())].concat(),
417				&vec![1],
418			);
419
420			Map::translate(|_k1, v: u64| Some(v * 2));
421			assert_eq!(Map::iter().collect::<Vec<_>>(), vec![]);
422		})
423	}
424
425	#[test]
426	fn map_reversible_reversible_iteration() {
427		sp_io::TestExternalities::default().execute_with(|| {
428			type Map = self::frame_system::Map<Runtime>;
429
430			// All map iterator
431			let prefix = Map::prefix_hash().to_vec();
432
433			unhashed::put(&key_before_prefix(prefix.clone()), &1u64);
434			unhashed::put(&key_after_prefix(prefix.clone()), &1u64);
435
436			for i in 0..4 {
437				Map::insert(i as u16, i as u64);
438			}
439
440			assert_eq!(Map::iter().collect::<Vec<_>>(), vec![(3, 3), (0, 0), (2, 2), (1, 1)]);
441
442			assert_eq!(Map::iter_keys().collect::<Vec<_>>(), vec![3, 0, 2, 1]);
443
444			assert_eq!(Map::iter_values().collect::<Vec<_>>(), vec![3, 0, 2, 1]);
445
446			assert_eq!(Map::drain().collect::<Vec<_>>(), vec![(3, 3), (0, 0), (2, 2), (1, 1)]);
447
448			assert_eq!(Map::iter().collect::<Vec<_>>(), vec![]);
449			assert_eq!(unhashed::get(&key_before_prefix(prefix.clone())), Some(1u64));
450			assert_eq!(unhashed::get(&key_after_prefix(prefix.clone())), Some(1u64));
451
452			// Translate
453			let prefix = Map::prefix_hash().to_vec();
454
455			unhashed::put(&key_before_prefix(prefix.clone()), &1u64);
456			unhashed::put(&key_after_prefix(prefix.clone()), &1u64);
457			for i in 0..4 {
458				Map::insert(i as u16, i as u64);
459			}
460
461			Map::translate(|_k1, v: u64| Some(v * 2));
462			assert_eq!(Map::iter().collect::<Vec<_>>(), vec![(3, 6), (0, 0), (2, 4), (1, 2)]);
463		})
464	}
465}
frame_support/storage/generator/map.rs

frame_support/storage/generator/
map.rs
