Add total nb to trie migration rpc

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<(u64, u64, u64, u64), 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 (nb_to_migrate, total_nb, trie) = count_migrate(essence, essence.root())?;
 
 	let mut nb_to_migrate_child = 0;
+	let mut total_nb_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,10 +98,12 @@ 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_nb, _) = count_migrate(&storage, &child_root)?;
+		nb_to_migrate_child += nb;
+		total_nb_child += total_nb;
 	}
 
-	Ok((nb_to_migrate, nb_to_migrate_child))
+	Ok((nb_to_migrate, nb_to_migrate_child, total_nb, total_nb_child))
 }
 
 #[derive(Serialize, Deserialize)]
@@ -106,6 +112,8 @@ where
 pub struct MigrationStatusResult {
 	top_remaining_to_migrate: u64,
 	child_remaining_to_migrate: u64,
+	total_top: u64,
+	total_child: u64,
 }
 
 /// Migration RPC methods.
@@ -146,11 +154,14 @@ 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)?;
+		let (top, child, total_top, total_child) =
+			migration_status(&state).map_err(error_into_rpc_err)?;
 
 		Ok(MigrationStatusResult {
 			top_remaining_to_migrate: top,
 			child_remaining_to_migrate: child,
+			total_top,
+			total_child,
 		})
 	}
 }