1use crate::{
21 hash::ReversibleStorageHasher,
22 storage::{storage_prefix, unhashed},
23 StorageHasher, Twox128,
24};
25use alloc::{vec, vec::Vec};
26use codec::{Decode, Encode};
27
28use super::PrefixIterator;
29
30pub fn storage_iter<T: Decode + Sized>(module: &[u8], item: &[u8]) -> PrefixIterator<(Vec<u8>, T)> {
32 storage_iter_with_suffix(module, item, &[][..])
33}
34
35pub fn storage_iter_with_suffix<T: Decode + Sized>(
37 module: &[u8],
38 item: &[u8],
39 suffix: &[u8],
40) -> PrefixIterator<(Vec<u8>, T)> {
41 let mut prefix = Vec::new();
42 let storage_prefix = storage_prefix(module, item);
43 prefix.extend_from_slice(&storage_prefix);
44 prefix.extend_from_slice(suffix);
45 let previous_key = prefix.clone();
46 let closure = |raw_key_without_prefix: &[u8], mut raw_value: &[u8]| {
47 let value = T::decode(&mut raw_value)?;
48 Ok((raw_key_without_prefix.to_vec(), value))
49 };
50
51 PrefixIterator { prefix, previous_key, drain: false, closure, phantom: Default::default() }
52}
53
54pub fn storage_key_iter<K: Decode + Sized, T: Decode + Sized, H: ReversibleStorageHasher>(
56 module: &[u8],
57 item: &[u8],
58) -> PrefixIterator<(K, T)> {
59 storage_key_iter_with_suffix::<K, T, H>(module, item, &[][..])
60}
61
62pub fn storage_key_iter_with_suffix<
64 K: Decode + Sized,
65 T: Decode + Sized,
66 H: ReversibleStorageHasher,
67>(
68 module: &[u8],
69 item: &[u8],
70 suffix: &[u8],
71) -> PrefixIterator<(K, T)> {
72 let mut prefix = Vec::new();
73 let storage_prefix = storage_prefix(module, item);
74
75 prefix.extend_from_slice(&storage_prefix);
76 prefix.extend_from_slice(suffix);
77 let previous_key = prefix.clone();
78 let closure = |raw_key_without_prefix: &[u8], mut raw_value: &[u8]| {
79 let mut key_material = H::reverse(raw_key_without_prefix);
80 let key = K::decode(&mut key_material)?;
81 let value = T::decode(&mut raw_value)?;
82 Ok((key, value))
83 };
84 PrefixIterator { prefix, previous_key, drain: false, closure, phantom: Default::default() }
85}
86
87pub fn have_storage_value(module: &[u8], item: &[u8], hash: &[u8]) -> bool {
89 get_storage_value::<()>(module, item, hash).is_some()
90}
91
92pub fn get_storage_value<T: Decode + Sized>(module: &[u8], item: &[u8], hash: &[u8]) -> Option<T> {
94 let mut key = vec![0u8; 32 + hash.len()];
95 let storage_prefix = storage_prefix(module, item);
96 key[0..32].copy_from_slice(&storage_prefix);
97 key[32..].copy_from_slice(hash);
98 frame_support::storage::unhashed::get::<T>(&key)
99}
100
101pub fn take_storage_value<T: Decode + Sized>(module: &[u8], item: &[u8], hash: &[u8]) -> Option<T> {
103 let mut key = vec![0u8; 32 + hash.len()];
104 let storage_prefix = storage_prefix(module, item);
105 key[0..32].copy_from_slice(&storage_prefix);
106 key[32..].copy_from_slice(hash);
107 frame_support::storage::unhashed::take::<T>(&key)
108}
109
110pub fn put_storage_value<T: Encode>(module: &[u8], item: &[u8], hash: &[u8], value: T) {
112 let mut key = vec![0u8; 32 + hash.len()];
113 let storage_prefix = storage_prefix(module, item);
114 key[0..32].copy_from_slice(&storage_prefix);
115 key[32..].copy_from_slice(hash);
116 frame_support::storage::unhashed::put(&key, &value);
117}
118
119pub fn clear_storage_prefix(
136 module: &[u8],
137 item: &[u8],
138 hash: &[u8],
139 maybe_limit: Option<u32>,
140 maybe_cursor: Option<&[u8]>,
141) -> sp_io::MultiRemovalResults {
142 let mut key = vec![0u8; 32 + hash.len()];
143 let storage_prefix = storage_prefix(module, item);
144 key[0..32].copy_from_slice(&storage_prefix);
145 key[32..].copy_from_slice(hash);
146 frame_support::storage::unhashed::clear_prefix(&key, maybe_limit, maybe_cursor)
147}
148
149pub fn take_storage_item<K: Encode + Sized, T: Decode + Sized, H: StorageHasher>(
151 module: &[u8],
152 item: &[u8],
153 key: K,
154) -> Option<T> {
155 take_storage_value(module, item, key.using_encoded(H::hash).as_ref())
156}
157
158pub fn move_storage_from_pallet(
179 storage_name: &[u8],
180 old_pallet_name: &[u8],
181 new_pallet_name: &[u8],
182) {
183 let new_prefix = storage_prefix(new_pallet_name, storage_name);
184 let old_prefix = storage_prefix(old_pallet_name, storage_name);
185
186 move_prefix(&old_prefix, &new_prefix);
187
188 if let Some(value) = unhashed::get_raw(&old_prefix) {
189 unhashed::put_raw(&new_prefix, &value);
190 unhashed::kill(&old_prefix);
191 }
192}
193
194pub fn move_pallet(old_pallet_name: &[u8], new_pallet_name: &[u8]) {
215 move_prefix(&Twox128::hash(old_pallet_name), &Twox128::hash(new_pallet_name))
216}
217
218pub fn move_prefix(from_prefix: &[u8], to_prefix: &[u8]) {
225 if from_prefix == to_prefix {
226 return;
227 }
228
229 let iter = PrefixIterator::<_> {
230 prefix: from_prefix.to_vec(),
231 previous_key: from_prefix.to_vec(),
232 drain: true,
233 closure: |key, value| Ok((key.to_vec(), value.to_vec())),
234 phantom: Default::default(),
235 };
236
237 for (key, value) in iter {
238 let full_key = [to_prefix, &key].concat();
239 unhashed::put_raw(&full_key, &value);
240 }
241}
242
243#[cfg(test)]
244mod tests {
245 use super::{
246 move_pallet, move_prefix, move_storage_from_pallet, storage_iter, storage_key_iter,
247 };
248 use crate::{
249 hash::StorageHasher,
250 pallet_prelude::{StorageMap, StorageValue, Twox128, Twox64Concat},
251 };
252 use sp_io::TestExternalities;
253
254 struct OldPalletStorageValuePrefix;
255 impl frame_support::traits::StorageInstance for OldPalletStorageValuePrefix {
256 const STORAGE_PREFIX: &'static str = "foo_value";
257 fn pallet_prefix() -> &'static str {
258 "my_old_pallet"
259 }
260 }
261 type OldStorageValue = StorageValue<OldPalletStorageValuePrefix, u32>;
262
263 struct OldPalletStorageMapPrefix;
264 impl frame_support::traits::StorageInstance for OldPalletStorageMapPrefix {
265 const STORAGE_PREFIX: &'static str = "foo_map";
266 fn pallet_prefix() -> &'static str {
267 "my_old_pallet"
268 }
269 }
270 type OldStorageMap = StorageMap<OldPalletStorageMapPrefix, Twox64Concat, u32, u32>;
271
272 struct NewPalletStorageValuePrefix;
273 impl frame_support::traits::StorageInstance for NewPalletStorageValuePrefix {
274 const STORAGE_PREFIX: &'static str = "foo_value";
275 fn pallet_prefix() -> &'static str {
276 "my_new_pallet"
277 }
278 }
279 type NewStorageValue = StorageValue<NewPalletStorageValuePrefix, u32>;
280
281 struct NewPalletStorageMapPrefix;
282 impl frame_support::traits::StorageInstance for NewPalletStorageMapPrefix {
283 const STORAGE_PREFIX: &'static str = "foo_map";
284 fn pallet_prefix() -> &'static str {
285 "my_new_pallet"
286 }
287 }
288 type NewStorageMap = StorageMap<NewPalletStorageMapPrefix, Twox64Concat, u32, u32>;
289
290 #[test]
291 fn test_move_prefix() {
292 TestExternalities::new_empty().execute_with(|| {
293 OldStorageValue::put(3);
294 OldStorageMap::insert(1, 2);
295 OldStorageMap::insert(3, 4);
296
297 move_prefix(&Twox128::hash(b"my_old_pallet"), &Twox128::hash(b"my_new_pallet"));
298
299 assert_eq!(OldStorageValue::get(), None);
300 assert_eq!(OldStorageMap::iter().collect::<Vec<_>>(), vec![]);
301 assert_eq!(NewStorageValue::get(), Some(3));
302 assert_eq!(NewStorageMap::iter().collect::<Vec<_>>(), vec![(1, 2), (3, 4)]);
303 })
304 }
305
306 #[test]
307 fn test_move_storage() {
308 TestExternalities::new_empty().execute_with(|| {
309 OldStorageValue::put(3);
310 OldStorageMap::insert(1, 2);
311 OldStorageMap::insert(3, 4);
312
313 move_storage_from_pallet(b"foo_map", b"my_old_pallet", b"my_new_pallet");
314
315 assert_eq!(OldStorageValue::get(), Some(3));
316 assert_eq!(OldStorageMap::iter().collect::<Vec<_>>(), vec![]);
317 assert_eq!(NewStorageValue::get(), None);
318 assert_eq!(NewStorageMap::iter().collect::<Vec<_>>(), vec![(1, 2), (3, 4)]);
319
320 move_storage_from_pallet(b"foo_value", b"my_old_pallet", b"my_new_pallet");
321
322 assert_eq!(OldStorageValue::get(), None);
323 assert_eq!(OldStorageMap::iter().collect::<Vec<_>>(), vec![]);
324 assert_eq!(NewStorageValue::get(), Some(3));
325 assert_eq!(NewStorageMap::iter().collect::<Vec<_>>(), vec![(1, 2), (3, 4)]);
326 })
327 }
328
329 #[test]
330 fn test_move_pallet() {
331 TestExternalities::new_empty().execute_with(|| {
332 OldStorageValue::put(3);
333 OldStorageMap::insert(1, 2);
334 OldStorageMap::insert(3, 4);
335
336 move_pallet(b"my_old_pallet", b"my_new_pallet");
337
338 assert_eq!(OldStorageValue::get(), None);
339 assert_eq!(OldStorageMap::iter().collect::<Vec<_>>(), vec![]);
340 assert_eq!(NewStorageValue::get(), Some(3));
341 assert_eq!(NewStorageMap::iter().collect::<Vec<_>>(), vec![(1, 2), (3, 4)]);
342 })
343 }
344
345 #[test]
346 fn test_storage_iter() {
347 TestExternalities::new_empty().execute_with(|| {
348 OldStorageValue::put(3);
349 OldStorageMap::insert(1, 2);
350 OldStorageMap::insert(3, 4);
351
352 assert_eq!(
353 storage_key_iter::<i32, i32, Twox64Concat>(b"my_old_pallet", b"foo_map")
354 .collect::<Vec<_>>(),
355 vec![(1, 2), (3, 4)],
356 );
357
358 assert_eq!(
359 storage_iter(b"my_old_pallet", b"foo_map")
360 .drain()
361 .map(|t| t.1)
362 .collect::<Vec<i32>>(),
363 vec![2, 4],
364 );
365 assert_eq!(OldStorageMap::iter().collect::<Vec<_>>(), vec![]);
366
367 assert_eq!(storage_iter::<i32>(b"my_old_pallet", b"foo_value").drain().next(), None);
369 assert_eq!(OldStorageValue::get(), Some(3));
370 });
371 }
372}