Add total nb to trie migration rpc (#12770)

* Add total nb to trie migration rpc

* fix and format

* Use struct instead of tuple

* fixes

Co-authored-by: parity-processbot <>

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

@@ -44,30 +44,33 @@ 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,
@@ -75,9 +78,10 @@ where
 {
 	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)
 	}
 }