Pure state sync refactoring (part-1)

substrate/client/api/src/backend.rs

@@ -37,7 +37,9 @@ pub use sp_trie::MerkleValue;
 
 use crate::{blockchain::Backend as BlockchainBackend, UsageInfo};
 
-pub use sp_state_machine::{Backend as StateBackend, BackendTransaction, KeyValueStates};
+pub use sp_state_machine::{
+	Backend as StateBackend, BackendTransaction, KeyValueStates, KeyValueStorageLevel,
+};
 
 /// Extracts the state backend type for the given backend.
 pub type StateBackendFor<B, Block> = <B as Backend<Block>>::State;

substrate/client/network/sync/src/strategy/state_sync.rs

@@ -19,12 +19,12 @@
 //! State sync support.
 
 use crate::{
-	schema::v1::{StateEntry, StateRequest, StateResponse},
+	schema::v1::{KeyValueStateEntry, StateEntry, StateRequest, StateResponse},
 	LOG_TARGET,
 };
 use codec::{Decode, Encode};
 use log::debug;
-use sc_client_api::{CompactProof, ProofProvider};
+use sc_client_api::{CompactProof, KeyValueStates, KeyValueStorageLevel, ProofProvider};
 use sc_consensus::ImportedState;
 use smallvec::SmallVec;
 use sp_core::storage::well_known_keys;
@@ -132,6 +132,78 @@ where
 			skip_proof,
 		}
 	}
+
+	fn process_state_key_values(
+		&mut self,
+		state_root: Vec<u8>,
+		key_values: impl IntoIterator<Item = (Vec<u8>, Vec<u8>)>,
+	) {
+		let is_top = state_root.is_empty();
+
+		let (child_key_values, top_key_values): (Vec<_>, Vec<_>) =
+			key_values.into_iter().partition(|key_value| {
+				is_top && well_known_keys::is_child_storage_key(key_value.0.as_slice())
+			});
+
+		let entry = self.state.entry(state_root).or_default();
+
+		if entry.0.len() > 0 && entry.1.len() > 1 {
+			// Already imported child_trie with same root.
+			// Warning this will not work with parallel download.
+		} else {
+			self.imported_bytes +=
+				top_key_values.iter().map(|(key, _value)| key.len()).sum::<usize>() as u64;
+
+			entry.0.extend(top_key_values);
+
+			for key_value in child_key_values {
+				self.state.entry(key_value.1).or_default().1.push(key_value.0);
+			}
+		}
+	}
+
+	fn process_state_verified(&mut self, values: KeyValueStates) {
+		for values in values.0 {
+			let KeyValueStorageLevel { state_root, parent_storage_keys: _, key_values } = values;
+			self.process_state_key_values(state_root, key_values);
+		}
+	}
+
+	fn process_state_unverified(&mut self, response: StateResponse) -> bool {
+		let mut complete = true;
+		// if the trie is a child trie and one of its parent trie is empty,
+		// the parent cursor stays valid.
+		// Empty parent trie content only happens when all the response content
+		// is part of a single child trie.
+		if self.last_key.len() == 2 && response.entries[0].entries.is_empty() {
+			// Do not remove the parent trie position.
+			self.last_key.pop();
+		} else {
+			self.last_key.clear();
+		}
+		for state in response.entries {
+			debug!(
+				target: LOG_TARGET,
+				"Importing state from {:?} to {:?}",
+				state.entries.last().map(|e| sp_core::hexdisplay::HexDisplay::from(&e.key)),
+				state.entries.first().map(|e| sp_core::hexdisplay::HexDisplay::from(&e.key)),
+			);
+
+			if !state.complete {
+				if let Some(e) = state.entries.last() {
+					self.last_key.push(e.key.clone());
+				}
+				complete = false;
+			}
+
+			let KeyValueStateEntry { state_root, entries, complete: _ } = state;
+			self.process_state_key_values(
+				state_root,
+				entries.into_iter().map(|StateEntry { key, value }| (key, value)),
+			);
+		}
+		complete
+	}
 }
 
 impl<B, Client> StateSyncProvider<B> for StateSync<B, Client>
@@ -181,94 +253,11 @@ where
 				debug!(target: LOG_TARGET, "Error updating key cursor, depth: {}", completed);
 			};
 
-			for values in values.0 {
-				let key_values = if values.state_root.is_empty() {
-					// Read child trie roots.
-					values
-						.key_values
-						.into_iter()
-						.filter(|key_value| {
-							if well_known_keys::is_child_storage_key(key_value.0.as_slice()) {
-								self.state
-									.entry(key_value.1.clone())
-									.or_default()
-									.1
-									.push(key_value.0.clone());
-								false
-							} else {
-								true
-							}
-						})
-						.collect()
-				} else {
-					values.key_values
-				};
-				let entry = self.state.entry(values.state_root).or_default();
-				if entry.0.len() > 0 && entry.1.len() > 1 {
-					// Already imported child_trie with same root.
-					// Warning this will not work with parallel download.
-				} else if entry.0.is_empty() {
-					for (key, _value) in key_values.iter() {
-						self.imported_bytes += key.len() as u64;
-					}
-
-					entry.0 = key_values;
-				} else {
-					for (key, value) in key_values {
-						self.imported_bytes += key.len() as u64;
-						entry.0.push((key, value))
-					}
-				}
-			}
+			self.process_state_verified(values);
 			self.imported_bytes += proof_size;
 			complete
 		} else {
-			let mut complete = true;
-			// if the trie is a child trie and one of its parent trie is empty,
-			// the parent cursor stays valid.
-			// Empty parent trie content only happens when all the response content
-			// is part of a single child trie.
-			if self.last_key.len() == 2 && response.entries[0].entries.is_empty() {
-				// Do not remove the parent trie position.
-				self.last_key.pop();
-			} else {
-				self.last_key.clear();
-			}
-			for state in response.entries {
-				debug!(
-					target: LOG_TARGET,
-					"Importing state from {:?} to {:?}",
-					state.entries.last().map(|e| sp_core::hexdisplay::HexDisplay::from(&e.key)),
-					state.entries.first().map(|e| sp_core::hexdisplay::HexDisplay::from(&e.key)),
-				);
-
-				if !state.complete {
-					if let Some(e) = state.entries.last() {
-						self.last_key.push(e.key.clone());
-					}
-					complete = false;
-				}
-				let is_top = state.state_root.is_empty();
-				let entry = self.state.entry(state.state_root).or_default();
-				if entry.0.len() > 0 && entry.1.len() > 1 {
-					// Already imported child trie with same root.
-				} else {
-					let mut child_roots = Vec::new();
-					for StateEntry { key, value } in state.entries {
-						// Skip all child key root (will be recalculated on import).
-						if is_top && well_known_keys::is_child_storage_key(key.as_slice()) {
-							child_roots.push((value, key));
-						} else {
-							self.imported_bytes += key.len() as u64;
-							entry.0.push((key, value))
-						}
-					}
-					for (root, storage_key) in child_roots {
-						self.state.entry(root).or_default().1.push(storage_key);
-					}
-				}
-			}
-			complete
+			self.process_state_unverified(response)
 		};
 		if complete {
 			self.complete = true;