feat(prover): Add Keccak precompile (#519)

* save work

* Wire keccak chips into prover pipeline: trace builder, AIR
  creation, VmAirs registration, and E2E test scaffold

* Fix keccak round trace generation

* Use from_elf_and_logs for keccak test

* Align with spec: shared ECALL bus, MEMW read_addr for x10, pi
  verification and state_ptr ADD constraints

* Add IS_BIT constraints for rbc and equality constraints for rnc constants

* save work

* align keccak with spec

* lint

* revolve conflicts

* Remove dead keccak CPU columns and BusId; verify KAT in test_keccak

* Apply spec Cxz_right Byte Bit optimization in KECCAK_RND

* Reject unaligned/overflowing Keccak state_addr in
  executor and AIR

Mirror alignment + no-overflow in the AIR, dedup
  KECCAK_RC, fix chi indexer, and add a multi-call tiny-keccak cross-check test.

* Range-check Keccak addr bytes with IS_BYTE to close alignment bypass in the AIR

Author: jotabulacios

executor/programs/asm/test_keccak.s

@@ -0,0 +1,38 @@
+	.attribute	5, "rv64i2p1_m2p0_zmmul1p0"
+.Lfunc_end0:
+	.globl	main
+main:
+	# Allocate 200 bytes on the stack for the Keccak state (25 × u64)
+	addi	sp, sp, -200
+
+	# Zero out the state (200 bytes = 25 doublewords)
+	mv	t0, sp
+	li	t1, 25
+.Lzero_loop:
+	sd	zero, 0(t0)
+	addi	t0, t0, 8
+	addi	t1, t1, -1
+	bnez	t1, .Lzero_loop
+
+	# Call keccak-f[1600] permutation
+	# a0 = pointer to 200-byte state
+	# a7 = syscall number (0xFFFFFFFFFFFFFFFE = u64::MAX - 1)
+	mv	a0, sp
+	li	a7, -2
+	ecall
+
+	# Commit the post-permutation state so the test can verify the KAT.
+	# Commit syscall: a0=fd(1), a1=buf_addr, a2=count, a7=64
+	li	a0, 1
+	mv	a1, sp
+	li	a2, 200
+	li	a7, 64
+	ecall
+
+	# Restore stack and halt
+	addi	sp, sp, 200
+	li	a0, 0
+	li	a7, 93
+	ecall
+.Lfunc_end1:
+	.size	main, .Lfunc_end1-main

executor/programs/asm/test_keccak_multi.s

@@ -0,0 +1,48 @@
+	.attribute	5, "rv64i2p1_m2p0_zmmul1p0"
+.Lfunc_end0:
+	.globl	main
+main:
+	# Allocate 200 bytes on the stack for the Keccak state (25 × u64).
+	addi	sp, sp, -200
+
+	# Initialize a non-zero, deterministic state: lane[i] = i + 1.
+	# Used by the host test as the initial state for tiny-keccak::keccakf
+	# cross-checking.
+	mv	t0, sp
+	li	t1, 1
+	li	t2, 26
+.Linit_loop:
+	sd	t1, 0(t0)
+	addi	t0, t0, 8
+	addi	t1, t1, 1
+	bne	t1, t2, .Linit_loop
+
+	# First keccak-f[1600] call.
+	mv	a0, sp
+	li	a7, -2
+	ecall
+
+	# Second keccak-f[1600] call on the result.
+	mv	a0, sp
+	li	a7, -2
+	ecall
+
+	# Third keccak-f[1600] call on the result.
+	mv	a0, sp
+	li	a7, -2
+	ecall
+
+	# Commit the final 200-byte state.
+	li	a0, 1
+	mv	a1, sp
+	li	a2, 200
+	li	a7, 64
+	ecall
+
+	# Restore stack and halt.
+	addi	sp, sp, 200
+	li	a0, 0
+	li	a7, 93
+	ecall
+.Lfunc_end1:
+	.size	main, .Lfunc_end1-main

executor/src/vm/instruction/execution.rs

@@ -8,12 +8,20 @@ use crate::vm::{
 const REGULAR_PC_UPDATE: u64 = 4;
 
 pub enum SyscallNumbers {
+    // Placeholder discriminant. The actual syscall value is KECCAK_SYSCALL_NUMBER.
+    KeccakPermute = 0,
     Print = 1,
     Panic = 2,
     Commit = 64,
     Halt = 93,
 }
 
+/// Syscall number for KeccakPermute (u64::MAX - 1 = 0xFFFF_FFFF_FFFF_FFFE).
+///
+/// Cannot be an enum discriminant because it exceeds isize::MAX.
+pub const KECCAK_SYSCALL_NUMBER: u64 = u64::MAX - 1;
+const KECCAK_STATE_BYTES: u64 = 25 * 8;
+
 impl TryFrom<u64> for SyscallNumbers {
     type Error = ();
     fn try_from(value: u64) -> Result<Self, Self::Error> {
@@ -22,6 +30,7 @@ impl TryFrom<u64> for SyscallNumbers {
             2 => Ok(SyscallNumbers::Panic),
             64 => Ok(SyscallNumbers::Commit),
             93 => Ok(SyscallNumbers::Halt),
+            v if v == KECCAK_SYSCALL_NUMBER => Ok(SyscallNumbers::KeccakPermute),
             _ => Err(()),
         }
     }
@@ -324,6 +333,32 @@ impl Instruction {
                         src2_val = buf_addr;
                         dst_val = count;
                     }
+                    SyscallNumbers::KeccakPermute => {
+                        // keccak-f[1600] permutation on 200 bytes (25 × u64) at address in x10
+                        let state_addr = registers.read(10)?;
+                        if !state_addr.is_multiple_of(8) {
+                            return Err(ExecutionError::UnalignedKeccakStateAddress(state_addr));
+                        }
+                        state_addr
+                            .checked_add(KECCAK_STATE_BYTES - 1)
+                            .ok_or(ExecutionError::KeccakStateAddressOverflow(state_addr))?;
+
+                        let mut state = [0u64; 25];
+                        for (i, lane) in state.iter_mut().enumerate() {
+                            let lane_addr = state_addr
+                                .checked_add((i as u64) * 8)
+                                .ok_or(ExecutionError::KeccakStateAddressOverflow(state_addr))?;
+                            *lane = memory.load_doubleword(lane_addr)?;
+                        }
+                        keccak_f1600(&mut state);
+                        for (i, &lane) in state.iter().enumerate() {
+                            let lane_addr = state_addr
+                                .checked_add((i as u64) * 8)
+                                .ok_or(ExecutionError::KeccakStateAddressOverflow(state_addr))?;
+                            memory.store_doubleword(lane_addr, lane)?;
+                        }
+                        src2_val = state_addr;
+                    }
                     SyscallNumbers::Halt => {
                         // halt
                         return Ok(Log {
@@ -496,4 +531,177 @@ pub enum ExecutionError {
     InvalidWSuffixOperation(ArithOp),
     #[error("Invalid commit fd: expected 1 (stdout), got {0}")]
     InvalidCommitFd(u64),
+    #[error("Unaligned Keccak state address: {0:#018x}")]
+    UnalignedKeccakStateAddress(u64),
+    #[error("Keccak state address range overflows: {0:#018x}")]
+    KeccakStateAddressOverflow(u64),
+}
+
+// =============================================================================
+// Keccak-f[1600] permutation
+// =============================================================================
+
+/// Round constants for Keccak-f[1600] (24 rounds).
+pub const KECCAK_RC: [u64; 24] = [
+    0x0000000000000001,
+    0x0000000000008082,
+    0x800000000000808A,
+    0x8000000080008000,
+    0x000000000000808B,
+    0x0000000080000001,
+    0x8000000080008081,
+    0x8000000000008009,
+    0x000000000000008A,
+    0x0000000000000088,
+    0x0000000080008009,
+    0x000000008000000A,
+    0x000000008000808B,
+    0x800000000000008B,
+    0x8000000000008089,
+    0x8000000000008003,
+    0x8000000000008002,
+    0x8000000000000080,
+    0x000000000000800A,
+    0x800000008000000A,
+    0x8000000080008081,
+    0x8000000000008080,
+    0x0000000080000001,
+    0x8000000080008008,
+];
+
+/// Rotation offsets R[x][y] for the rho step of Keccak-f[1600].
+pub const KECCAK_RHO: [[u32; 5]; 5] = [
+    [0, 36, 3, 41, 18],
+    [1, 44, 10, 45, 2],
+    [62, 6, 43, 15, 61],
+    [28, 55, 25, 21, 56],
+    [27, 20, 39, 8, 14],
+];
+
+/// Apply the Keccak-f[1600] permutation (24 rounds) to a 25-word state.
+///
+/// The state is indexed as `state[x + 5*y]` where `x, y ∈ {0..4}`.
+pub fn keccak_f1600(state: &mut [u64; 25]) {
+    for &rc in &KECCAK_RC {
+        // θ (theta)
+        let mut c = [0u64; 5];
+        for x in 0..5 {
+            c[x] = state[x] ^ state[x + 5] ^ state[x + 10] ^ state[x + 15] ^ state[x + 20];
+        }
+        let mut d = [0u64; 5];
+        for x in 0..5 {
+            d[x] = c[(x + 4) % 5] ^ c[(x + 1) % 5].rotate_left(1);
+        }
+        for x in 0..5 {
+            for y in 0..5 {
+                state[x + 5 * y] ^= d[x];
+            }
+        }
+
+        // ρ (rho) and π (pi)
+        let mut b = [0u64; 25];
+        for x in 0..5 {
+            for y in 0..5 {
+                b[y + 5 * ((2 * x + 3 * y) % 5)] = state[x + 5 * y].rotate_left(KECCAK_RHO[x][y]);
+            }
+        }
+
+        // χ (chi)
+        for x in 0..5 {
+            for y in 0..5 {
+                state[x + 5 * y] =
+                    b[x + 5 * y] ^ (!b[(x + 1) % 5 + 5 * y] & b[(x + 2) % 5 + 5 * y]);
+            }
+        }
+
+        // ι (iota)
+        state[0] ^= rc;
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_keccak_f1600_zero_input() {
+        let mut state = [0u64; 25];
+        keccak_f1600(&mut state);
+
+        let expected: [u64; 25] = [
+            0xF1258F7940E1DDE7,
+            0x84D5CCF933C0478A,
+            0xD598261EA65AA9EE,
+            0xBD1547306F80494D,
+            0x8B284E056253D057,
+            0xFF97A42D7F8E6FD4,
+            0x90FEE5A0A44647C4,
+            0x8C5BDA0CD6192E76,
+            0xAD30A6F71B19059C,
+            0x30935AB7D08FFC64,
+            0xEB5AA93F2317D635,
+            0xA9A6E6260D712103,
+            0x81A57C16DBCF555F,
+            0x43B831CD0347C826,
+            0x01F22F1A11A5569F,
+            0x05E5635A21D9AE61,
+            0x64BEFEF28CC970F2,
+            0x613670957BC46611,
+            0xB87C5A554FD00ECB,
+            0x8C3EE88A1CCF32C8,
+            0x940C7922AE3A2614,
+            0x1841F924A2C509E4,
+            0x16F53526E70465C2,
+            0x75F644E97F30A13B,
+            0xEAF1FF7B5CECA249,
+        ];
+
+        assert_eq!(state, expected, "keccak-f[1600] on zero input mismatch");
+    }
+
+    #[test]
+    fn test_keccak_f1600_nonzero_input() {
+        let mut state = [0u64; 25];
+        state[0] = 1;
+        let original = state;
+        keccak_f1600(&mut state);
+        assert_ne!(state, original);
+        assert!(state.iter().any(|&x| x != 0));
+    }
+
+    #[test]
+    fn test_keccak_syscall_rejects_unaligned_state_addr() {
+        let mut pc = 0;
+        let mut registers = Registers::default();
+        let mut memory = Memory::default();
+
+        registers.write(17, KECCAK_SYSCALL_NUMBER).unwrap();
+        registers.write(10, 0x1001).unwrap();
+
+        let err = Instruction::EcallEbreak
+            .run(&mut pc, &mut registers, &mut memory)
+            .unwrap_err();
+        assert!(matches!(
+            err,
+            ExecutionError::UnalignedKeccakStateAddress(0x1001)
+        ));
+    }
+
+    #[test]
+    fn test_keccak_syscall_rejects_overflowing_state_range() {
+        let mut pc = 0;
+        let mut registers = Registers::default();
+        let mut memory = Memory::default();
+
+        registers.write(17, KECCAK_SYSCALL_NUMBER).unwrap();
+        registers.write(10, u64::MAX - 191).unwrap();
+
+        let err = Instruction::EcallEbreak
+            .run(&mut pc, &mut registers, &mut memory)
+            .unwrap_err();
+        assert!(matches!(
+            err,
+            ExecutionError::KeccakStateAddressOverflow(addr) if addr == u64::MAX - 191
+        ));
+    }
 }

executor/tests/asm.rs

@@ -801,3 +801,49 @@ fn test_sub_64bit() {
 fn test_sub_underflow() {
     run_program("./program_artifacts/asm/sub_underflow.elf");
 }
+
+// ==================== Keccak Precompile ====================
+
+#[test]
+fn test_keccak() {
+    // Runs keccak-f[1600] on a zeroed state and commits the 200-byte result.
+    // Expected output is the FIPS-202 zero-input KAT.
+    let elf_data = std::fs::read("./program_artifacts/asm/test_keccak.elf").unwrap();
+    let program = Elf::load(&elf_data).unwrap();
+    let executor = Executor::new(&program, vec![]).expect("Failed to create executor");
+    let result = executor.run().expect("Failed to run program");
+
+    let expected_state: [u64; 25] = [
+        0xF1258F7940E1DDE7,
+        0x84D5CCF933C0478A,
+        0xD598261EA65AA9EE,
+        0xBD1547306F80494D,
+        0x8B284E056253D057,
+        0xFF97A42D7F8E6FD4,
+        0x90FEE5A0A44647C4,
+        0x8C5BDA0CD6192E76,
+        0xAD30A6F71B19059C,
+        0x30935AB7D08FFC64,
+        0xEB5AA93F2317D635,
+        0xA9A6E6260D712103,
+        0x81A57C16DBCF555F,
+        0x43B831CD0347C826,
+        0x01F22F1A11A5569F,
+        0x05E5635A21D9AE61,
+        0x64BEFEF28CC970F2,
+        0x613670957BC46611,
+        0xB87C5A554FD00ECB,
+        0x8C3EE88A1CCF32C8,
+        0x940C7922AE3A2614,
+        0x1841F924A2C509E4,
+        0x16F53526E70465C2,
+        0x75F644E97F30A13B,
+        0xEAF1FF7B5CECA249,
+    ];
+    let mut expected_bytes = Vec::with_capacity(200);
+    for lane in expected_state {
+        expected_bytes.extend_from_slice(&lane.to_le_bytes());
+    }
+    assert_eq!(result.return_values.memory_values, expected_bytes);
+    assert_eq!(result.return_values.register_values.0, 0);
+}

prover/Cargo.toml

@@ -20,6 +20,7 @@ rayon = { version = "1.8.0", optional = true }
 [dev-dependencies]
 env_logger = "*"
 criterion = { version = "0.5", default-features = false }
+tiny-keccak = { version = "2.0", features = ["keccak"] }
 
 [[bench]]
 name = "vm_prover_benchmark"

prover/src/constraints/cpu.rs

@@ -1033,7 +1033,7 @@ pub fn create_jalr_constraints(constraint_idx_start: usize) -> (Vec<AddConstrain
 
 /// Total number of CPU constraints.
 ///
-/// - IS_BIT: 32 (all bit flags, including read_register1/2)
+/// - IS_BIT: 34 (all bit flags, including read_register1/2 and inline-PC columns)
 /// - ADD carry: 2 (for ADD + LOAD)
 /// - STORE ADD carry: 2 (for STORE: res = arg1 + imm)
 /// - SUB carry: 2 (for SUB + BEQ)