Refactor trie node handling to use enum for type safety

- Replace three separate fields (trie_node_writer, trie_nodes,
  trie_nodes_written) with TrieNodeImportMode enum that enforces
  mutual exclusivity at compile time
- Change import_state_from_key_values from associated function to
  method taking &self since it always uses self.executor

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <[email protected]>

Author: liuchengxu

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

@@ -139,20 +139,22 @@ impl<B: BlockT> StateSyncMetadata<B> {
 	}
 }
 
+/// Defines how trie nodes are imported during state sync.
+enum TrieNodeImportMode {
+	/// Incremental mode: write nodes directly to storage as they arrive.
+	Incremental { writer: Arc<dyn TrieNodeWriter>, nodes_written: u64 },
+	/// Legacy mode: accumulate nodes in memory for later import.
+	Accumulated { nodes: Vec<(Vec<u8>, Vec<u8>)> },
+}
+
 /// State sync state machine.
 ///
 /// Accumulates partial state data until it is ready to be imported.
 pub struct StateSync<B: BlockT, Client> {
 	metadata: StateSyncMetadata<B>,
 	state: HashMap<Vec<u8>, (Vec<(Vec<u8>, Vec<u8>)>, Vec<Vec<u8>>)>,
-	/// Optional writer for incremental trie node import.
-	/// When set, trie nodes are written directly to storage as each chunk arrives.
-	/// When None, trie nodes are accumulated in memory (legacy behavior).
-	trie_node_writer: Option<Arc<dyn TrieNodeWriter>>,
-	/// Accumulated trie nodes when no writer is provided (legacy fallback).
-	trie_nodes: Vec<(Vec<u8>, Vec<u8>)>,
-	/// Count of trie nodes written incrementally.
-	trie_nodes_written: u64,
+	/// How trie nodes are imported during sync.
+	trie_node_import_mode: TrieNodeImportMode,
 	client: Arc<Client>,
 }
 
@@ -174,6 +176,11 @@ where
 		skip_proof: bool,
 		trie_node_writer: Option<Arc<dyn TrieNodeWriter>>,
 	) -> Self {
+		let trie_node_import_mode = match trie_node_writer {
+			Some(writer) => TrieNodeImportMode::Incremental { writer, nodes_written: 0 },
+			None => TrieNodeImportMode::Accumulated { nodes: Vec::new() },
+		};
+
 		Self {
 			client,
 			metadata: StateSyncMetadata {
@@ -186,30 +193,29 @@ where
 				skip_proof,
 			},
 			state: HashMap::default(),
-			trie_node_writer,
-			trie_nodes: Vec::new(),
-			trie_nodes_written: 0,
+			trie_node_import_mode,
 		}
 	}
 
 	/// Process trie nodes from a chunk - either write incrementally or accumulate.
 	fn process_trie_nodes(&mut self, chunk_nodes: Vec<(Vec<u8>, Vec<u8>)>) {
 		let node_count = chunk_nodes.len() as u64;
 
-		if let Some(ref writer) = self.trie_node_writer {
-			// Incremental mode: write directly to storage
-			if let Err(e) = writer.write_trie_nodes(chunk_nodes) {
-				// Log error but continue - the final state root check will catch issues
-				debug!(
-					target: LOG_TARGET,
-					"Failed to write trie nodes incrementally: {}", e
-				);
-			} else {
-				self.trie_nodes_written += node_count;
-			}
-		} else {
-			// Legacy mode: accumulate in memory
-			self.trie_nodes.extend(chunk_nodes);
+		match &mut self.trie_node_import_mode {
+			TrieNodeImportMode::Incremental { writer, nodes_written } => {
+				if let Err(e) = writer.write_trie_nodes(chunk_nodes) {
+					// Log error but continue - the final state root check will catch issues
+					debug!(
+						target: LOG_TARGET,
+						"Failed to write trie nodes incrementally: {e}"
+					);
+				} else {
+					*nodes_written += node_count;
+				}
+			},
+			TrieNodeImportMode::Accumulated { nodes } => {
+				nodes.extend(chunk_nodes);
+			},
 		}
 	}
 
@@ -356,22 +362,22 @@ where
 			let target_hash = self.metadata.target_hash();
 
 			// Determine trie_nodes based on whether incremental writing was used
-			let trie_nodes = if self.trie_node_writer.is_some() {
-				// Incremental mode: nodes already written to storage
-				info!(
-					target: LOG_TARGET,
-					"State sync complete: {} trie nodes written incrementally",
-					self.trie_nodes_written
-				);
-				None
-			} else {
-				// Legacy mode: return accumulated nodes
-				info!(
-					target: LOG_TARGET,
-					"State sync complete: {} trie nodes accumulated",
-					self.trie_nodes.len()
-				);
-				Some(std::mem::take(&mut self.trie_nodes))
+			let trie_nodes = match &mut self.trie_node_import_mode {
+				TrieNodeImportMode::Incremental { nodes_written, .. } => {
+					info!(
+						target: LOG_TARGET,
+						"State sync complete: {nodes_written} trie nodes written incrementally"
+					);
+					None
+				},
+				TrieNodeImportMode::Accumulated { ref mut nodes } => {
+					info!(
+						target: LOG_TARGET,
+						"State sync complete: {} trie nodes accumulated",
+						nodes.len()
+					);
+					Some(std::mem::take(nodes))
+				},
 			};
 
 			ImportResult::Import(

substrate/client/service/src/client/client.rs

@@ -576,9 +576,9 @@ where
 
 		// the block is lower than our last finalized block so it must revert
 		// finality, refusing import.
-		if status == blockchain::BlockStatus::Unknown &&
-			*import_headers.post().number() <= info.finalized_number &&
-			!gap_block
+		if status == blockchain::BlockStatus::Unknown
+			&& *import_headers.post().number() <= info.finalized_number
+			&& !gap_block
 		{
 			return Err(sp_blockchain::Error::NotInFinalizedChain);
 		}
@@ -587,8 +587,9 @@ where
 		// but the general goal is to only make notifications when we are already fully synced
 		// and get a new chain head.
 		let make_notifications = match origin {
-			BlockOrigin::NetworkBroadcast | BlockOrigin::Own | BlockOrigin::ConsensusBroadcast =>
-				true,
+			BlockOrigin::NetworkBroadcast | BlockOrigin::Own | BlockOrigin::ConsensusBroadcast => {
+				true
+			},
 			BlockOrigin::Genesis | BlockOrigin::NetworkInitialSync | BlockOrigin::File => false,
 		};
 
@@ -627,19 +628,17 @@ where
 								None
 							} else {
 								// Fall back to key-value import if no trie nodes
-								Self::import_state_from_key_values(
+								self.import_state_from_key_values(
 									changes.state,
 									&mut operation.op,
-									&self.executor,
 									import_headers.post().state_root(),
 								)?
 							}
 						} else {
 							// Fall back to key-value import
-							Self::import_state_from_key_values(
+							self.import_state_from_key_values(
 								changes.state,
 								&mut operation.op,
-								&self.executor,
 								import_headers.post().state_root(),
 							)?
 						}
@@ -663,11 +662,12 @@ where
 			)?;
 		}
 
-		let is_new_best = !gap_block &&
-			(finalized ||
-				match fork_choice {
-					ForkChoiceStrategy::LongestChain =>
-						import_headers.post().number() > &info.best_number,
+		let is_new_best = !gap_block
+			&& (finalized
+				|| match fork_choice {
+					ForkChoiceStrategy::LongestChain => {
+						import_headers.post().number() > &info.best_number
+					},
 					ForkChoiceStrategy::Custom(v) => v,
 				});
 
@@ -786,18 +786,21 @@ where
 		let state_action = std::mem::replace(&mut import_block.state_action, StateAction::Skip);
 		let (enact_state, storage_changes) = match (self.block_status(*parent_hash)?, state_action)
 		{
-			(BlockStatus::KnownBad, _) =>
-				return Ok(PrepareStorageChangesResult::Discard(ImportResult::KnownBad)),
+			(BlockStatus::KnownBad, _) => {
+				return Ok(PrepareStorageChangesResult::Discard(ImportResult::KnownBad))
+			},
 			(
 				BlockStatus::InChainPruned,
 				StateAction::ApplyChanges(sc_consensus::StorageChanges::Changes(_)),
 			) => return Ok(PrepareStorageChangesResult::Discard(ImportResult::MissingState)),
 			(_, StateAction::ApplyChanges(changes)) => (true, Some(changes)),
-			(BlockStatus::Unknown, _) =>
-				return Ok(PrepareStorageChangesResult::Discard(ImportResult::UnknownParent)),
+			(BlockStatus::Unknown, _) => {
+				return Ok(PrepareStorageChangesResult::Discard(ImportResult::UnknownParent))
+			},
 			(_, StateAction::Skip) => (false, None),
-			(BlockStatus::InChainPruned, StateAction::Execute) =>
-				return Ok(PrepareStorageChangesResult::Discard(ImportResult::MissingState)),
+			(BlockStatus::InChainPruned, StateAction::Execute) => {
+				return Ok(PrepareStorageChangesResult::Discard(ImportResult::MissingState))
+			},
 			(BlockStatus::InChainPruned, StateAction::ExecuteIfPossible) => (false, None),
 			(_, StateAction::Execute) => (true, None),
 			(_, StateAction::ExecuteIfPossible) => (true, None),
@@ -1096,12 +1099,13 @@ where
 
 		let hash_and_number = self.backend.blockchain().number(hash)?.map(|n| (hash, n));
 		match hash_and_number {
-			Some((hash, number)) =>
+			Some((hash, number)) => {
 				if self.backend.have_state_at(hash, number) {
 					Ok(BlockStatus::InChainWithState)
 				} else {
 					Ok(BlockStatus::InChainPruned)
-				},
+				}
+			},
 			None => Ok(BlockStatus::Unknown),
 		}
 	}
@@ -1167,9 +1171,9 @@ where
 
 	/// Helper to import state from key-values (original reset_storage path).
 	fn import_state_from_key_values(
+		&self,
 		state: KeyValueStates,
 		op: &mut B::BlockImportOperation,
-		executor: &E,
 		expected_state_root: &Block::Hash,
 	) -> sp_blockchain::Result<Option<(StorageCollection, ChildStorageCollection)>> {
 		let mut storage = sp_storage::Storage::default();
@@ -1183,8 +1187,9 @@ where
 					let storage_key = PrefixedStorageKey::new_ref(&parent_storage);
 					let storage_key = match ChildType::from_prefixed_key(storage_key) {
 						Some((ChildType::ParentKeyId, storage_key)) => storage_key,
-						None =>
-							return Err(Error::Backend("Invalid child storage key.".to_string())),
+						None => {
+							return Err(Error::Backend("Invalid child storage key.".to_string()))
+						},
 					};
 					let entry = storage
 						.children_default
@@ -1202,7 +1207,7 @@ where
 
 		// This is used by fast sync for runtime version to be resolvable from changes.
 		let state_version =
-			resolve_state_version_from_wasm::<_, HashingFor<Block>>(&storage, executor)?;
+			resolve_state_version_from_wasm::<_, HashingFor<Block>>(&storage, &self.executor)?;
 		let state_root = op.reset_storage(storage, state_version)?;
 		if state_root != *expected_state_root {
 			// State root mismatch when importing state. This should not happen in
@@ -1292,8 +1297,9 @@ where
 		let child_info = |storage_key: &Vec<u8>| -> sp_blockchain::Result<ChildInfo> {
 			let storage_key = PrefixedStorageKey::new_ref(storage_key);
 			match ChildType::from_prefixed_key(storage_key) {
-				Some((ChildType::ParentKeyId, storage_key)) =>
-					Ok(ChildInfo::new_default(storage_key)),
+				Some((ChildType::ParentKeyId, storage_key)) => {
+					Ok(ChildInfo::new_default(storage_key))
+				},
 				None => Err(Error::Backend("Invalid child storage key.".to_string())),
 			}
 		};
@@ -1353,9 +1359,9 @@ where
 				}
 				total_size += size;
 
-				if current_child.is_none() &&
-					sp_core::storage::well_known_keys::is_child_storage_key(next_key.as_slice()) &&
-					!child_roots.contains(value.as_slice())
+				if current_child.is_none()
+					&& sp_core::storage::well_known_keys::is_child_storage_key(next_key.as_slice())
+					&& !child_roots.contains(value.as_slice())
 				{
 					child_roots.insert(value.clone());
 					switch_child_key = Some((next_key.clone(), value.clone()));
@@ -1860,10 +1866,12 @@ where
 			.block_status(hash)
 			.map_err(|e| ConsensusError::ClientImport(e.to_string()))?
 		{
-			BlockStatus::InChainWithState | BlockStatus::Queued =>
-				return Ok(ImportResult::AlreadyInChain),
-			BlockStatus::InChainPruned if !import_existing =>
-				return Ok(ImportResult::AlreadyInChain),
+			BlockStatus::InChainWithState | BlockStatus::Queued => {
+				return Ok(ImportResult::AlreadyInChain)
+			},
+			BlockStatus::InChainPruned if !import_existing => {
+				return Ok(ImportResult::AlreadyInChain)
+			},
 			BlockStatus::InChainPruned => {},
 			BlockStatus::Unknown => {},
 			BlockStatus::KnownBad => return Ok(ImportResult::KnownBad),
@@ -2029,8 +2037,9 @@ where
 
 	fn block(&self, hash: Block::Hash) -> sp_blockchain::Result<Option<SignedBlock<Block>>> {
 		Ok(match (self.header(hash)?, self.body(hash)?, self.justifications(hash)?) {
-			(Some(header), Some(extrinsics), justifications) =>
-				Some(SignedBlock { block: Block::new(header, extrinsics), justifications }),
+			(Some(header), Some(extrinsics), justifications) => {
+				Some(SignedBlock { block: Block::new(header, extrinsics), justifications })
+			},
 			_ => None,
 		})
 	}