feat(type2): Implement bytecode padding rule in the kernel

evm_arithmetization/src/cpu/kernel/asm/account_code.asm

@@ -139,3 +139,66 @@ remove_padding_loop:
 remove_padding_after:
     %stack (offset, ctx, retdest) -> (retdest, offset)
     JUMP
+
+// Convenience macro to call poseidon_hash_code_unpadded and return where we left off.
+%macro poseidon_hash_code_unpadded
+    %stack (addr, len) -> (addr, len, %%after)
+    %jump(poseidon_hash_code_unpadded)
+%%after:
+%endmacro
+
+/// Applies the padding rule to the code located at the provided address before hashing it.
+/// Memory cells after the last code byte will be overwritten.
+global poseidon_hash_code_unpadded:
+    // stack: addr, len, retdest
+    DUP2 ISZERO %jumpi(poseidon_empty_code)
+    DUP2 DUP2 ADD
+    // stack: padding_addr, addr, len, retdest
+
+    // write 1 after the last code byte
+    DUP1 PUSH 1 MSTORE_GENERAL
+    // stack: padding_addr, addr, len, retdest
+    %increment
+    // stack: padding_addr, addr, len, retdest
+
+    // Pad with 0s until the length is a multiple of 56
+    PUSH 56
+    DUP4 %increment
+    // stack: curr_len, 56, padding_addr, addr, len, retdest
+    PUSH 56 SUB
+    // stack: 56 - curr_len, 56, padding_addr, addr, len, retdest
+    MOD
+    // stack: padding_len, padding_addr, addr, len, retdest
+    SWAP3 DUP4
+    // stack: padding_len, len, padding_addr, addr, padding_len, retdest
+    ADD
+    // stack: last_byte_offset, padding_addr, addr, padding_len, retdest
+    %stack (last_byte_offset, padding_addr, addr, padding_len)
+        -> (padding_addr, padding_len, after_padding, addr, last_byte_offset)
+    %jump(memset)
+after_padding:
+    // stack: addr, last_byte_offset, retdest
+
+    // Xor the last element with 0x80
+    PUSH 1 DUP3 ADD
+    // stack: total_code_len, addr, last_byte_offset, retdest
+    SWAP2
+    // stack: last_byte_offset, addr, total_code_len, retdest
+    DUP2 ADD
+    // stack: last_byte_addr, addr, total_code_len, retdest
+    DUP1 MLOAD_GENERAL
+    // stack: last_byte, last_byte_addr, addr, total_code_len, retdest
+    PUSH 0x80 ADD
+    // stack: last_byte_updated, last_byte_addr, addr, total_code_len, retdest
+    MSTORE_GENERAL
+    // stack: addr, total_code_len, retdest
+
+    POSEIDON_GENERAL
+    // stack: codehash, retdest
+    SWAP1
+    JUMP
+
+global poseidon_empty_code:
+    // stack: addr, len, retdest
+    %stack (addr, len, retdest) -> (retdest, @EMPTY_STRING_POSEIDON_HASH)
+    JUMP

evm_arithmetization/src/cpu/kernel/asm/core/create.asm

@@ -168,9 +168,8 @@ after_constructor:
     // Store the code hash of the new contract.
     %returndatasize
     PUSH @SEGMENT_RETURNDATA GET_CONTEXT %build_address_no_offset
-    // stack: addr, len
-    PROVER_INPUT(poseidon_code) // TODO: FIX THIS!
-    %stack (codehash, addr, len) -> (codehash)
+    // stack: addr, len, leftover_gas, success, address, kexit_info
+    %poseidon_hash_code_unpadded
     // stack: codehash, leftover_gas, success, address, kexit_info
     %observe_new_contract
     DUP4

evm_arithmetization/src/cpu/kernel/asm/core/process_txn.asm

@@ -203,9 +203,8 @@ global process_contract_creation_txn_after_constructor:
     PUSH @SEGMENT_RETURNDATA
     GET_CONTEXT
     %build_address_no_offset
-    // stack: addr, len
-    PROVER_INPUT(poseidon_code) // TODO: FIX THIS!
-    %stack (codehash, addr, len) -> (codehash)
+    // stack: addr, len, leftover_gas, new_ctx, address, retdest, success
+    %poseidon_hash_code_unpadded
     // stack: codehash, leftover_gas, new_ctx, address, retdest, success
     %observe_new_contract
     DUP4

evm_arithmetization/src/generation/prover_input.rs

@@ -9,7 +9,6 @@ use itertools::Itertools;
 use num_bigint::BigUint;
 use plonky2::hash::hash_types::RichField;
 use serde::{Deserialize, Serialize};
-use smt_trie::code::hash_bytecode_u256;
 
 use crate::cpu::kernel::cancun_constants::KZG_VERSIONED_HASH;
 use crate::cpu::kernel::constants::cancun_constants::{
@@ -64,7 +63,6 @@ impl<F: RichField> GenerationState<F> {
             "num_bits" => self.run_num_bits(),
             "jumpdest_table" => self.run_jumpdest_table(input_fn),
             "access_lists" => self.run_access_lists(input_fn),
-            "poseidon_code" => self.run_poseidon_code(),
             "ger" => self.run_global_exit_roots(),
             "kzg_point_eval" => self.run_kzg_point_eval(),
             "kzg_point_eval_2" => self.run_kzg_point_eval_2(),
@@ -415,21 +413,6 @@ impl<F: RichField> GenerationState<F> {
         Ok((Segment::AccessedStorageKeys as usize).into())
     }
 
-    fn run_poseidon_code(&mut self) -> Result<U256, ProgramError> {
-        let addr = stack_peek(self, 0)?;
-        let len = stack_peek(self, 1)?.as_usize();
-        let addr = MemoryAddress::new_bundle(addr)?;
-        let code = (0..len)
-            .map(|i| {
-                let mut a = addr;
-                a.virt += i;
-                self.memory.get_with_init(a).as_usize() as u8
-            })
-            .collect_vec();
-
-        Ok(hash_bytecode_u256(code))
-    }
-
     /// Returns the first part of the KZG precompile output.
     fn run_kzg_point_eval(&mut self) -> Result<U256, ProgramError> {
         let versioned_hash = stack_peek(self, 0)?;