Fill in hash_kernel

evm/src/cpu/bootstrap_kernel.rs

@@ -9,44 +9,37 @@ use plonky2::field::packed::PackedField;
 use plonky2::field::types::Field;
 use plonky2::hash::hash_types::RichField;
 use plonky2::plonk::circuit_builder::CircuitBuilder;
-use plonky2_util::ceil_div_usize;
 
 use crate::constraint_consumer::{ConstraintConsumer, RecursiveConstraintConsumer};
 use crate::cpu::columns::{CpuColumnsView, NUM_CPU_COLUMNS};
 use crate::cpu::kernel::aggregator::KERNEL;
 use crate::cpu::kernel::keccak_util::keccakf_u32s;
 use crate::generation::state::GenerationState;
+use crate::keccak_sponge::columns::KECCAK_RATE_U32S;
 use crate::memory::segments::Segment;
 use crate::vars::{StarkEvaluationTargets, StarkEvaluationVars};
 
-/// The Keccak rate (1088 bits), measured in bytes.
-const KECCAK_RATE_BYTES: usize = 1088 / 8;
-
-/// The Keccak rate (1088 bits), measured in u32 limbs.
-const KECCAK_RATE_LIMBS: usize = 1088 / 32;
-
 /// We can't process more than `NUM_CHANNELS` bytes per row, since that's all the memory bandwidth
 /// we have. We also can't process more than 4 bytes (or the number of bytes in a `u32`), since we
 /// want them to fit in a single limb of Keccak input.
 const BYTES_PER_ROW: usize = 4;
 
 pub(crate) fn generate_bootstrap_kernel<F: Field>(state: &mut GenerationState<F>) {
-    let mut code = KERNEL.code.clone();
-
-    // Zero-pad the code such that its size is a multiple of the Keccak rate.
-    let padded_size = ceil_div_usize(code.len(), KECCAK_RATE_BYTES) * KECCAK_RATE_BYTES;
-    code.resize(padded_size, 0);
-
     let mut sponge_state = [0u32; 50];
     let mut sponge_input_pos: usize = 0;
 
     // Iterate through chunks of the code, such that we can write one chunk to memory per row.
-    for chunk in &code.into_iter().enumerate().chunks(BYTES_PER_ROW) {
+    for chunk in &KERNEL
+        .padded_code()
+        .iter()
+        .enumerate()
+        .chunks(BYTES_PER_ROW)
+    {
         state.current_cpu_row.is_bootstrap_kernel = F::ONE;
 
         // Write this chunk to memory, while simultaneously packing its bytes into a u32 word.
         let mut packed_bytes: u32 = 0;
-        for (channel, (addr, byte)) in chunk.enumerate() {
+        for (channel, (addr, &byte)) in chunk.enumerate() {
             state.set_mem_cpu_current(channel, Segment::Code, addr, byte.into());
 
             packed_bytes = (packed_bytes << 8) | byte as u32;
@@ -57,7 +50,7 @@ pub(crate) fn generate_bootstrap_kernel<F: Field>(state: &mut GenerationState<F>
         keccak.input_limbs = sponge_state.map(F::from_canonical_u32);
         state.commit_cpu_row();
 
-        sponge_input_pos = (sponge_input_pos + 1) % KECCAK_RATE_LIMBS;
+        sponge_input_pos = (sponge_input_pos + 1) % KECCAK_RATE_U32S;
         // If we just crossed a multiple of KECCAK_RATE_LIMBS, then we've filled the Keccak input
         // buffer, so it's time to absorb.
         if sponge_input_pos == 0 {

evm/src/cpu/kernel/assembler.rs

@@ -3,6 +3,7 @@ use std::collections::HashMap;
 use ethereum_types::U256;
 use itertools::izip;
 use log::debug;
+use plonky2_util::ceil_div_usize;
 
 use super::ast::PushTarget;
 use crate::cpu::kernel::ast::Item::LocalLabelDeclaration;
@@ -16,6 +17,7 @@ use crate::cpu::kernel::{
     opcodes::{get_opcode, get_push_opcode},
 };
 use crate::generation::prover_input::ProverInputFn;
+use crate::keccak_sponge::columns::KECCAK_RATE_BYTES;
 
 /// The number of bytes to push when pushing an offset within the code (i.e. when assembling jumps).
 /// Ideally we would automatically use the minimal number of bytes required, but that would be
@@ -42,7 +44,8 @@ impl Kernel {
         global_labels: HashMap<String, usize>,
         prover_inputs: HashMap<usize, ProverInputFn>,
     ) -> Self {
-        let code_hash = hash_kernel(&code);
+        let code_hash = hash_kernel(&Self::padded_code_helper(&code));
+
         Self {
             code,
             code_hash,
@@ -50,6 +53,18 @@ impl Kernel {
             prover_inputs,
         }
     }
+
+    /// Zero-pads the code such that its length is a multiple of the Keccak rate.
+    pub(crate) fn padded_code(&self) -> Vec<u8> {
+        Self::padded_code_helper(&self.code)
+    }
+
+    fn padded_code_helper(code: &[u8]) -> Vec<u8> {
+        let padded_len = ceil_div_usize(code.len(), KECCAK_RATE_BYTES) * KECCAK_RATE_BYTES;
+        let mut padded_code = code.to_vec();
+        padded_code.resize(padded_len, 0);
+        padded_code
+    }
 }
 
 #[derive(Eq, PartialEq, Hash, Clone, Debug)]

evm/src/cpu/kernel/keccak_util.rs

@@ -1,13 +1,27 @@
 use tiny_keccak::keccakf;
 
+use crate::keccak_sponge::columns::{KECCAK_RATE_BYTES, KECCAK_RATE_U32S};
+
 /// A Keccak-f based hash.
 ///
 /// This hash does not use standard Keccak padding, since we don't care about extra zeros at the
-/// end of the code.
-pub(crate) fn hash_kernel(_code: &[u8]) -> [u32; 8] {
-    let state = [0u32; 50];
-    // TODO: absorb code
-    state[0..8].try_into().unwrap()
+/// end of the code. It also uses an overwrite-mode sponge, rather than a standard sponge where
+/// inputs are xor'ed in.
+pub(crate) fn hash_kernel(code: &[u8]) -> [u32; 8] {
+    debug_assert_eq!(
+        code.len() % KECCAK_RATE_BYTES,
+        0,
+        "Code should have been padded to a multiple of the Keccak rate."
+    );
+
+    let mut state = [0u32; 50];
+    for chunk in code.chunks(KECCAK_RATE_BYTES) {
+        for i in 0..KECCAK_RATE_U32S {
+            state[i] = u32::from_le_bytes(std::array::from_fn(|j| chunk[i * 4 + j]));
+        }
+        keccakf_u32s(&mut state);
+    }
+    state[..8].try_into().unwrap()
 }
 
 /// Like tiny-keccak's `keccakf`, but deals with `u32` limbs instead of `u64` limbs.