fix(engine): Try to always compute storage root in V2 proofs when account proof is present, fallback if not (#21579)

crates/trie/parallel/src/proof_task.rs

@@ -1810,11 +1810,24 @@ fn dispatch_storage_proofs(
 fn dispatch_v2_storage_proofs(
     storage_work_tx: &CrossbeamSender<StorageWorkerJob>,
     account_targets: &Vec<proof_v2::Target>,
-    storage_targets: B256Map<Vec<proof_v2::Target>>,
+    mut storage_targets: B256Map<Vec<proof_v2::Target>>,
 ) -> Result<B256Map<CrossbeamReceiver<StorageProofResultMessage>>, ParallelStateRootError> {
     let mut storage_proof_receivers =
         B256Map::with_capacity_and_hasher(account_targets.len(), Default::default());
 
+    // Collect hashed addresses from account targets that need their storage roots computed
+    let account_target_addresses: B256Set = account_targets.iter().map(|t| t.key()).collect();
+
+    // For storage targets with associated account proofs, ensure the first target has
+    // min_len(0) so the root node is returned for storage root computation
+    for (hashed_address, targets) in &mut storage_targets {
+        if account_target_addresses.contains(hashed_address) &&
+            let Some(first) = targets.first_mut()
+        {
+            *first = first.with_min_len(0);
+        }
+    }
+
     // Dispatch all proofs for targeted storage slots
     for (hashed_address, targets) in storage_targets {
         // Create channel for receiving StorageProofResultMessage

crates/trie/parallel/src/proof_task_metrics.rs

@@ -26,6 +26,8 @@ pub struct ProofTaskTrieMetrics {
     deferred_encoder_from_cache: Histogram,
     /// Histogram for `Sync` deferred encoder variant count.
     deferred_encoder_sync: Histogram,
+    /// Histogram for dispatched storage proofs that fell back to sync due to missing root.
+    deferred_encoder_dispatched_missing_root: Histogram,
 }
 
 impl ProofTaskTrieMetrics {
@@ -54,6 +56,8 @@ impl ProofTaskTrieMetrics {
         self.deferred_encoder_dispatched.record(stats.dispatched_count as f64);
         self.deferred_encoder_from_cache.record(stats.from_cache_count as f64);
         self.deferred_encoder_sync.record(stats.sync_count as f64);
+        self.deferred_encoder_dispatched_missing_root
+            .record(stats.dispatched_missing_root_count as f64);
     }
 }
 

crates/trie/parallel/src/value_encoder.rs

@@ -32,6 +32,9 @@ pub(crate) struct ValueEncoderStats {
     pub(crate) from_cache_count: u64,
     /// Number of times the `Sync` variant was used (synchronous computation).
     pub(crate) sync_count: u64,
+    /// Number of times a dispatched storage proof had no root node and fell back to sync
+    /// computation.
+    pub(crate) dispatched_missing_root_count: u64,
 }
 
 impl ValueEncoderStats {
@@ -41,6 +44,7 @@ impl ValueEncoderStats {
         self.dispatched_count += other.dispatched_count;
         self.from_cache_count += other.from_cache_count;
         self.sync_count += other.sync_count;
+        self.dispatched_missing_root_count += other.dispatched_missing_root_count;
     }
 }
 
@@ -55,6 +59,11 @@ pub(crate) enum AsyncAccountDeferredValueEncoder<TC, HC> {
         storage_proof_results: Rc<RefCell<B256Map<Vec<ProofTrieNode>>>>,
         /// Shared stats for tracking wait time and counts.
         stats: Rc<RefCell<ValueEncoderStats>>,
+        /// Shared storage proof calculator for synchronous fallback when dispatched proof has no
+        /// root.
+        storage_calculator: Rc<RefCell<StorageProofCalculator<TC, HC>>>,
+        /// Cache to store computed storage roots for future reuse.
+        cached_storage_roots: Arc<DashMap<B256, B256>>,
     },
     /// The storage root was found in cache.
     FromCache { account: Account, root: B256 },
@@ -82,6 +91,8 @@ where
                 proof_result_rx,
                 storage_proof_results,
                 stats,
+                storage_calculator,
+                cached_storage_roots,
             } => {
                 let wait_start = Instant::now();
                 let result = proof_result_rx?
@@ -94,12 +105,35 @@ where
                     .result?;
                 stats.borrow_mut().storage_wait_time += wait_start.elapsed();
 
-                let StorageProofResult::V2 { root: Some(root), proof } = result else {
-                    panic!("StorageProofResult is not V2 with root: {result:?}")
+                let StorageProofResult::V2 { root, proof } = result else {
+                    panic!("StorageProofResult is not V2: {result:?}")
                 };
 
                 storage_proof_results.borrow_mut().insert(hashed_address, proof);
 
+                let root = match root {
+                    Some(root) => root,
+                    None => {
+                        // In `compute_v2_account_multiproof` we ensure that all dispatched storage
+                        // proofs computations for which there is also an account proof will return
+                        // a root node, but it could happen randomly that an account which is not in
+                        // the account proof targets, but _is_ in storage proof targets, will need
+                        // to be encoded as part of general trie traversal, so we need to handle
+                        // that case here.
+                        stats.borrow_mut().dispatched_missing_root_count += 1;
+
+                        let mut calculator = storage_calculator.borrow_mut();
+                        let proof =
+                            calculator.storage_proof(hashed_address, &mut [B256::ZERO.into()])?;
+                        let storage_root = calculator
+                            .compute_root_hash(&proof)?
+                            .expect("storage_proof with dummy target always returns root");
+
+                        cached_storage_roots.insert(hashed_address, storage_root);
+                        storage_root
+                    }
+                };
+
                 (account, root)
             }
             Self::FromCache { account, root } => (account, root),
@@ -235,6 +269,8 @@ where
                 proof_result_rx: Ok(rx),
                 storage_proof_results: self.storage_proof_results.clone(),
                 stats: self.stats.clone(),
+                storage_calculator: self.storage_calculator.clone(),
+                cached_storage_roots: self.cached_storage_roots.clone(),
             }
         }