perf(levm): store valid jump targets with code

CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## Perf
 
+### 2025-10-21
+
+- Instead of lazy computation of blocklist, do greedy computation of allowlist and store the result, fetch it with the DB. [#4961](https://github.com/lambdaclass/ethrex/pull/4961)
+
 ### 2025-10-17
 
 - Replaces incremental iteration with a one-time precompute method that scans the entire bytecode, building a `BitVec<u8, Msb0>` where bits mark valid `JUMPDEST` positions, skipping `PUSH1..PUSH32` data bytes.

crates/blockchain/blockchain.rs

@@ -313,7 +313,7 @@ impl Blockchain {
                     .ok_or(ChainError::WitnessGeneration(
                         "Failed to get account code".to_string(),
                     ))?;
-                codes.push(code.to_vec());
+                codes.push(code.bytecode.to_vec());
             }
 
             // Apply account updates to the trie recording all the necessary nodes to do so

crates/blockchain/vm.rs

@@ -1,8 +1,7 @@
-use bytes::Bytes;
 use ethrex_common::{
     Address, H256, U256,
     constants::EMPTY_KECCACK_HASH,
-    types::{AccountState, BlockHash, BlockNumber, ChainConfig},
+    types::{AccountState, BlockHash, BlockNumber, ChainConfig, Code},
 };
 use ethrex_storage::Store;
 use ethrex_vm::{EvmError, VmDatabase};
@@ -104,9 +103,9 @@ impl VmDatabase for StoreVmDatabase {
     }
 
     #[instrument(level = "trace", name = "Account code read", skip_all)]
-    fn get_account_code(&self, code_hash: H256) -> Result<Bytes, EvmError> {
+    fn get_account_code(&self, code_hash: H256) -> Result<Code, EvmError> {
         if code_hash == *EMPTY_KECCACK_HASH {
-            return Ok(Bytes::new());
+            return Ok(Code::default());
         }
         match self.store.get_account_code(code_hash) {
             Ok(Some(code)) => Ok(code),

crates/common/constants.rs

@@ -35,14 +35,8 @@ pub static EMPTY_KECCACK_HASH: LazyLock<H256> = LazyLock::new(|| {
     )
 });
 
-pub static EMPTY_TRIE_HASH: LazyLock<H256> = LazyLock::new(|| {
-    H256::from_slice(
-        Keccak256::new()
-            .chain_update([RLP_NULL])
-            .finalize()
-            .as_slice(),
-    )
-});
+pub static EMPTY_TRIE_HASH: LazyLock<H256> =
+    LazyLock::new(|| H256::from_slice(&Keccak256::digest([RLP_NULL])));
 
 // Request related
 pub static DEPOSIT_TOPIC: LazyLock<H256> = LazyLock::new(|| {

crates/common/trie/node_hash.rs

@@ -27,7 +27,7 @@ impl NodeHash {
     pub fn from_encoded_raw(encoded: &[u8]) -> NodeHash {
         if encoded.len() >= 32 {
             let hash = Keccak256::new_with_prefix(encoded).finalize();
-            NodeHash::Hashed(H256::from_slice(hash.as_slice()))
+            NodeHash::Hashed(H256::from_slice(&hash))
         } else {
             NodeHash::from_slice(encoded)
         }
@@ -51,12 +51,7 @@ impl NodeHash {
     /// NOTE: This will hash smaller nodes, only use to get the final root hash, not for intermediate node hashes
     pub fn finalize(self) -> H256 {
         match self {
-            NodeHash::Inline(_) => H256::from_slice(
-                Keccak256::new()
-                    .chain_update(self.as_ref())
-                    .finalize()
-                    .as_slice(),
-            ),
+            NodeHash::Inline(_) => H256::from_slice(&Keccak256::digest(self.as_ref())),
             NodeHash::Hashed(x) => x,
         }
     }

crates/common/trie/trie.rs

@@ -34,10 +34,7 @@ use lazy_static::lazy_static;
 lazy_static! {
     // Hash value for an empty trie, equal to keccak(RLP_NULL)
     pub static ref EMPTY_TRIE_HASH: H256 = H256::from_slice(
-        Keccak256::new()
-            .chain_update([RLP_NULL])
-            .finalize()
-            .as_slice(),
+        &Keccak256::digest([RLP_NULL]),
     );
 }
 

crates/common/types/account.rs

@@ -19,11 +19,58 @@ use crate::{
     utils::keccak,
 };
 
-#[allow(unused)]
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
+pub struct Code {
+    pub hash: H256,
+    pub bytecode: Bytes,
+    // TODO: Consider using Arc<u16> (needs to enable serde rc feature)
+    pub jump_targets: Vec<u16>,
+}
+
+impl Code {
+    // TODO: also add `from_hashed_bytecode` to optimize the download pipeline,
+    // where hash is already known and checked.
+    pub fn from_bytecode(code: Bytes) -> Self {
+        let jump_targets = Self::compute_jump_targets(&code);
+        Self {
+            hash: keccak(code.as_ref()),
+            bytecode: code,
+            jump_targets,
+        }
+    }
+
+    fn compute_jump_targets(code: &[u8]) -> Vec<u16> {
+        debug_assert!(code.len() <= u16::MAX as usize);
+        let mut targets = Vec::new();
+        let mut i = 0;
+        while i < code.len() {
+            match code[i] {
+                // OP_JUMPDEST
+                0x5B => {
+                    targets.push(i as u16);
+                }
+                // OP_PUSH1..32
+                c @ 0x60..0x80 => {
+                    i += (c - 0x5F) as usize;
+                }
+                _ => (),
+            }
+            i += 1;
+        }
+        targets
+    }
+}
+
+impl AsRef<Bytes> for Code {
+    fn as_ref(&self) -> &Bytes {
+        &self.bytecode
+    }
+}
+
 #[derive(Clone, Default, Debug, PartialEq, Eq, Serialize, Deserialize)]
 pub struct Account {
     pub info: AccountInfo,
-    pub code: Bytes,
+    pub code: Code,
     pub storage: BTreeMap<H256, U256>,
 }
 
@@ -63,6 +110,16 @@ impl Default for AccountState {
     }
 }
 
+impl Default for Code {
+    fn default() -> Self {
+        Self {
+            bytecode: Bytes::new(),
+            hash: *EMPTY_KECCACK_HASH,
+            jump_targets: Vec::new(),
+        }
+    }
+}
+
 impl From<GenesisAccount> for Account {
     fn from(genesis: GenesisAccount) -> Self {
         Self {
@@ -71,7 +128,7 @@ impl From<GenesisAccount> for Account {
                 balance: genesis.balance,
                 nonce: genesis.nonce,
             },
-            code: genesis.code,
+            code: Code::from_bytecode(genesis.code),
             storage: genesis
                 .storage
                 .iter()
@@ -160,11 +217,11 @@ impl From<&GenesisAccount> for AccountState {
 }
 
 impl Account {
-    pub fn new(balance: U256, code: Bytes, nonce: u64, storage: BTreeMap<H256, U256>) -> Self {
+    pub fn new(balance: U256, code: Code, nonce: u64, storage: BTreeMap<H256, U256>) -> Self {
         Self {
             info: AccountInfo {
                 balance,
-                code_hash: keccak(code.as_ref()).0.into(),
+                code_hash: code.hash,
                 nonce,
             },
             code,

crates/common/types/account_update.rs

@@ -1,15 +1,17 @@
 use std::collections::BTreeMap;
 
-use crate::{Address, H256, U256, types::AccountInfo};
-use bytes::Bytes;
+use crate::{
+    Address, H256, U256,
+    types::{AccountInfo, Code},
+};
 use serde::{Deserialize, Serialize};
 
 #[derive(Default, Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
 pub struct AccountUpdate {
     pub address: Address,
     pub removed: bool,
     pub info: Option<AccountInfo>,
-    pub code: Option<Bytes>,
+    pub code: Option<Code>,
     pub added_storage: BTreeMap<H256, U256>,
     /// If account was destroyed and then modified we need this for removing its storage but not the entire account.
     pub removed_storage: bool,

crates/common/types/block_execution_witness.rs

@@ -2,7 +2,7 @@ use std::collections::BTreeMap;
 use std::fmt;
 use std::str::FromStr;
 
-use crate::types::Block;
+use crate::types::{Block, Code};
 use crate::{
     H160,
     constants::EMPTY_KECCACK_HASH,
@@ -34,7 +34,7 @@ pub struct GuestProgramState {
     /// Map of code hashes to their corresponding bytecode.
     /// This is computed during guest program execution inside the zkVM,
     /// before the stateless validation.
-    pub codes_hashed: BTreeMap<H256, Vec<u8>>,
+    pub codes_hashed: BTreeMap<H256, Code>,
     /// Map of block numbers to their corresponding block headers.
     /// The block headers are pushed to the zkVM RLP-encoded, and then
     /// decoded and stored in this map during guest program execution,
@@ -148,10 +148,14 @@ impl TryFrom<ExecutionWitness> for GuestProgramState {
             .collect();
 
         // hash codes
+        // TODO: codes here probably needs to be Vec<Code>, rather than recomputing here. This requires rkyv implementation.
         let codes_hashed = value
             .codes
             .into_iter()
-            .map(|code| (keccak(&code), code))
+            .map(|code| {
+                let code = Code::from_bytecode(code.into());
+                (code.hash, code)
+            })
             .collect();
 
         let mut guest_program_state = GuestProgramState {
@@ -255,7 +259,7 @@ impl GuestProgramState {
                     account_state.code_hash = info.code_hash;
                     // Store updated code in DB
                     if let Some(code) = &update.code {
-                        self.codes_hashed.insert(info.code_hash, code.to_vec());
+                        self.codes_hashed.insert(info.code_hash, code.clone());
                     }
                 }
                 // Store the added storage in the account's storage trie and compute its new root
@@ -446,15 +450,12 @@ impl GuestProgramState {
 
     /// Retrieves the account code for a specific account.
     /// Returns an Err if the code is not found.
-    pub fn get_account_code(
-        &self,
-        code_hash: H256,
-    ) -> Result<bytes::Bytes, GuestProgramStateError> {
+    pub fn get_account_code(&self, code_hash: H256) -> Result<Code, GuestProgramStateError> {
         if code_hash == *EMPTY_KECCACK_HASH {
-            return Ok(Bytes::new());
+            return Ok(Code::default());
         }
         match self.codes_hashed.get(&code_hash) {
-            Some(code) => Ok(Bytes::copy_from_slice(code)),
+            Some(code) => Ok(code.clone()),
             None => {
                 // We do this because what usually happens is that the Witness doesn't have the code we asked for but it is because it isn't relevant for that particular case.
                 // In client implementations there are differences and it's natural for some clients to access more/less information in some edge cases.
@@ -463,7 +464,7 @@ impl GuestProgramState {
                     "Missing bytecode for hash {} in witness. Defaulting to empty code.", // If there's a state root mismatch and this prints we have to see if it's the cause or not.
                     hex::encode(code_hash)
                 );
-                Ok(Bytes::new())
+                Ok(Code::default())
             }
         }
     }

crates/l2/common/src/state_diff.rs

@@ -3,7 +3,7 @@ use std::collections::{BTreeMap, HashMap};
 use bytes::Bytes;
 use ethereum_types::{Address, H256, U256};
 use ethrex_common::types::{
-    AccountInfo, AccountState, AccountUpdate, BlockHeader, PrivilegedL2Transaction, TxKind,
+    AccountInfo, AccountState, AccountUpdate, BlockHeader, Code, PrivilegedL2Transaction, TxKind,
     code_hash,
 };
 use ethrex_rlp::decode::RLPDecode;
@@ -297,7 +297,7 @@ impl StateDiff {
                     address: *address,
                     removed: false,
                     info: account_info,
-                    code: diff.bytecode.clone(),
+                    code: diff.bytecode.clone().map(Code::from_bytecode),
                     added_storage: diff.storage.clone().into_iter().collect(),
                     removed_storage: false,
                 },
@@ -583,7 +583,7 @@ pub fn prepare_state_diff(
                 new_balance: account_update.info.clone().map(|info| info.balance),
                 nonce_diff,
                 storage: account_update.added_storage.clone().into_iter().collect(),
-                bytecode: account_update.code.clone(),
+                bytecode: account_update.code.map(|b| b.bytecode).clone(),
                 bytecode_hash: None,
             },
         );

crates/l2/networking/rpc/l2/transaction.rs

@@ -101,13 +101,14 @@ impl RpcHandler for SponsoredTx {
                 .map_err(RpcErr::from)?
                 .unwrap_or_default();
 
-            let prefix: Vec<u8> = code.iter().take(3).copied().collect();
-            if code.len() != EIP7702_DELEGATED_CODE_LEN || prefix != DELGATION_PREFIX {
+            if code.bytecode.len() != EIP7702_DELEGATED_CODE_LEN
+                || code.bytecode[..3] != DELGATION_PREFIX
+            {
                 return Err(RpcErr::InvalidEthrexL2Message(
                     "Invalid tx trying to call non delegated account".to_string(),
                 ));
             }
-            let address = Address::from_slice(&code[3..]);
+            let address = Address::from_slice(&code.bytecode[3..]);
             if address.is_zero() {
                 return Err(RpcErr::InvalidEthrexL2Message(
                     "Invalid tx trying to call non delegated account".to_string(),

crates/l2/prover/src/guest_program/src/execution.rs

@@ -199,8 +199,14 @@ pub fn stateless_validation_l2(
 
     // Check state diffs are valid
     let blob_versioned_hash = if !validium {
+        use bytes::Bytes;
+        use ethrex_common::types::Code;
+
         let mut guest_program_state = GuestProgramState {
-            codes_hashed,
+            codes_hashed: codes_hashed
+                .into_iter()
+                .map(|(h, c)| (h, Code::from_bytecode(Bytes::from_owner(c))))
+                .collect(),
             parent_block_header,
             first_block_number: initial_db.first_block_number,
             chain_config: initial_db.chain_config,
@@ -279,7 +285,11 @@ fn execute_stateless(
     #[cfg(feature = "l2")]
     let nodes_hashed = guest_program_state.nodes_hashed.clone();
     #[cfg(feature = "l2")]
-    let codes_hashed = guest_program_state.codes_hashed.clone();
+    let codes_hashed = guest_program_state
+        .codes_hashed
+        .iter()
+        .map(|(h, c)| (*h, c.bytecode.to_vec()))
+        .collect();
     #[cfg(feature = "l2")]
     let parent_block_header_clone = guest_program_state.parent_block_header.clone();
     #[cfg(feature = "l2")]

crates/l2/sequencer/block_producer/payload_builder.rs

@@ -326,7 +326,7 @@ fn get_account_diffs_in_tx(
             new_balance,
             nonce_diff,
             storage: BTreeMap::new(), // We add the storage later
-            bytecode,
+            bytecode: bytecode.map(|c| c.bytecode),
             bytecode_hash: None,
         };
 

crates/networking/p2p/snap.rs

@@ -112,7 +112,7 @@ pub fn process_byte_codes_request(
     let mut codes = vec![];
     let mut bytes_used = 0;
     for code_hash in request.hashes {
-        if let Some(code) = store.get_account_code(code_hash)? {
+        if let Some(code) = store.get_account_code(code_hash)?.map(|c| c.bytecode) {
             bytes_used += code.len() as u64;
             codes.push(code);
         }