add iterable functions, some test and fixes

frame/example/src/tests.rs

@@ -181,41 +181,6 @@ fn counted_map_works() {
 		assert_eq!(CountedMap::<Test>::count(), 0);
 		CountedMap::<Test>::insert(3, 3);
 		assert_eq!(CountedMap::<Test>::count(), 1);
-		CountedMap::<Test>::insert(3, 5);
-		assert_eq!(CountedMap::<Test>::count(), 1);
-		CountedMap::<Test>::insert(5, 5);
-		assert_eq!(CountedMap::<Test>::count(), 2);
-		CountedMap::<Test>::remove(5);
-		assert_eq!(CountedMap::<Test>::count(), 1);
-		// Does nothing
-		CountedMap::<Test>::take(4);
-		assert_eq!(CountedMap::<Test>::count(), 1);
-		CountedMap::<Test>::take(3);
-		assert_eq!(CountedMap::<Test>::count(), 0);
-
-		CountedMap::<Test>::mutate(3, |old| {
-			if old.is_none() {
-				*old = Some(69)
-			} else {
-				*old = Some(420)
-			}
-		});
-		assert_eq!(CountedMap::<Test>::get(3), Some(69));
-		assert_eq!(CountedMap::<Test>::count(), 1);
-
-		CountedMap::<Test>::mutate(3, |old| {
-			if old.is_none() {
-				*old = Some(69)
-			} else {
-				*old = Some(420)
-			}
-		});
-		assert_eq!(CountedMap::<Test>::get(3), Some(420));
-		assert_eq!(CountedMap::<Test>::count(), 1);
-
-		CountedMap::<Test>::mutate(3, |old| *old = None);
-		assert_eq!(CountedMap::<Test>::get(3), None);
-		assert_eq!(CountedMap::<Test>::count(), 0);
 	})
 }
 

frame/support/src/storage/generator/double_map.rs

@@ -225,6 +225,7 @@ impl<K1, K2, V, G> storage::StorageDoubleMap<K1, K2, V> for G where
 			previous_key: prefix,
 			drain: false,
 			closure: |_raw_key, mut raw_value| V::decode(&mut raw_value),
+			phantom: Default::default(),
 		}
 	}
 
@@ -352,6 +353,7 @@ impl<
 				let mut key_material = G::Hasher2::reverse(raw_key_without_prefix);
 				Ok((K2::decode(&mut key_material)?, V::decode(&mut raw_value)?))
 			},
+			phantom: Default::default(),
 		}
 	}
 
@@ -374,6 +376,7 @@ impl<
 				let k2 = K2::decode(&mut k2_material)?;
 				Ok((k1, k2, V::decode(&mut raw_value)?))
 			},
+			phantom: Default::default(),
 		}
 	}
 

frame/support/src/storage/generator/map.rs

@@ -152,6 +152,7 @@ impl<
 				let mut key_material = G::Hasher::reverse(raw_key_without_prefix);
 				Ok((K::decode(&mut key_material)?, V::decode(&mut raw_value)?))
 			},
+			phantom: Default::default(),
 		}
 	}
 

frame/support/src/storage/generator/nmap.rs

@@ -213,6 +213,7 @@ where
 			previous_key: prefix,
 			drain: false,
 			closure: |_raw_key, mut raw_value| V::decode(&mut raw_value),
+			phantom: Default::default(),
 		}
 	}
 
@@ -325,6 +326,7 @@ impl<K: ReversibleKeyGenerator, V: FullCodec, G: StorageNMap<K, V>>
 				let partial_key = K::decode_partial_key(raw_key_without_prefix)?;
 				Ok((partial_key, V::decode(&mut raw_value)?))
 			},
+			phantom: Default::default(),
 		}
 	}
 
@@ -347,6 +349,7 @@ impl<K: ReversibleKeyGenerator, V: FullCodec, G: StorageNMap<K, V>>
 				let (final_key, _) = K::decode_final_key(raw_key_without_prefix)?;
 				Ok((final_key, V::decode(&mut raw_value)?))
 			},
+			phantom: Default::default(),
 		}
 	}
 

frame/support/src/storage/migration.rs

@@ -175,7 +175,7 @@ pub fn storage_iter_with_suffix<T: Decode + Sized>(
 		Ok((raw_key_without_prefix.to_vec(), value))
 	};
 
-	PrefixIterator { prefix, previous_key, drain: false, closure }
+	PrefixIterator { prefix, previous_key, drain: false, closure, phantom: Default::default() }
 }
 
 /// Construct iterator to iterate over map items in `module` for the map called `item`.
@@ -203,7 +203,7 @@ pub fn storage_key_iter_with_suffix<K: Decode + Sized, T: Decode + Sized, H: Rev
 		let value = T::decode(&mut &raw_value[..])?;
 		Ok((key, value))
 	};
-	PrefixIterator { prefix, previous_key, drain: false, closure }
+	PrefixIterator { prefix, previous_key, drain: false, closure, phantom: Default::default() }
 }
 
 /// Get a particular value in storage by the `module`, the map's `item` name and the key `hash`.
@@ -332,11 +332,12 @@ pub fn move_prefix(from_prefix: &[u8], to_prefix: &[u8]) {
 		return
 	}
 
-	let iter = PrefixIterator {
+	let iter = PrefixIterator::<_> {
 		prefix: from_prefix.to_vec(),
 		previous_key: from_prefix.to_vec(),
 		drain: true,
 		closure: |key, value| Ok((key.to_vec(), value.to_vec())),
+		phantom: Default::default(),
 	};
 
 	for (key, value) in iter {

frame/support/src/storage/mod.rs

@@ -666,28 +666,41 @@ pub trait StorageNMap<K: KeyGenerator, V: FullCodec> {
 		KArg: EncodeLikeTuple<K::KArg> + TupleToEncodedIter;
 }
 
-/// Iterate over a prefix and decode raw_key and raw_value into `T`.
+/// Iterate or drain over a prefix and decode raw_key and raw_value into `T`.
 ///
 /// If any decoding fails it skips it and continues to the next key.
-pub struct PrefixIterator<T> {
+///
+/// If draining, then the hook `OnRemoval::on_removal` is called after each removal.
+pub struct PrefixIterator<T, OnRemoval=()> {
 	prefix: Vec<u8>,
 	previous_key: Vec<u8>,
 	/// If true then value are removed while iterating
 	drain: bool,
 	/// Function that take `(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>,
+	phantom: core::marker::PhantomData<OnRemoval>,
+}
+
+/// Trait for specialising on removal logic of [`PrefixIterator`].
+pub trait PrefixIteratorOnRemoval {
+	fn on_removal();
+}
+
+/// No-op implement.
+impl PrefixIteratorOnRemoval for () {
+	fn on_removal() {}
 }
 
-impl<T> PrefixIterator<T> {
+impl<T, OnRemoval> PrefixIterator<T, OnRemoval> {
 	/// 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> {
+impl<T, OnRemoval: PrefixIteratorOnRemoval> Iterator for PrefixIterator<T, OnRemoval> {
 	type Item = T;
 
 	fn next(&mut self) -> Option<Self::Item> {
@@ -708,7 +721,8 @@ impl<T> Iterator for PrefixIterator<T> {
 						}
 					};
 					if self.drain {
-						unhashed::kill(&self.previous_key)
+						unhashed::kill(&self.previous_key);
+						OnRemoval::on_removal();
 					}
 					let raw_key_without_prefix = &self.previous_key[self.prefix.len()..];
 					let item = match (self.closure)(raw_key_without_prefix, &raw_value[..]) {
@@ -894,6 +908,7 @@ pub trait StoragePrefixedMap<Value: FullCodec> {
 			previous_key: prefix.to_vec(),
 			drain: false,
 			closure: |_raw_key, mut raw_value| Value::decode(&mut raw_value),
+			phantom: Default::default(),
 		}
 	}