Add total nb to trie migration rpc

frame/state-trie-migration/src/lib.rs

@@ -1606,15 +1606,14 @@ mod test {
 pub(crate) mod remote_tests {
 	use crate::{AutoLimits, MigrationLimits, Pallet as StateTrieMigration, LOG_TARGET};
 	use codec::Encode;
-	use frame_benchmarking::Zero;
 	use frame_support::{
 		traits::{Get, Hooks},
 		weights::Weight,
 	};
 	use frame_system::Pallet as System;
 	use remote_externalities::Mode;
 	use sp_core::H256;
-	use sp_runtime::traits::{Block as BlockT, HashFor, Header as _, One};
+	use sp_runtime::traits::{Block as BlockT, HashFor, Header as _, One, Zero};
 	use thousands::Separable;
 
 	#[allow(dead_code)]
@@ -1663,18 +1662,20 @@ pub(crate) mod remote_tests {
 		// set the version to 1, as if the upgrade happened.
 		ext.state_version = sp_core::storage::StateVersion::V1;
 
-		let (top_left, child_left) =
+		let status =
 			substrate_state_trie_migration_rpc::migration_status(&ext.as_backend()).unwrap();
 		assert!(
-			top_left > 0,
+			status.top_remaining_to_migrate > 0,
 			"no node needs migrating, this probably means that state was initialized with `StateVersion::V1`",
 		);
 
 		log::info!(
 			target: LOG_TARGET,
-			"initial check: top_left: {}, child_left: {}",
-			top_left.separate_with_commas(),
-			child_left.separate_with_commas(),
+			"initial check: top_left: {}, child_left: {}, total_top {}, total_child {}",
+			status.top_remaining_to_migrate.separate_with_commas(),
+			status.child_remaining_to_migrate.separate_with_commas(),
+			status.total_top.separate_with_commas(),
+			status.total_child.separate_with_commas(),
 		);
 
 		loop {
@@ -1722,17 +1723,17 @@ pub(crate) mod remote_tests {
 			)
 		});
 
-		let (top_left, child_left) =
+		let status =
 			substrate_state_trie_migration_rpc::migration_status(&ext.as_backend()).unwrap();
-		assert_eq!(top_left, 0);
-		assert_eq!(child_left, 0);
+		assert_eq!(status.top_remaining_to_migrate, 0);
+		assert_eq!(status.child_remaining_to_migrate, 0);
 	}
 }
 
 #[cfg(all(test, feature = "remote-test"))]
 mod remote_tests_local {
 	use super::{
-		mock::{Call as MockCall, *},
+		mock::{RuntimeCall as MockCall, *},
 		remote_tests::run_with_limits,
 		*,
 	};

utils/frame/rpc/state-trie-migration-rpc/src/lib.rs

@@ -44,40 +44,44 @@ use trie_db::{
 fn count_migrate<'a, H: Hasher>(
 	storage: &'a dyn trie_db::HashDBRef<H, Vec<u8>>,
 	root: &'a H::Out,
-) -> std::result::Result<(u64, TrieDB<'a, 'a, H>), String> {
+) -> std::result::Result<(u64, u64, TrieDB<'a, 'a, H>), String> {
 	let mut nb = 0u64;
+	let mut total_nb = 0u64;
 	let trie = TrieDBBuilder::new(storage, root).build();
 	let iter_node =
 		TrieDBNodeIterator::new(&trie).map_err(|e| format!("TrieDB node iterator error: {}", e))?;
 	for node in iter_node {
 		let node = node.map_err(|e| format!("TrieDB node iterator error: {}", e))?;
 		match node.2.node_plan() {
-			NodePlan::Leaf { value, .. } | NodePlan::NibbledBranch { value: Some(value), .. } =>
+			NodePlan::Leaf { value, .. } | NodePlan::NibbledBranch { value: Some(value), .. } => {
+				total_nb += 1;
 				if let ValuePlan::Inline(range) = value {
 					if (range.end - range.start) as u32 >=
 						sp_core::storage::TRIE_VALUE_NODE_THRESHOLD
 					{
 						nb += 1;
 					}
-				},
+				}
+			},
 			_ => (),
 		}
 	}
-	Ok((nb, trie))
+	Ok((nb, total_nb, trie))
 }
 
 /// Check trie migration status.
-pub fn migration_status<H, B>(backend: &B) -> std::result::Result<(u64, u64), String>
+pub fn migration_status<H, B>(backend: &B) -> std::result::Result<MigrationStatusResult, String>
 where
 	H: Hasher,
 	H::Out: codec::Codec,
 	B: AsTrieBackend<H>,
 {
 	let trie_backend = backend.as_trie_backend();
 	let essence = trie_backend.essence();
-	let (nb_to_migrate, trie) = count_migrate(essence, essence.root())?;
+	let (top_remaining_to_migrate, total_top, trie) = count_migrate(essence, essence.root())?;
 
-	let mut nb_to_migrate_child = 0;
+	let mut child_remaining_to_migrate = 0;
+	let mut total_child = 0;
 	let mut child_roots: Vec<(ChildInfo, Vec<u8>)> = Vec::new();
 	// get all child trie roots
 	for key_value in trie.iter().map_err(|e| format!("TrieDB node iterator error: {}", e))? {
@@ -94,18 +98,32 @@ where
 		let storage = KeySpacedDB::new(essence, child_info.keyspace());
 
 		child_root.as_mut()[..].copy_from_slice(&root[..]);
-		nb_to_migrate_child += count_migrate(&storage, &child_root)?.0;
+		let (nb, total_top, _) = count_migrate(&storage, &child_root)?;
+		child_remaining_to_migrate += nb;
+		total_child += total_top;
 	}
 
-	Ok((nb_to_migrate, nb_to_migrate_child))
+	Ok(MigrationStatusResult {
+		top_remaining_to_migrate,
+		child_remaining_to_migrate,
+		total_top,
+		total_child,
+	})
 }
 
+/// Current state migration status.
 #[derive(Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
 #[serde(deny_unknown_fields)]
 pub struct MigrationStatusResult {
-	top_remaining_to_migrate: u64,
-	child_remaining_to_migrate: u64,
+	/// Number of top items that should migrate.
+	pub top_remaining_to_migrate: u64,
+	/// Number of child items that should migrate.
+	pub child_remaining_to_migrate: u64,
+	/// Number of top items that we will iterate on.
+	pub total_top: u64,
+	/// Number of child items that we will iterate on.
+	pub total_child: u64,
 }
 
 /// Migration RPC methods.
@@ -146,12 +164,7 @@ where
 
 		let hash = at.unwrap_or_else(|| self.client.info().best_hash);
 		let state = self.backend.state_at(hash).map_err(error_into_rpc_err)?;
-		let (top, child) = migration_status(&state).map_err(error_into_rpc_err)?;
-
-		Ok(MigrationStatusResult {
-			top_remaining_to_migrate: top,
-			child_remaining_to_migrate: child,
-		})
+		migration_status(&state).map_err(error_into_rpc_err)
 	}
 }