From 0f561aabe9c2b1f5b953cf3ce984e6c321f4af69 Mon Sep 17 00:00:00 2001 From: Robin Salen Date: Thu, 11 Jul 2024 15:00:38 -0400 Subject: [PATCH 1/3] Implement code padding for Poseidon --- .../src/cpu/kernel/asm/account_code.asm | 49 +++++++++++++++++++ .../src/cpu/kernel/asm/core/create.asm | 5 +- .../src/cpu/kernel/asm/core/process_txn.asm | 5 +- .../src/generation/prover_input.rs | 17 ------- 4 files changed, 53 insertions(+), 23 deletions(-) diff --git a/evm_arithmetization/src/cpu/kernel/asm/account_code.asm b/evm_arithmetization/src/cpu/kernel/asm/account_code.asm index 6dac4c32a..af0651d1f 100644 --- a/evm_arithmetization/src/cpu/kernel/asm/account_code.asm +++ b/evm_arithmetization/src/cpu/kernel/asm/account_code.asm @@ -139,3 +139,52 @@ remove_padding_loop: remove_padding_after: %stack (offset, ctx, retdest) -> (retdest, offset) JUMP + +/// 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. +%macro poseidon_hash_code_unpadded + // stack: addr, len + DUP2 DUP2 ADD + // stack: padding_addr, addr, len + + // write 1 after the last code byte + DUP1 PUSH 1 MSTORE_GENERAL + // stack: padding_addr, addr, len + %increment + // stack: padding_addr, addr, len + + // Pad with 0s until the length is a multiple of 56 + PUSH 56 + DUP4 %increment + // stack: curr_len, 56, padding_addr, addr, len + MOD + // stack: remainder, padding_addr, addr, len + PUSH 56 SUB + // stack: padding_len, padding_addr, addr, len + SWAP3 DUP4 + // stack: padding_len, len, padding_addr, addr, padding_len + ADD + // stack: last_byte_offset, padding_addr, addr, padding_len + %stack (last_byte_offset, padding_addr, addr, padding_len) + -> (padding_addr, padding_len, %%after, addr, last_byte_offset) + %jump(memset) +%%after: + // stack: addr, last_byte_offset + + // Xor the last element with 0x80 + PUSH 1 DUP3 ADD + // stack: total_code_len, addr, last_byte_offset + SWAP2 + // stack: last_byte_offset, addr, total_code_len + DUP2 ADD + // stack: last_byte_addr, addr, total_code_len + DUP1 MLOAD_GENERAL + // stack: last_byte, last_byte_addr, addr, total_code_len + PUSH 0x80 XOR + // stack: last_byte_updated, last_byte_addr, addr, total_code_len + MSTORE_GENERAL + // stack: addr, total_code_len + + POSEIDON_GENERAL + // stack: codehash +%endmacro diff --git a/evm_arithmetization/src/cpu/kernel/asm/core/create.asm b/evm_arithmetization/src/cpu/kernel/asm/core/create.asm index 07b2af6d9..99644f2fa 100644 --- a/evm_arithmetization/src/cpu/kernel/asm/core/create.asm +++ b/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 diff --git a/evm_arithmetization/src/cpu/kernel/asm/core/process_txn.asm b/evm_arithmetization/src/cpu/kernel/asm/core/process_txn.asm index d1a898264..242a266f8 100644 --- a/evm_arithmetization/src/cpu/kernel/asm/core/process_txn.asm +++ b/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 diff --git a/evm_arithmetization/src/generation/prover_input.rs b/evm_arithmetization/src/generation/prover_input.rs index 5b664ceab..408222945 100644 --- a/evm_arithmetization/src/generation/prover_input.rs +++ b/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 GenerationState { "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 GenerationState { Ok((Segment::AccessedStorageKeys as usize).into()) } - fn run_poseidon_code(&mut self) -> Result { - 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 { let versioned_hash = stack_peek(self, 0)?; From 767e904302fef4a3eb83adcc6663aa95a8621708 Mon Sep 17 00:00:00 2001 From: Robin Salen Date: Thu, 11 Jul 2024 15:18:43 -0400 Subject: [PATCH 2/3] Check empty code --- .../src/cpu/kernel/asm/account_code.asm | 56 ++++++++++++------- 1 file changed, 35 insertions(+), 21 deletions(-) diff --git a/evm_arithmetization/src/cpu/kernel/asm/account_code.asm b/evm_arithmetization/src/cpu/kernel/asm/account_code.asm index af0651d1f..597d877fb 100644 --- a/evm_arithmetization/src/cpu/kernel/asm/account_code.asm +++ b/evm_arithmetization/src/cpu/kernel/asm/account_code.asm @@ -140,51 +140,65 @@ 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. -%macro poseidon_hash_code_unpadded - // stack: addr, len +global poseidon_hash_code_unpadded: + // stack: addr, len, retdest + DUP2 ISZERO %jumpi(poseidon_empty_code) DUP2 DUP2 ADD - // stack: padding_addr, addr, len + // stack: padding_addr, addr, len, retdest // write 1 after the last code byte DUP1 PUSH 1 MSTORE_GENERAL - // stack: padding_addr, addr, len + // stack: padding_addr, addr, len, retdest %increment - // stack: padding_addr, addr, len + // 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 + // stack: curr_len, 56, padding_addr, addr, len, retdest MOD - // stack: remainder, padding_addr, addr, len + // stack: remainder, padding_addr, addr, len, retdest PUSH 56 SUB - // stack: padding_len, padding_addr, addr, len + // stack: padding_len, padding_addr, addr, len, retdest SWAP3 DUP4 - // stack: padding_len, len, padding_addr, addr, padding_len + // stack: padding_len, len, padding_addr, addr, padding_len, retdest ADD - // stack: last_byte_offset, padding_addr, addr, padding_len + // stack: last_byte_offset, padding_addr, addr, padding_len, retdest %stack (last_byte_offset, padding_addr, addr, padding_len) - -> (padding_addr, padding_len, %%after, addr, last_byte_offset) + -> (padding_addr, padding_len, after_padding, addr, last_byte_offset) %jump(memset) -%%after: - // stack: addr, last_byte_offset +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 + // stack: total_code_len, addr, last_byte_offset, retdest SWAP2 - // stack: last_byte_offset, addr, total_code_len + // stack: last_byte_offset, addr, total_code_len, retdest DUP2 ADD - // stack: last_byte_addr, addr, total_code_len + // stack: last_byte_addr, addr, total_code_len, retdest DUP1 MLOAD_GENERAL - // stack: last_byte, last_byte_addr, addr, total_code_len + // stack: last_byte, last_byte_addr, addr, total_code_len, retdest PUSH 0x80 XOR - // stack: last_byte_updated, last_byte_addr, addr, total_code_len + // stack: last_byte_updated, last_byte_addr, addr, total_code_len, retdest MSTORE_GENERAL - // stack: addr, total_code_len + // stack: addr, total_code_len, retdest POSEIDON_GENERAL - // stack: codehash -%endmacro + // stack: codehash, retdest + SWAP1 + JUMP + +global poseidon_empty_code: + // stack: addr, len, retdest + %stack (addr, len, retdest) -> (retdest, @EMPTY_STRING_POSEIDON_HASH) + JUMP From 5b956929172fbe0d14035bed1c139ef1108ca146 Mon Sep 17 00:00:00 2001 From: Robin Salen Date: Mon, 15 Jul 2024 18:36:32 -0400 Subject: [PATCH 3/3] Apply comments --- evm_arithmetization/src/cpu/kernel/asm/account_code.asm | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/evm_arithmetization/src/cpu/kernel/asm/account_code.asm b/evm_arithmetization/src/cpu/kernel/asm/account_code.asm index 597d877fb..bc8abc5f8 100644 --- a/evm_arithmetization/src/cpu/kernel/asm/account_code.asm +++ b/evm_arithmetization/src/cpu/kernel/asm/account_code.asm @@ -165,9 +165,9 @@ global poseidon_hash_code_unpadded: PUSH 56 DUP4 %increment // stack: curr_len, 56, padding_addr, addr, len, retdest - MOD - // stack: remainder, 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 @@ -188,7 +188,7 @@ after_padding: // 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 XOR + PUSH 0x80 ADD // stack: last_byte_updated, last_byte_addr, addr, total_code_len, retdest MSTORE_GENERAL // stack: addr, total_code_len, retdest