perf(core): improve u256 handling, improving PUSH and other opcodes

CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## Perf
 
+### 2025-06-27
+
+- Improve U256 decoding and PUSHX [#3332](https://github.com/lambdaclass/ethrex/pull/3332)
+
 ### 2025-06-26
 
 - Refactor jump opcodes to use a blacklist on invalid targets.

crates/common/common.rs

@@ -6,3 +6,4 @@ pub use bytes::Bytes;
 pub mod base64;
 pub use ethrex_trie::{TrieLogger, TrieWitness};
 pub mod tracing;
+pub mod utils;

crates/common/utils.rs

@@ -0,0 +1,37 @@
+use ethereum_types::U256;
+
+/// Converts a big endian slice to a u256, faster than `u256::from_big_endian`.
+pub fn u256_from_big_endian(slice: &[u8]) -> U256 {
+    let mut padded = [0u8; 32];
+    padded[32 - slice.len()..32].copy_from_slice(slice);
+
+    let mut ret = [0; 4];
+
+    let mut u64_bytes = [0u8; 8];
+    for i in 0..4 {
+        u64_bytes.copy_from_slice(&padded[8 * i..(8 * i + 8)]);
+        ret[4 - i - 1] = u64::from_be_bytes(u64_bytes);
+    }
+
+    U256(ret)
+}
+
+/// Converts a constant big endian slice to a u256, faster than `u256::from_big_endian` and `u256_from_big_endian`.
+///
+/// Note: N should not exceed 32.
+pub fn u256_from_big_endian_const<const N: usize>(slice: [u8; N]) -> U256 {
+    const { assert!(N <= 32, "N must be less or equal to 32") };
+
+    let mut padded = [0u8; 32];
+    padded[32 - N..32].copy_from_slice(&slice);
+
+    let mut ret = [0u64; 4];
+
+    let mut u64_bytes = [0u8; 8];
+    for i in 0..4 {
+        u64_bytes.copy_from_slice(&padded[8 * i..(8 * i + 8)]);
+        ret[4 - i - 1] = u64::from_be_bytes(u64_bytes);
+    }
+
+    U256(ret)
+}

crates/vm/levm/src/execution_handlers.rs

@@ -80,14 +80,38 @@ impl<'a> VM<'a> {
             Opcode::BLOBHASH => self.op_blobhash(),
             Opcode::BLOBBASEFEE => self.op_blobbasefee(),
             Opcode::PUSH0 => self.op_push0(),
-            Opcode::PUSH1 => self.op_push1(),
-            Opcode::PUSH2 => self.op_push2(),
-            // PUSHn
-            op if (Opcode::PUSH3..=Opcode::PUSH32).contains(&op) => {
-                // The following conversions and operation cannot fail due to known operand ranges.
-                #[expect(clippy::arithmetic_side_effects, clippy::as_conversions)]
-                self.op_push(op as usize - Opcode::PUSH0 as usize)
-            }
+            Opcode::PUSH1 => self.op_push::<1>(),
+            Opcode::PUSH2 => self.op_push::<2>(),
+            Opcode::PUSH3 => self.op_push::<3>(),
+            Opcode::PUSH4 => self.op_push::<4>(),
+            Opcode::PUSH5 => self.op_push::<5>(),
+            Opcode::PUSH6 => self.op_push::<6>(),
+            Opcode::PUSH7 => self.op_push::<7>(),
+            Opcode::PUSH8 => self.op_push::<8>(),
+            Opcode::PUSH9 => self.op_push::<9>(),
+            Opcode::PUSH10 => self.op_push::<10>(),
+            Opcode::PUSH11 => self.op_push::<11>(),
+            Opcode::PUSH12 => self.op_push::<12>(),
+            Opcode::PUSH13 => self.op_push::<13>(),
+            Opcode::PUSH14 => self.op_push::<14>(),
+            Opcode::PUSH15 => self.op_push::<15>(),
+            Opcode::PUSH16 => self.op_push::<16>(),
+            Opcode::PUSH17 => self.op_push::<17>(),
+            Opcode::PUSH18 => self.op_push::<18>(),
+            Opcode::PUSH19 => self.op_push::<19>(),
+            Opcode::PUSH20 => self.op_push::<20>(),
+            Opcode::PUSH21 => self.op_push::<21>(),
+            Opcode::PUSH22 => self.op_push::<22>(),
+            Opcode::PUSH23 => self.op_push::<23>(),
+            Opcode::PUSH24 => self.op_push::<24>(),
+            Opcode::PUSH25 => self.op_push::<25>(),
+            Opcode::PUSH26 => self.op_push::<26>(),
+            Opcode::PUSH27 => self.op_push::<27>(),
+            Opcode::PUSH28 => self.op_push::<28>(),
+            Opcode::PUSH29 => self.op_push::<29>(),
+            Opcode::PUSH30 => self.op_push::<30>(),
+            Opcode::PUSH31 => self.op_push::<31>(),
+            Opcode::PUSH32 => self.op_push::<32>(),
             Opcode::AND => self.op_and(),
             Opcode::OR => self.op_or(),
             Opcode::XOR => self.op_xor(),

crates/vm/levm/src/memory.rs

@@ -4,7 +4,7 @@ use crate::{
 };
 use ExceptionalHalt::OutOfBounds;
 use ExceptionalHalt::OutOfGas;
-use ethrex_common::U256;
+use ethrex_common::{U256, utils::u256_from_big_endian};
 
 /// Memory of the EVM, a volatile byte array.
 pub type Memory = Vec<u8>;
@@ -32,7 +32,7 @@ pub fn try_resize(memory: &mut Memory, unchecked_new_size: usize) -> Result<(),
 }
 
 pub fn load_word(memory: &mut Memory, offset: U256) -> Result<U256, VMError> {
-    load_range(memory, offset, WORD_SIZE_IN_BYTES_USIZE).map(U256::from_big_endian)
+    load_range(memory, offset, WORD_SIZE_IN_BYTES_USIZE).map(u256_from_big_endian)
 }
 
 pub fn load_range(memory: &mut Memory, offset: U256, size: usize) -> Result<&[u8], VMError> {

crates/vm/levm/src/opcode_handlers/block.rs

@@ -5,7 +5,7 @@ use crate::{
     utils::*,
     vm::VM,
 };
-use ethrex_common::{U256, types::Fork};
+use ethrex_common::{U256, types::Fork, utils::u256_from_big_endian_const};
 
 // Block Information (11)
 // Opcodes: BLOCKHASH, COINBASE, TIMESTAMP, NUMBER, PREVRANDAO, GASLIMIT, CHAINID, SELFBALANCE, BASEFEE, BLOBHASH, BLOBBASEFEE
@@ -34,7 +34,7 @@ impl<'a> VM<'a> {
         let block_hash = self.db.store.get_block_hash(block_number)?;
         self.current_call_frame_mut()?
             .stack
-            .push(&[U256::from_big_endian(block_hash.as_bytes())])?;
+            .push(&[u256_from_big_endian_const(block_hash.to_fixed_bytes())])?;
 
         Ok(OpcodeResult::Continue { pc_increment: 1 })
     }
@@ -78,7 +78,8 @@ impl<'a> VM<'a> {
     pub fn op_prevrandao(&mut self) -> Result<OpcodeResult, VMError> {
         // https://eips.ethereum.org/EIPS/eip-4399
         // After Paris the prev randao is the prev_randao (or current_random) field
-        let randao = U256::from_big_endian(self.env.prev_randao.unwrap_or_default().0.as_slice());
+        let randao =
+            u256_from_big_endian_const(self.env.prev_randao.unwrap_or_default().to_fixed_bytes());
 
         let current_call_frame = self.current_call_frame_mut()?;
         current_call_frame.increase_consumed_gas(gas_cost::PREVRANDAO)?;
@@ -160,7 +161,7 @@ impl<'a> VM<'a> {
 
         //This should never fail because we check if the index fits above
         let blob_hash = blob_hashes.get(index).ok_or(InternalError::Slicing)?;
-        let hash = U256::from_big_endian(blob_hash.as_bytes());
+        let hash = u256_from_big_endian_const(blob_hash.to_fixed_bytes());
 
         self.current_call_frame_mut()?.stack.push(&[hash])?;
 

crates/vm/levm/src/opcode_handlers/environment.rs

@@ -5,7 +5,7 @@ use crate::{
     utils::word_to_address,
     vm::VM,
 };
-use ethrex_common::U256;
+use ethrex_common::{U256, utils::u256_from_big_endian_const};
 
 // Environmental Information (16)
 // Opcodes: ADDRESS, BALANCE, ORIGIN, CALLER, CALLVALUE, CALLDATALOAD, CALLDATASIZE, CALLDATACOPY, CODESIZE, CODECOPY, GASPRICE, EXTCODESIZE, EXTCODECOPY, RETURNDATASIZE, RETURNDATACOPY, EXTCODEHASH
@@ -20,7 +20,7 @@ impl<'a> VM<'a> {
 
         current_call_frame
             .stack
-            .push(&[U256::from_big_endian(addr.as_bytes())])?;
+            .push(&[u256_from_big_endian_const(addr.to_fixed_bytes())])?;
 
         Ok(OpcodeResult::Continue { pc_increment: 1 })
     }
@@ -49,7 +49,7 @@ impl<'a> VM<'a> {
 
         current_call_frame
             .stack
-            .push(&[U256::from_big_endian(origin.as_bytes())])?;
+            .push(&[u256_from_big_endian_const(origin.to_fixed_bytes())])?;
 
         Ok(OpcodeResult::Continue { pc_increment: 1 })
     }
@@ -62,7 +62,7 @@ impl<'a> VM<'a> {
         let caller = current_call_frame.msg_sender;
         current_call_frame
             .stack
-            .push(&[U256::from_big_endian(caller.as_bytes())])?;
+            .push(&[u256_from_big_endian_const(caller.to_fixed_bytes())])?;
 
         Ok(OpcodeResult::Continue { pc_increment: 1 })
     }
@@ -111,7 +111,7 @@ impl<'a> VM<'a> {
                 *data_byte = *byte;
             }
         }
-        let result = U256::from_big_endian(&data);
+        let result = u256_from_big_endian_const(data);
 
         current_call_frame.stack.push(&[result])?;
 
@@ -402,7 +402,7 @@ impl<'a> VM<'a> {
             return Ok(OpcodeResult::Continue { pc_increment: 1 });
         }
 
-        let hash = U256::from_big_endian(&account_code_hash);
+        let hash = u256_from_big_endian_const(account_code_hash);
         current_call_frame.stack.push(&[hash])?;
 
         Ok(OpcodeResult::Continue { pc_increment: 1 })

crates/vm/levm/src/opcode_handlers/keccak.rs

@@ -4,7 +4,7 @@ use crate::{
     memory::{self, calculate_memory_size},
     vm::VM,
 };
-use ethrex_common::U256;
+use ethrex_common::utils::u256_from_big_endian;
 use sha3::{Digest, Keccak256};
 
 // KECCAK256 (1)
@@ -34,7 +34,7 @@ impl<'a> VM<'a> {
         )?);
         current_call_frame
             .stack
-            .push(&[U256::from_big_endian(&hasher.finalize())])?;
+            .push(&[u256_from_big_endian(&hasher.finalize())])?;
 
         Ok(OpcodeResult::Continue { pc_increment: 1 })
     }

crates/vm/levm/src/opcode_handlers/push.rs

@@ -5,67 +5,30 @@ use crate::{
     vm::VM,
 };
 use ExceptionalHalt::OutOfBounds;
-use ethrex_common::{U256, types::Fork};
+use ethrex_common::{U256, types::Fork, utils::u256_from_big_endian_const};
 
 // Push Operations
 // Opcodes: PUSH0, PUSH1 ... PUSH32
 
 impl<'a> VM<'a> {
-    // Generic PUSH operation
-    pub fn op_push(&mut self, n_bytes: usize) -> Result<OpcodeResult, VMError> {
+    // Generic PUSH operation, optimized at compile time for the given N.
+    pub fn op_push<const N: usize>(&mut self) -> Result<OpcodeResult, VMError> {
         let current_call_frame = self.current_call_frame_mut()?;
         current_call_frame.increase_consumed_gas(gas_cost::PUSHN)?;
 
-        let read_n_bytes = read_bytcode_slice(current_call_frame, n_bytes)?;
+        let read_n_bytes = read_bytcode_slice::<N>(current_call_frame)?;
 
-        current_call_frame
-            .stack
-            .push(&[U256::from_big_endian(read_n_bytes)])?;
+        let value = u256_from_big_endian_const(read_n_bytes);
+        current_call_frame.stack.push(&[value])?;
 
         // The n_bytes that you push to the stack + 1 for the next instruction
-        let increment_pc_by = n_bytes.wrapping_add(1);
+        let increment_pc_by = N.wrapping_add(1);
 
         Ok(OpcodeResult::Continue {
             pc_increment: increment_pc_by,
         })
     }
 
-    /// Specialized PUSH1 operation
-    ///
-    /// We use specialized push1 and push2 implementations because they are way more frequent than the others,
-    /// so their impact on performance is significant.
-    /// These implementations allow using U256::from, which is considerable more performant than U256::from_big_endian)
-    pub fn op_push1(&mut self) -> Result<OpcodeResult, VMError> {
-        let current_call_frame = self.current_call_frame_mut()?;
-        current_call_frame.increase_consumed_gas(gas_cost::PUSHN)?;
-
-        let value = read_bytcode_slice_const::<1>(current_call_frame)?[0];
-
-        current_call_frame.stack.push(&[U256::from(value)])?;
-
-        Ok(OpcodeResult::Continue {
-            // The 1 byte that you push to the stack + 1 for the next instruction
-            pc_increment: 2,
-        })
-    }
-
-    // Specialized PUSH2 operation
-    pub fn op_push2(&mut self) -> Result<OpcodeResult, VMError> {
-        let current_call_frame = self.current_call_frame_mut()?;
-        current_call_frame.increase_consumed_gas(gas_cost::PUSHN)?;
-
-        let read_n_bytes = read_bytcode_slice_const::<2>(current_call_frame)?;
-
-        let value = u16::from_be_bytes(read_n_bytes);
-
-        current_call_frame.stack.push(&[U256::from(value)])?;
-
-        Ok(OpcodeResult::Continue {
-            // The 2 bytes that you push to the stack + 1 for the next instruction
-            pc_increment: 3,
-        })
-    }
-
     // PUSH0
     pub fn op_push0(&mut self) -> Result<OpcodeResult, VMError> {
         // [EIP-3855] - PUSH0 is only available from SHANGHAI
@@ -82,26 +45,8 @@ impl<'a> VM<'a> {
     }
 }
 
-fn read_bytcode_slice(current_call_frame: &CallFrame, n_bytes: usize) -> Result<&[u8], VMError> {
-    let current_pc = current_call_frame.pc;
-    let pc_offset = current_pc
-        // Add 1 to the PC because we don't want to include the
-        // Bytecode of the current instruction in the data we're about
-        // to read. We only want to read the data _NEXT_ to that
-        // bytecode
-        .checked_add(1)
-        .ok_or(InternalError::Overflow)?;
-
-    Ok(current_call_frame
-        .bytecode
-        .get(pc_offset..pc_offset.checked_add(n_bytes).ok_or(OutOfBounds)?)
-        .unwrap_or_default())
-}
-
 // Like `read_bytcode_slice` but using a const generic and returning a fixed size array.
-fn read_bytcode_slice_const<const N: usize>(
-    current_call_frame: &CallFrame,
-) -> Result<[u8; N], VMError> {
+fn read_bytcode_slice<const N: usize>(current_call_frame: &CallFrame) -> Result<[u8; N], VMError> {
     let current_pc = current_call_frame.pc;
     let pc_offset = current_pc
         // Add 1 to the PC because we don't want to include the