From 28f56b6938f8a5be7bc7d6b94cba5a919b59545d Mon Sep 17 00:00:00 2001
From: PG Herveou <pgherveou@gmail.com>
Date: Fri, 2 Jun 2023 16:11:11 +0200
Subject: [PATCH] frame-support Add translate_next (#14043)

* Frame Add translate_next

This works similarly to to `translate` but only translate a single entry.
This function will be useful in the context of multi-block migration.

* Add test

* add None return case

* fixes

* PR comment use `?`
---
 frame/support/src/storage/generator/map.rs | 64 +++++++++++++---------
 frame/support/src/storage/mod.rs           |  9 +++
 frame/support/src/storage/types/map.rs     | 11 +++-
 3 files changed, 58 insertions(+), 26 deletions(-)

diff --git a/frame/support/src/storage/generator/map.rs b/frame/support/src/storage/generator/map.rs
index 3cd7bea2fe7b2..90fac4b41c759 100644
--- a/frame/support/src/storage/generator/map.rs
+++ b/frame/support/src/storage/generator/map.rs
@@ -178,34 +178,48 @@ where
 	}
 
 	fn translate<O: Decode, F: FnMut(K, O) -> Option<V>>(mut f: F) {
+		let mut previous_key = None;
+		loop {
+			previous_key = Self::translate_next(previous_key, &mut f);
+			if previous_key.is_none() {
+				break
+			}
+		}
+	}
+
+	fn translate_next<O: Decode, F: FnMut(K, O) -> Option<V>>(
+		previous_key: Option<Vec<u8>>,
+		mut f: F,
+	) -> Option<Vec<u8>> {
 		let prefix = G::prefix_hash();
-		let mut previous_key = prefix.clone();
-		while let Some(next) =
-			sp_io::storage::next_key(&previous_key).filter(|n| n.starts_with(&prefix))
-		{
-			previous_key = next;
-			let value = match unhashed::get::<O>(&previous_key) {
-				Some(value) => value,
-				None => {
-					log::error!("Invalid translate: fail to decode old value");
-					continue
-				},
-			};
-
-			let mut key_material = G::Hasher::reverse(&previous_key[prefix.len()..]);
-			let key = match K::decode(&mut key_material) {
-				Ok(key) => key,
-				Err(_) => {
-					log::error!("Invalid translate: fail to decode key");
-					continue
-				},
-			};
+		let previous_key = previous_key.unwrap_or_else(|| prefix.clone());
 
-			match f(key, value) {
-				Some(new) => unhashed::put::<V>(&previous_key, &new),
-				None => unhashed::kill(&previous_key),
-			}
+		let current_key =
+			sp_io::storage::next_key(&previous_key).filter(|n| n.starts_with(&prefix))?;
+
+		let value = match unhashed::get::<O>(&current_key) {
+			Some(value) => value,
+			None => {
+				log::error!("Invalid translate: fail to decode old value");
+				return Some(current_key)
+			},
+		};
+
+		let mut key_material = G::Hasher::reverse(&current_key[prefix.len()..]);
+		let key = match K::decode(&mut key_material) {
+			Ok(key) => key,
+			Err(_) => {
+				log::error!("Invalid translate: fail to decode key");
+				return Some(current_key)
+			},
+		};
+
+		match f(key, value) {
+			Some(new) => unhashed::put::<V>(&current_key, &new),
+			None => unhashed::kill(&current_key),
 		}
+
+		Some(current_key)
 	}
 }
 
diff --git a/frame/support/src/storage/mod.rs b/frame/support/src/storage/mod.rs
index 4c6ea943c6920..6cb8b5915ca85 100644
--- a/frame/support/src/storage/mod.rs
+++ b/frame/support/src/storage/mod.rs
@@ -303,6 +303,15 @@ pub trait IterableStorageMap<K: FullEncode, V: FullCodec>: StorageMap<K, V> {
 	///
 	/// NOTE: If a value fail to decode because storage is corrupted then it is skipped.
 	fn translate<O: Decode, F: FnMut(K, O) -> Option<V>>(f: F);
+
+	/// Translate the next entry following `previous_key` by a function `f`.
+	/// By returning `None` from `f` for an element, you'll remove it from the map.
+	///
+	/// Returns the next key to iterate from in lexicographical order of the encoded key.
+	fn translate_next<O: Decode, F: FnMut(K, O) -> Option<V>>(
+		previous_key: Option<Vec<u8>>,
+		f: F,
+	) -> Option<Vec<u8>>;
 }
 
 /// A strongly-typed double map in storage whose secondary keys and values can be iterated over.
diff --git a/frame/support/src/storage/types/map.rs b/frame/support/src/storage/types/map.rs
index 2110732b2f69c..ba6615bb8d381 100644
--- a/frame/support/src/storage/types/map.rs
+++ b/frame/support/src/storage/types/map.rs
@@ -484,7 +484,7 @@ mod test {
 	use crate::{
 		hash::*,
 		metadata_ir::{StorageEntryModifierIR, StorageEntryTypeIR, StorageHasherIR},
-		storage::types::ValueQuery,
+		storage::{types::ValueQuery, IterableStorageMap},
 	};
 	use sp_io::{hashing::twox_128, TestExternalities};
 
@@ -700,6 +700,15 @@ mod test {
 			A::translate::<u8, _>(|k, v| Some((k * v as u16).into()));
 			assert_eq!(A::iter().collect::<Vec<_>>(), vec![(4, 40), (3, 30)]);
 
+			let translate_next = |k: u16, v: u8| Some((v as u16 / k).into());
+			let k = A::translate_next::<u8, _>(None, translate_next);
+			let k = A::translate_next::<u8, _>(k, translate_next);
+			assert_eq!(None, A::translate_next::<u8, _>(k, translate_next));
+			assert_eq!(A::iter().collect::<Vec<_>>(), vec![(4, 10), (3, 10)]);
+
+			let _ = A::translate_next::<u8, _>(None, |_, _| None);
+			assert_eq!(A::iter().collect::<Vec<_>>(), vec![(3, 10)]);
+
 			let mut entries = vec![];
 			A::build_metadata(vec![], &mut entries);
 			AValueQueryWithAnOnEmpty::build_metadata(vec![], &mut entries);