fix(l1,l2): fix occasional error with debug_executionWitness

crates/blockchain/blockchain.rs

@@ -495,21 +495,6 @@ impl Blockchain {
             ))?
             .header;
 
-        // Later on, we need to access block hashes by number. To avoid needing
-        // to apply fork choice for each block, we cache them here.
-        // The clone is not redundant, the hash function modifies the original block.
-        #[allow(clippy::redundant_clone)]
-        let mut block_hashes_map: BTreeMap<u64, H256> = blocks
-            .iter()
-            .cloned()
-            .map(|block| (block.header.number, block.hash()))
-            .collect();
-
-        block_hashes_map.insert(
-            first_block_header.number.saturating_sub(1),
-            first_block_header.parent_hash,
-        );
-
         // Get state at previous block
         let trie = self
             .storage
@@ -538,7 +523,7 @@ impl Blockchain {
             ChainError::WitnessGeneration("Failed to get root state node".to_string())
         })?;
 
-        let mut block_hashes = HashMap::new();
+        let mut blockhash_opcode_references = HashMap::new();
         let mut codes = Vec::new();
 
         for (i, block) in blocks.iter().enumerate() {
@@ -602,7 +587,7 @@ impl Blockchain {
                 })?
                 .clone();
 
-            block_hashes.extend(logger_block_hashes);
+            blockhash_opcode_references.extend(logger_block_hashes);
 
             // Access all the accounts needed for withdrawals
             if let Some(withdrawals) = block.body.withdrawals.as_ref() {
@@ -696,14 +681,7 @@ impl Blockchain {
 
             let (new_state_trie_witness, updated_trie) = TrieLogger::open_trie(
                 self.storage
-                    .state_trie(
-                        block_hashes_map
-                            .get(&block.header.number)
-                            .ok_or(ChainError::WitnessGeneration(
-                                "Block hash not found for witness generation".to_string(),
-                            ))?
-                            .to_owned(),
-                    )
+                    .state_trie(block.header.hash())
                     .map_err(|_| ChainError::ParentStateNotFound)?
                     .ok_or(ChainError::ParentStateNotFound)?,
             );
@@ -734,39 +712,41 @@ impl Blockchain {
             used_trie_nodes.push(root.encode_to_vec());
         }
 
-        let mut needed_block_numbers = block_hashes.keys().collect::<Vec<_>>();
-
-        needed_block_numbers.sort();
+        // - We now need necessary block headers, these go from the first block referenced (via BLOCKHASH or just the first block to execute) up to the parent of the last block to execute.
+        let mut block_headers_bytes = Vec::new();
 
-        // Last needed block header for the witness is the parent of the last block we need to execute
-        let last_needed_block_number = blocks
+        let first_blockhash_opcode_number = blockhash_opcode_references.keys().min();
+        let first_needed_block_hash = first_blockhash_opcode_number
+            .and_then(|n| {
+                (*n < first_block_header.number.saturating_sub(1))
+                    .then(|| blockhash_opcode_references.get(n))?
+                    .copied()
+            })
+            .unwrap_or(first_block_header.parent_hash);
+
+        // At the beginning this is the header of the last block to execute.
+        let mut current_header = blocks
             .last()
-            .ok_or(ChainError::WitnessGeneration("Empty batch".to_string()))?
+            .ok_or_else(|| ChainError::WitnessGeneration("Empty batch".to_string()))?
             .header
-            .number
-            .saturating_sub(1);
-
-        // The first block number we need is either the parent of the first block number or the earliest block number used by BLOCKHASH
-        let mut first_needed_block_number = first_block_header.number.saturating_sub(1);
+            .clone();
 
-        if let Some(block_number_from_logger) = needed_block_numbers.first()
-            && **block_number_from_logger < first_needed_block_number
-        {
-            first_needed_block_number = **block_number_from_logger;
-        }
+        // Headers from latest - 1 until we reach first block header we need.
+        // We do it this way because we want to fetch headers by hash, not by number
+        while current_header.hash() != first_needed_block_hash {
+            let parent_hash = current_header.parent_hash;
+            let current_number = current_header.number - 1;
 
-        let mut block_headers_bytes = Vec::new();
+            current_header = self
+                .storage
+                .get_block_header_by_hash(parent_hash)?
+                .ok_or_else(|| {
+                    ChainError::WitnessGeneration(format!(
+                        "Failed to get block {current_number} header"
+                    ))
+                })?;
 
-        for block_number in first_needed_block_number..=last_needed_block_number {
-            let hash = block_hashes_map
-                .get(&block_number)
-                .ok_or(ChainError::WitnessGeneration(format!(
-                    "Failed to get block {block_number} hash"
-                )))?;
-            let header = self.storage.get_block_header_by_hash(*hash)?.ok_or(
-                ChainError::WitnessGeneration(format!("Failed to get block {block_number} header")),
-            )?;
-            block_headers_bytes.push(header.encode_to_vec());
+            block_headers_bytes.push(current_header.encode_to_vec());
         }
 
         let chain_config = self.storage.get_chain_config();

crates/common/types/block.rs

@@ -317,7 +317,7 @@ impl RLPDecode for BlockBody {
 }
 
 impl BlockHeader {
-    fn compute_block_hash(&self) -> H256 {
+    pub fn compute_block_hash(&self) -> H256 {
         let mut buf = vec![];
         self.encode(&mut buf);
         keccak(buf)

crates/common/types/block_execution_witness.rs

@@ -1,4 +1,4 @@
-use std::collections::BTreeMap;
+use std::collections::{BTreeMap, BTreeSet};
 use std::fmt;
 use std::str::FromStr;
 
@@ -469,40 +469,33 @@ impl GuestProgramState {
         }
     }
 
-    /// Hashes headers in witness and in blocks only once if they are repeated to avoid double hashing.
+    /// When executing multiple blocks in the L2 it happens that the headers in block_headers correspond to the same block headers that we have in the blocks array. The main goal is to hash these only once and set them in both places.
+    /// We also initialize the remaining block headers hashes. If they are set, we check their validity.
     pub fn initialize_block_header_hashes(
         &self,
         blocks: &[Block],
     ) -> Result<(), GuestProgramStateError> {
-        // First we need to ensure that the block headers are initialized not before the guest program is executed
-        for header in self.block_headers.values() {
-            if header.hash.get().is_some() {
-                return Err(GuestProgramStateError::Custom(format!(
-                    "Block header hash is already set for {}",
-                    header.number
-                )));
+        let mut block_numbers_in_common = BTreeSet::new();
+        for block in blocks {
+            let hash = block.header.compute_block_hash();
+            set_hash_or_validate(&block.header, hash)?;
+
+            let number = block.header.number;
+            if let Some(header) = self.block_headers.get(&number) {
+                block_numbers_in_common.insert(number);
+                set_hash_or_validate(header, hash)?;
             }
         }
 
-        // Now we initialize the block_headers hashes and check the remaining blocks hashes
-        for block in blocks {
-            // Verify each block's header hash is uninitialized
-            if block.header.hash.get().is_some() {
-                return Err(GuestProgramStateError::Custom(format!(
-                    "Block header hash is already set for {}",
-                    block.header.number
-                )));
+        for header in self.block_headers.values() {
+            if block_numbers_in_common.contains(&header.number) {
+                // We have already set this hash in the previous step
+                continue;
             }
-            let header = self
-                .block_headers
-                .get(&block.header.number)
-                .unwrap_or(&block.header);
-
-            let hash = header.hash();
-            // this returns err if it's already set, so we drop the Result as we don't
-            // care if it was already initialized.
-            let _ = block.header.hash.set(hash);
+            let hash = header.compute_block_hash();
+            set_hash_or_validate(header, hash)?;
         }
+
         Ok(())
     }
 }
@@ -581,3 +574,18 @@ pub fn hash_key(key: &H256) -> Vec<u8> {
         .finalize()
         .to_vec()
 }
+
+/// Initializes hash of header or validates the hash is correct in case it's already set
+/// Note that header doesn't need to be mutable because the hash is a OnceCell
+fn set_hash_or_validate(header: &BlockHeader, hash: H256) -> Result<(), GuestProgramStateError> {
+    // If it's already set the .set() method will return the current value
+    if let Err(prev_hash) = header.hash.set(hash)
+        && prev_hash != hash
+    {
+        return Err(GuestProgramStateError::Custom(format!(
+            "Block header hash was previously set for {} with the wrong value. It should be set correctly or left unset.",
+            header.number
+        )));
+    }
+    Ok(())
+}