secp256k1 and sha256_extend mocks + refactor (#826)

- all precompiles have the same interface as their counter-parts in
`sp1`; the `secp256k1` curve ops are backed by the `secp` crate.
- integrates appropriate dummy circuits for precompiles into the zkvm
- added `utils.rs` with some utilities for safely manipulating memory
segments of the VM
- tested the compatibility of the syscalls (+ keccak, done previously)
with sp1 (see the `syscalls.rs` example file).
- extended the docs pasted from sp1

Aurel has reviewed an earlier version of this
[here](Inversed-Tech#4).

Cargo.toml

@@ -46,6 +46,7 @@ rand_chacha = { version = "0.3", features = ["serde1"] }
 rand_core = "0.6"
 rand_xorshift = "0.3"
 rayon = "1.10"
+secp = "0.4.1"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 strum = "0.26"

ceno_emul/Cargo.toml

@@ -17,6 +17,7 @@ itertools.workspace = true
 num-derive.workspace = true
 num-traits.workspace = true
 rrs_lib = { package = "rrs-succinct", version = "0.1.0" }
+secp.workspace = true
 strum.workspace = true
 strum_macros.workspace = true
 tiny-keccak.workspace = true

ceno_emul/src/lib.rs

@@ -23,7 +23,17 @@ pub use elf::Program;
 pub mod disassemble;
 
 mod syscalls;
-pub use syscalls::{KECCAK_PERMUTE, keccak_permute::KECCAK_WORDS};
+pub use syscalls::{
+    KECCAK_PERMUTE, SECP256K1_ADD, SECP256K1_DECOMPRESS, SECP256K1_DOUBLE, SHA_EXTEND, SyscallSpec,
+    keccak_permute::{KECCAK_WORDS, KeccakSpec},
+    secp256k1::{
+        COORDINATE_WORDS, SECP256K1_ARG_WORDS, Secp256k1AddSpec, Secp256k1DecompressSpec,
+        Secp256k1DoubleSpec,
+    },
+    sha256::{SHA_EXTEND_WORDS, Sha256ExtendSpec},
+};
+
+pub mod utils;
 
 pub mod test_utils;
 

ceno_emul/src/syscalls.rs

@@ -2,16 +2,32 @@ use crate::{RegIdx, Tracer, VMState, Word, WordAddr, WriteOp};
 use anyhow::Result;
 
 pub mod keccak_permute;
+pub mod secp256k1;
+pub mod sha256;
 
 // Using the same function codes as sp1:
 // https://github.com/succinctlabs/sp1/blob/013c24ea2fa15a0e7ed94f7d11a7ada4baa39ab9/crates/core/executor/src/syscalls/code.rs
 
-pub use ceno_rt::syscalls::KECCAK_PERMUTE;
+pub use ceno_rt::syscalls::{
+    KECCAK_PERMUTE, SECP256K1_ADD, SECP256K1_DECOMPRESS, SECP256K1_DOUBLE, SHA_EXTEND,
+};
+
+pub trait SyscallSpec {
+    const NAME: &'static str;
+
+    const REG_OPS_COUNT: usize;
+    const MEM_OPS_COUNT: usize;
+    const CODE: u32;
+}
 
 /// Trace the inputs and effects of a syscall.
 pub fn handle_syscall(vm: &VMState, function_code: u32) -> Result<SyscallEffects> {
     match function_code {
         KECCAK_PERMUTE => Ok(keccak_permute::keccak_permute(vm)),
+        SECP256K1_ADD => Ok(secp256k1::secp256k1_add(vm)),
+        SECP256K1_DOUBLE => Ok(secp256k1::secp256k1_double(vm)),
+        SECP256K1_DECOMPRESS => Ok(secp256k1::secp256k1_decompress(vm)),
+        SHA_EXTEND => Ok(sha256::extend(vm)),
         // TODO: introduce error types.
         _ => Err(anyhow::anyhow!("Unknown syscall: {}", function_code)),
     }
@@ -22,6 +38,24 @@ pub fn handle_syscall(vm: &VMState, function_code: u32) -> Result<SyscallEffects
 pub struct SyscallWitness {
     pub mem_ops: Vec<WriteOp>,
     pub reg_ops: Vec<WriteOp>,
+    _marker: (),
+}
+
+impl SyscallWitness {
+    fn new(mem_ops: Vec<WriteOp>, reg_ops: Vec<WriteOp>) -> SyscallWitness {
+        for (i, op) in mem_ops.iter().enumerate() {
+            assert_eq!(
+                op.addr,
+                mem_ops[0].addr + i,
+                "Dummy circuit expects that mem_ops addresses are consecutive."
+            );
+        }
+        SyscallWitness {
+            mem_ops,
+            reg_ops,
+            _marker: (),
+        }
+    }
 }
 
 /// The effects of a syscall to apply on the VM.

ceno_emul/src/syscalls/keccak_permute.rs

@@ -1,65 +1,87 @@
-use itertools::{Itertools, izip};
+use itertools::Itertools;
 use tiny_keccak::keccakf;
 
-use crate::{Change, EmuContext, Platform, VMState, WORD_SIZE, WordAddr, WriteOp};
+use crate::{Change, EmuContext, Platform, VMState, Word, WriteOp, utils::MemoryView};
 
-use super::{SyscallEffects, SyscallWitness};
+use super::{SyscallEffects, SyscallSpec, SyscallWitness};
 
 const KECCAK_CELLS: usize = 25; // u64 cells
 pub const KECCAK_WORDS: usize = KECCAK_CELLS * 2; // u32 words
 
+pub struct KeccakSpec;
+
+impl SyscallSpec for KeccakSpec {
+    const NAME: &'static str = "KECCAK";
+
+    const REG_OPS_COUNT: usize = 2;
+    const MEM_OPS_COUNT: usize = KECCAK_WORDS;
+    const CODE: u32 = ceno_rt::syscalls::KECCAK_PERMUTE;
+}
+
+/// Wrapper type for the keccak_permute argument that implements conversions
+/// from and to VM word-representations according to the syscall spec
+pub struct KeccakState(pub [u64; KECCAK_CELLS]);
+
+impl From<[Word; KECCAK_WORDS]> for KeccakState {
+    fn from(words: [Word; KECCAK_WORDS]) -> Self {
+        KeccakState(
+            words
+                .chunks_exact(2)
+                .map(|chunk| (chunk[0] as u64 | ((chunk[1] as u64) << 32)))
+                .collect_vec()
+                .try_into()
+                .expect("failed to parse words into [u64; 25]"),
+        )
+    }
+}
+
+impl From<KeccakState> for [Word; KECCAK_WORDS] {
+    fn from(state: KeccakState) -> [Word; KECCAK_WORDS] {
+        state
+            .0
+            .iter()
+            .flat_map(|&elem| [elem as u32, (elem >> 32) as u32])
+            .collect_vec()
+            .try_into()
+            .unwrap()
+    }
+}
+
 /// Trace the execution of a Keccak permutation.
 ///
 /// Compatible with:
 /// https://github.com/succinctlabs/sp1/blob/013c24ea2fa15a0e7ed94f7d11a7ada4baa39ab9/crates/core/executor/src/syscalls/precompiles/keccak256/permute.rs
-///
-/// TODO: test compatibility.
 pub fn keccak_permute(vm: &VMState) -> SyscallEffects {
     let state_ptr = vm.peek_register(Platform::reg_arg0());
 
-    // Read the argument `state_ptr`.
-    let reg_ops = vec![WriteOp::new_register_op(
-        Platform::reg_arg0(),
-        Change::new(state_ptr, state_ptr),
-        0, // Cycle set later in finalize().
-    )];
-
-    let addrs = (state_ptr..)
-        .step_by(WORD_SIZE)
-        .take(KECCAK_WORDS)
-        .map(WordAddr::from)
-        .collect_vec();
-
-    // Read Keccak state.
-    let input = addrs
-        .iter()
-        .map(|&addr| vm.peek_memory(addr))
-        .collect::<Vec<_>>();
+    // for compatibility with sp1 spec
+    assert_eq!(vm.peek_register(Platform::reg_arg1()), 0);
 
-    // Compute Keccak permutation.
-    let output = {
-        let mut state = [0_u64; KECCAK_CELLS];
-        for (cell, (&lo, &hi)) in izip!(&mut state, input.iter().tuples()) {
-            *cell = lo as u64 | ((hi as u64) << 32);
-        }
-
-        keccakf(&mut state);
+    // Read the argument `state_ptr`.
+    let reg_ops = vec![
+        WriteOp::new_register_op(
+            Platform::reg_arg0(),
+            Change::new(state_ptr, state_ptr),
+            0, // Cycle set later in finalize().
+        ),
+        WriteOp::new_register_op(
+            Platform::reg_arg1(),
+            Change::new(0, 0),
+            0, // Cycle set later in finalize().
+        ),
+    ];
 
-        state.into_iter().flat_map(|c| [c as u32, (c >> 32) as u32])
-    };
+    let mut state_view = MemoryView::<KECCAK_WORDS>::new(vm, state_ptr);
+    let mut state = KeccakState::from(state_view.words());
+    keccakf(&mut state.0);
+    let output_words: [Word; KECCAK_WORDS] = state.into();
 
-    // Write permuted state.
-    let mem_ops = izip!(addrs, input, output)
-        .map(|(addr, before, after)| WriteOp {
-            addr,
-            value: Change { before, after },
-            previous_cycle: 0, // Cycle set later in finalize().
-        })
-        .collect_vec();
+    state_view.write(output_words);
+    let mem_ops: Vec<WriteOp> = state_view.mem_ops().to_vec();
 
     assert_eq!(mem_ops.len(), KECCAK_WORDS);
     SyscallEffects {
-        witness: SyscallWitness { mem_ops, reg_ops },
+        witness: SyscallWitness::new(mem_ops, reg_ops),
         next_pc: None,
     }
 }