feat: support Keccak precompile (#210)

* add runtime/emulator/prover/tests support for keccak syscall

* Fix keccak result when length of data is zero

* update test

* fix: syscall precompiles

* fix fmt and clippy

* remove unused binding

* fix keccak guest name

---------

Co-authored-by: eigmax <stephen@ieigen.com>

emulator/Cargo.toml

@@ -14,6 +14,7 @@ lazy_static = "1.4.0"
 elf = { version = "0.7", default-features = false }
 log = { version = "0.4.14", default-features = false }
 itertools = "0.13.0"
+keccak-hash = "0.11.0"
 
 [features]
 test = []

emulator/src/state.rs

@@ -25,6 +25,7 @@ pub const PAGE_CYCLES: u64 = PAGE_LOAD_CYCLES + PAGE_HASH_CYCLES;
 pub const IMAGE_ID_CYCLES: u64 = 3;
 pub const MAX_INSTRUCTION_CYCLES: u64 = PAGE_CYCLES * 6; //TOFIX
 pub const RESERVE_CYCLES: u64 = IMAGE_ID_CYCLES + MAX_INSTRUCTION_CYCLES;
+use keccak_hash::keccak;
 
 // image_id = keccak(page_hash_root || end_pc)
 #[derive(Serialize, Deserialize, Debug, PartialEq, Eq, Clone, Default)]
@@ -529,6 +530,23 @@ impl InstrumentedState {
         log::debug!("syscall {} {} {} {}", syscall_num, a0, a1, a2);
 
         match syscall_num {
+            0x010109 => {
+                assert!((a0 & 3) == 0);
+                assert!((a2 & 3) == 0);
+                let bytes = (0..a1)
+                    .map(|i| self.state.memory.byte(a0 + i))
+                    .collect::<Vec<u8>>();
+                log::debug!("keccak {:X?}", bytes);
+                let result = keccak(&bytes).0;
+                log::debug!("result {:X?}", result);
+                let result: [u32; 8] = core::array::from_fn(|i| {
+                    u32::from_be_bytes(core::array::from_fn(|j| result[i * 4 + j]))
+                });
+                assert!(result.len() == 8);
+                for (i, data) in result.iter().enumerate() {
+                    self.state.memory.set_memory(a2 + ((i << 2) as u32), *data);
+                }
+            }
             0xF0 => {
                 if self.state.input_stream_ptr >= self.state.input_stream.len() {
                     panic!("not enough vecs in hint input stream");
@@ -654,7 +672,9 @@ impl InstrumentedState {
                         v0 = a2;
                     }
                     FD_PUBLIC_VALUES => {
+                        log::debug!("commit {:X?}", slice);
                         self.state.public_values_stream.extend_from_slice(slice);
+
                         v0 = a2;
                     }
                     FD_HINT => {

prover/Cargo.toml

@@ -35,6 +35,7 @@ byteorder = "1.5.0"
 hex = "0.4"
 hashbrown = { version = "0.14.0", default-features = false, features = ["ahash", "serde"] } # NOTE: When upgrading, see `ahash` dependency.
 lazy_static = "1.4.0"
+keccak-hash = "0.10.0"
 
 elf = { version = "0.7", default-features = false }
 sha2 = { version = "0.10.8", default-features = false }
@@ -45,7 +46,6 @@ keccak-hash = "0.10.0"
 plonky2x = { git = "https://github.com/zkMIPS/succinctx.git", package = "plonky2x", branch = "zkm" }
 plonky2x-derive = { git = "https://github.com/zkMIPS/succinctx.git", package = "plonky2x-derive", branch = "zkm" }
 
-
 [features]
 test = []
 

prover/examples/Cargo.toml

@@ -4,6 +4,7 @@ members = [
   "sha2-rust/host",
   "sha2-precompile/host",
   "sha2-go/host",
+  "keccak/host",
   "split-seg",
   "prove-seg"
 ]

prover/examples/keccak/guest/Cargo.toml

@@ -0,0 +1,9 @@
+[workspace]
+[package]
+version = "0.1.0"
+name = "keccak-precompile"
+edition = "2021"
+
+[dependencies]
+#zkm-runtime = { git = "https://github.com/zkMIPS/zkm", package = "zkm-runtime" }
+zkm-runtime = { path = "../../../../runtime/entrypoint" }

prover/examples/keccak/guest/src/main.rs

@@ -0,0 +1,16 @@
+#![no_std]
+#![no_main]
+
+extern crate alloc;
+use alloc::vec::Vec;
+
+zkm_runtime::entrypoint!(main);
+
+pub fn main() {
+    let public_input: Vec<u8> = zkm_runtime::io::read();
+    let input: Vec<u8> = zkm_runtime::io::read();
+
+    let output = zkm_runtime::io::keccak(&input.as_slice()); 
+    assert_eq!(output.to_vec(), public_input);
+    zkm_runtime::io::commit::<[u8; 32]>(&output);
+}

prover/examples/keccak/host/Cargo.toml

@@ -0,0 +1,24 @@
+[package]
+name = "keccak-host"
+version = { workspace = true }
+edition = { workspace = true }
+publish = false
+
+[dependencies]
+zkm-prover = { workspace = true }
+zkm-emulator = { workspace = true }
+zkm-utils = { path = "../../utils" }
+
+
+log = { version = "0.4.14", default-features = false }
+serde = { version = "1.0.144", features = ["derive"] }
+serde_json = "1.0"
+byteorder = "1.5.0"
+hex = "0.4"
+env_logger = "0.11.5"
+anyhow = "1.0.75"
+
+[build-dependencies]
+zkm-build = { workspace = true }
+[features]
+test = ["zkm-prover/test"]

prover/examples/keccak/host/build.rs

@@ -0,0 +1,3 @@
+fn main() {
+    zkm_build::build_program(&format!("{}/../guest", env!("CARGO_MANIFEST_DIR")));
+}

prover/examples/keccak/host/src/main.rs

@@ -0,0 +1,46 @@
+use std::env;
+
+use zkm_emulator::utils::{load_elf_with_patch, split_prog_into_segs};
+use zkm_utils::utils::prove_segments;
+
+const ELF_PATH: &str = "../guest/elf/mips-zkm-zkvm-elf";
+
+fn prove_keccak_rust() {
+    // 1. split ELF into segs
+    let seg_path = env::var("SEG_OUTPUT").expect("Segment output path is missing");
+    let seg_size = env::var("SEG_SIZE").unwrap_or("65536".to_string());
+    let seg_size = seg_size.parse::<_>().unwrap_or(0);
+
+    let mut state = load_elf_with_patch(ELF_PATH, vec![]);
+    // load input
+    let args = env::var("ARGS").unwrap_or("data-to-hash".to_string());
+    // assume the first arg is the hash output(which is a public input), and the second is the input.
+    let args: Vec<&str> = args.split_whitespace().collect();
+    assert!(args.len() >= 1);
+
+    let public_input: Vec<u8> = hex::decode(args[0]).unwrap();
+    state.add_input_stream(&public_input);
+    log::info!("expected public value in hex: {:X?}", args[0]);
+    log::info!("expected public value: {:X?}", public_input);
+
+    let private_input: Vec<u8> = if args.len() > 1 {
+        hex::decode(args[1]).unwrap()
+    } else {
+        vec![]
+    };
+    log::info!("private input value: {:X?}", private_input);
+    state.add_input_stream(&private_input);
+
+    let (_total_steps, seg_num, mut state) = split_prog_into_segs(state, &seg_path, "", seg_size);
+
+    let value = state.read_public_values::<[u8; 32]>();
+    log::info!("public value: {:X?}", value);
+    log::info!("public value: {} in hex", hex::encode(value));
+
+    let _ = prove_segments(&seg_path, "", "", "", seg_num, 0, vec![]).unwrap();
+}
+
+fn main() {
+    env_logger::try_init().unwrap_or_default();
+    prove_keccak_rust();
+}

prover/examples/utils/src/utils.rs

@@ -16,7 +16,7 @@ use zkm_prover::cpu::kernel::assembler::segment_kernel;
 use zkm_prover::fixed_recursive_verifier::AllRecursiveCircuits;
 use zkm_prover::generation::state::{AssumptionReceipts, Receipt};
 
-const DEGREE_BITS_RANGE: [Range<usize>; 6] = [10..21, 12..22, 12..21, 8..21, 6..21, 13..23];
+const DEGREE_BITS_RANGE: [Range<usize>; 8] = [10..21, 12..22, 11..21, 8..21, 6..21, 6..21, 6..21, 13..23];
 
 const D: usize = 2;
 type C = PoseidonGoldilocksConfig;

prover/src/all_stark.rs

@@ -9,6 +9,11 @@ use plonky2::field::extension::Extendable;
 use plonky2::field::types::Field;
 use plonky2::hash::hash_types::RichField;
 
+use crate::keccak::keccak_stark;
+use crate::keccak::keccak_stark::KeccakStark;
+use crate::keccak_sponge::columns::KECCAK_RATE_BYTES;
+use crate::keccak_sponge::keccak_sponge_stark;
+use crate::keccak_sponge::keccak_sponge_stark::KeccakSpongeStark;
 use crate::logic;
 use crate::logic::LogicStark;
 use crate::memory::memory_stark;
@@ -26,6 +31,8 @@ pub struct AllStark<F: RichField + Extendable<D>, const D: usize> {
     pub cpu_stark: CpuStark<F, D>,
     pub poseidon_stark: PoseidonStark<F, D>,
     pub poseidon_sponge_stark: PoseidonSpongeStark<F, D>,
+    pub keccak_stark: KeccakStark<F, D>,
+    pub keccak_sponge_stark: KeccakSpongeStark<F, D>,
     pub logic_stark: LogicStark<F, D>,
     pub memory_stark: MemoryStark<F, D>,
     pub cross_table_lookups: Vec<CrossTableLookup<F>>,
@@ -38,6 +45,8 @@ impl<F: RichField + Extendable<D>, const D: usize> Default for AllStark<F, D> {
             cpu_stark: CpuStark::default(),
             poseidon_stark: PoseidonStark::default(),
             poseidon_sponge_stark: PoseidonSpongeStark::default(),
+            keccak_stark: KeccakStark::default(),
+            keccak_sponge_stark: KeccakSpongeStark::default(),
             logic_stark: LogicStark::default(),
             memory_stark: MemoryStark::default(),
             cross_table_lookups: all_cross_table_lookups(),
@@ -52,6 +61,8 @@ impl<F: RichField + Extendable<D>, const D: usize> AllStark<F, D> {
             self.cpu_stark.num_lookup_helper_columns(config),
             self.poseidon_stark.num_lookup_helper_columns(config),
             self.poseidon_sponge_stark.num_lookup_helper_columns(config),
+            self.keccak_stark.num_lookup_helper_columns(config),
+            self.keccak_sponge_stark.num_lookup_helper_columns(config),
             self.logic_stark.num_lookup_helper_columns(config),
             self.memory_stark.num_lookup_helper_columns(config),
         ]
@@ -64,8 +75,10 @@ pub enum Table {
     Cpu = 1,
     Poseidon = 2,
     PoseidonSponge = 3,
-    Logic = 4,
-    Memory = 5,
+    Keccak = 4,
+    KeccakSponge = 5,
+    Logic = 6,
+    Memory = 7,
 }
 
 pub(crate) const NUM_TABLES: usize = Table::Memory as usize + 1;
@@ -80,6 +93,8 @@ impl Table {
             Self::Cpu,
             Self::Poseidon,
             Self::PoseidonSponge,
+            Self::Keccak,
+            Self::KeccakSponge,
             Self::Logic,
             Self::Memory,
         ]
@@ -92,6 +107,9 @@ pub(crate) fn all_cross_table_lookups<F: Field>() -> Vec<CrossTableLookup<F>> {
         ctl_poseidon_sponge(),
         ctl_poseidon_inputs(),
         ctl_poseidon_outputs(),
+        ctl_keccak_sponge(),
+        ctl_keccak_inputs(),
+        ctl_keccak_outputs(),
         ctl_logic(),
         ctl_memory(),
     ]
@@ -152,16 +170,72 @@ fn ctl_poseidon_sponge<F: Field>() -> CrossTableLookup<F> {
     CrossTableLookup::new(vec![cpu_looking], poseidon_sponge_looked)
 }
 
+// We now need two different looked tables for `KeccakStark`:
+// one for the inputs and one for the outputs.
+// They are linked with the timestamp.
+fn ctl_keccak_inputs<F: Field>() -> CrossTableLookup<F> {
+    let keccak_sponge_looking = TableWithColumns::new(
+        Table::KeccakSponge,
+        keccak_sponge_stark::ctl_looking_keccak_inputs(),
+        Some(keccak_sponge_stark::ctl_looking_keccak_filter()),
+    );
+    let keccak_looked = TableWithColumns::new(
+        Table::Keccak,
+        keccak_stark::ctl_data_inputs(),
+        Some(keccak_stark::ctl_filter_inputs()),
+    );
+    CrossTableLookup::new(vec![keccak_sponge_looking], keccak_looked)
+}
+
+fn ctl_keccak_outputs<F: Field>() -> CrossTableLookup<F> {
+    let keccak_sponge_looking = TableWithColumns::new(
+        Table::KeccakSponge,
+        keccak_sponge_stark::ctl_looking_keccak_outputs(),
+        Some(keccak_sponge_stark::ctl_looking_keccak_filter()),
+    );
+    let keccak_looked = TableWithColumns::new(
+        Table::Keccak,
+        keccak_stark::ctl_data_outputs(),
+        Some(keccak_stark::ctl_filter_outputs()),
+    );
+    CrossTableLookup::new(vec![keccak_sponge_looking], keccak_looked)
+}
+
+fn ctl_keccak_sponge<F: Field>() -> CrossTableLookup<F> {
+    let cpu_looking = TableWithColumns::new(
+        Table::Cpu,
+        cpu_stark::ctl_data_keccak_sponge(),
+        Some(cpu_stark::ctl_filter_keccak_sponge()),
+    );
+    let keccak_sponge_looked = TableWithColumns::new(
+        Table::KeccakSponge,
+        keccak_sponge_stark::ctl_looked_data(),
+        Some(keccak_sponge_stark::ctl_looked_filter()),
+    );
+    CrossTableLookup::new(vec![cpu_looking], keccak_sponge_looked)
+}
+
 pub(crate) fn ctl_logic<F: Field>() -> CrossTableLookup<F> {
     let cpu_looking = TableWithColumns::new(
         Table::Cpu,
         cpu_stark::ctl_data_logic(),
         Some(cpu_stark::ctl_filter_logic()),
     );
+
+    let mut all_lookers = vec![cpu_looking];
+    for i in 0..keccak_sponge_stark::num_logic_ctls() {
+        let keccak_sponge_looking = TableWithColumns::new(
+            Table::KeccakSponge,
+            keccak_sponge_stark::ctl_looking_logic(i),
+            Some(keccak_sponge_stark::ctl_looking_logic_filter()),
+        );
+        all_lookers.push(keccak_sponge_looking);
+    }
+
     let logic_looked =
         TableWithColumns::new(Table::Logic, logic::ctl_data(), Some(logic::ctl_filter()));
 
-    CrossTableLookup::new(vec![cpu_looking], logic_looked)
+    CrossTableLookup::new(all_lookers, logic_looked)
 }
 
 fn ctl_memory<F: Field>() -> CrossTableLookup<F> {
@@ -179,10 +253,19 @@ fn ctl_memory<F: Field>() -> CrossTableLookup<F> {
             Some(poseidon_sponge_stark::ctl_looking_memory_filter(i)),
         )
     });
+
+    let keccak_sponge_reads = (0..KECCAK_RATE_BYTES).map(|i| {
+        TableWithColumns::new(
+            Table::KeccakSponge,
+            keccak_sponge_stark::ctl_looking_memory(i),
+            Some(keccak_sponge_stark::ctl_looking_memory_filter(i)),
+        )
+    });
     let all_lookers = []
         .into_iter()
         .chain(cpu_memory_gp_ops)
         .chain(poseidon_sponge_reads)
+        .chain(keccak_sponge_reads)
         .collect();
     let memory_looked = TableWithColumns::new(
         Table::Memory,

prover/src/cpu/columns/general.rs

@@ -11,6 +11,7 @@ pub(crate) union CpuGeneralColumnsView<T: Copy> {
     shift: CpuShiftView<T>,
     io: CpuIOAuxView<T>,
     hash: CpuHashView<T>,
+    khash: CpuKHashView<T>,
     misc: CpuMiscView<T>,
 }
 
@@ -25,6 +26,16 @@ impl<T: Copy> CpuGeneralColumnsView<T> {
         unsafe { &mut self.hash }
     }
 
+    // SAFETY: Each view is a valid interpretation of the underlying array.
+    pub(crate) fn khash(&self) -> &CpuKHashView<T> {
+        unsafe { &self.khash }
+    }
+
+    // SAFETY: Each view is a valid interpretation of the underlying array.
+    pub(crate) fn khash_mut(&mut self) -> &mut CpuKHashView<T> {
+        unsafe { &mut self.khash }
+    }
+
     // SAFETY: Each view is a valid interpretation of the underlying array.
     pub(crate) fn syscall(&self) -> &CpuSyscallView<T> {
         unsafe { &self.syscall }
@@ -152,5 +163,10 @@ pub(crate) struct CpuHashView<T: Copy> {
     pub(crate) value: [T; 4],
 }
 
+#[derive(Copy, Clone)]
+pub(crate) struct CpuKHashView<T: Copy> {
+    pub(crate) value: [T; 8],
+}
+
 // `u8` is guaranteed to have a `size_of` of 1.
 pub const NUM_SHARED_COLUMNS: usize = size_of::<CpuGeneralColumnsView<u8>>();

prover/src/cpu/columns/mod.rs

@@ -102,6 +102,7 @@ pub struct CpuColumnsView<T: Copy> {
     // inst_index: [rs_bits, rt_bits, rd_bits, shamt_bits, func_bits]
     /// Filter. 1 iff a Poseidon sponge lookup is performed on this row.
     pub is_poseidon_sponge: T,
+    pub is_keccak_sponge: T,
 
     pub(crate) general: CpuGeneralColumnsView<T>,