Add uint256 emulation program

Author: dreamATD

Cargo.lock

@@ -17,16 +17,19 @@ elf = "0.7"
 ff_ext.workspace = true
 itertools.workspace = true
 multilinear_extensions.workspace = true
+num.workspace = true
 num-derive.workspace = true
 num-traits.workspace = true
 rrs_lib = { package = "rrs-succinct", version = "0.1.0" }
 secp.workspace = true
 serde.workspace = true
+sp1-curves.workspace = true
 strum.workspace = true
 strum_macros.workspace = true
 substrate-bn.workspace = true
 tiny-keccak.workspace = true
 tracing.workspace = true
+typenum = "1.19.0"
 
 [features]
 default = ["forbid_overflow"]

ceno_emul/Cargo.toml

@@ -38,6 +38,7 @@ pub use syscalls::{
         Secp256k1DecompressSpec, Secp256k1DoubleSpec,
     },
     sha256::{SHA_EXTEND_WORDS, Sha256ExtendSpec},
+    uint256::UINT256_WORDS_FIELD_ELEMENT,
 };
 
 pub mod utils;

ceno_emul/src/lib.rs

@@ -5,6 +5,7 @@ pub mod bn254;
 pub mod keccak_permute;
 pub mod secp256k1;
 pub mod sha256;
+pub mod uint256;
 
 // Using the same function codes as sp1:
 // https://github.com/succinctlabs/sp1/blob/013c24ea2fa15a0e7ed94f7d11a7ada4baa39ab9/crates/core/executor/src/syscalls/code.rs
@@ -13,7 +14,7 @@ pub use ceno_syscall::{
     BLS12381_ADD, BLS12381_DECOMPRESS, BLS12381_DOUBLE, BN254_ADD, BN254_DOUBLE, BN254_FP_ADD,
     BN254_FP_MUL, BN254_FP2_ADD, BN254_FP2_MUL, KECCAK_PERMUTE, SECP256K1_ADD,
     SECP256K1_DECOMPRESS, SECP256K1_DOUBLE, SECP256R1_ADD, SECP256R1_DECOMPRESS, SECP256R1_DOUBLE,
-    SHA_EXTEND,
+    SHA_EXTEND, UINT256_MUL,
 };
 
 pub trait SyscallSpec {
@@ -42,6 +43,7 @@ pub fn handle_syscall(vm: &VMState, function_code: u32) -> Result<SyscallEffects
         BN254_FP_MUL => Ok(bn254::bn254_fp_mul(vm)),
         BN254_FP2_ADD => Ok(bn254::bn254_fp2_add(vm)),
         BN254_FP2_MUL => Ok(bn254::bn254_fp2_mul(vm)),
+        UINT256_MUL => Ok(uint256::uint256_mul(vm)),
         // TODO: introduce error types.
         _ => Err(anyhow::anyhow!("Unknown syscall: {}", function_code)),
     }

ceno_emul/src/syscalls.rs

@@ -0,0 +1,91 @@
+use crate::{
+    Change, EmuContext, Platform, SyscallSpec, VMState, Word, WriteOp,
+    syscalls::{SyscallEffects, SyscallWitness},
+    utils::MemoryView,
+};
+
+use super::UINT256_MUL;
+use itertools::Itertools;
+use num::{BigUint, One, Zero};
+use sp1_curves::{params::NumWords, uint256::U256Field};
+use typenum::marker_traits::Unsigned;
+
+type WordsFieldElement = <U256Field as NumWords>::WordsFieldElement;
+pub const UINT256_WORDS_FIELD_ELEMENT: usize = WordsFieldElement::USIZE;
+const WORD_SIZE: usize = 4;
+
+pub(crate) struct Uint256MulSpec;
+
+impl SyscallSpec for Uint256MulSpec {
+    const NAME: &'static str = "UINT256_MUL";
+
+    const REG_OPS_COUNT: usize = 2;
+    const MEM_OPS_COUNT: usize = 3 * UINT256_WORDS_FIELD_ELEMENT; // x, y, modulus
+    const CODE: u32 = ceno_rt::syscalls::UINT256_MUL;
+}
+
+pub fn uint256_mul(vm: &VMState) -> SyscallEffects {
+    let x_ptr = vm.peek_register(Platform::reg_arg0());
+    let y_ptr = vm.peek_register(Platform::reg_arg1());
+    let mod_ptr = y_ptr + UINT256_WORDS_FIELD_ELEMENT as u32 * WORD_SIZE as u32;
+
+    // Read the argument pointers
+    let reg_ops = vec![
+        WriteOp::new_register_op(
+            Platform::reg_arg0(),
+            Change::new(x_ptr, x_ptr),
+            0, // Cycle set later in finalize().
+        ),
+        WriteOp::new_register_op(
+            Platform::reg_arg1(),
+            Change::new(y_ptr, y_ptr),
+            0, // Cycle set later in finalize().
+        ),
+    ];
+
+    // Memory segments of x, y, and modulus
+    let [mut x_view, y_view, mod_view] =
+        [x_ptr, y_ptr, mod_ptr].map(|start| MemoryView::<UINT256_WORDS_FIELD_ELEMENT>::new(vm, start));
+
+    // Read x and y from words via wrapper type
+    let [x, y, modulus] = [&x_view, &y_view, &mod_view].map(|view| {
+        BigUint::from_bytes_le(
+            &view
+                .words()
+                .into_iter()
+                .flat_map(|w| w.to_le_bytes())
+                .collect_vec(),
+        )
+    });
+
+    // Perform the multiplication and take the result modulo the modulus.
+    let result: BigUint = if modulus.is_zero() {
+        let modulus = BigUint::one() << 256;
+        (x * y) % modulus
+    } else {
+        (x * y) % modulus
+    };
+    let mut result_bytes = result.to_bytes_le();
+    result_bytes.resize(32, 0u8); // Pad the result to 32 bytes.
+
+    // Convert the result to little endian u32 words.
+    let result: [u32; 8] = {
+        let mut iter = result_bytes
+            .chunks_exact(4)
+            .map(|chunk| u32::from_le_bytes(chunk.try_into().unwrap()));
+        core::array::from_fn(|_| iter.next().unwrap())
+    };
+    x_view.write(result);
+
+    let mem_ops = x_view
+        .mem_ops()
+        .into_iter()
+        .chain(y_view.mem_ops())
+        .chain(mod_view.mem_ops())
+        .collect_vec();
+
+    SyscallEffects {
+        witness: SyscallWitness::new(mem_ops, reg_ops),
+        next_pc: None,
+    }
+}

ceno_emul/src/syscalls/uint256.rs

@@ -4,7 +4,7 @@ use anyhow::Result;
 use ceno_emul::{
     BN254_FP_WORDS, BN254_FP2_WORDS, BN254_POINT_WORDS, CENO_PLATFORM, EmuContext, InsnKind,
     Platform, Program, SECP256K1_ARG_WORDS, SECP256K1_COORDINATE_WORDS, SHA_EXTEND_WORDS,
-    StepRecord, VMState, WORD_SIZE, Word, WordAddr, WriteOp,
+    StepRecord, UINT256_WORDS_FIELD_ELEMENT, VMState, WORD_SIZE, Word, WordAddr, WriteOp,
     host_utils::{read_all_messages, read_all_messages_as_words},
 };
 use ceno_host::CenoStdin;
@@ -623,6 +623,61 @@ fn test_bn254_precompile() -> Result<()> {
     let program_elf = ceno_examples::bn254_precompile;
     let mut state = VMState::new_from_elf(unsafe_platform(), program_elf)?;
     let _ = run(&mut state)?;
+    Ok(())
+}
+
+fn test_uint256_mul() -> Result<()> {
+    let mut state = VMState::new_from_elf(unsafe_platform(), program_elf)?;
+
+    let steps = run(&mut state)?;
+
+    let syscalls = steps.iter().filter_map(|step| step.syscall()).collect_vec();
+    assert_eq!(syscalls.len(), 1);
+
+    let witness = syscalls[0];
+    assert_eq!(witness.reg_ops.len(), 2);
+    assert_eq!(witness.reg_ops[0].register_index(), Platform::reg_arg0());
+    assert_eq!(witness.reg_ops[1].register_index(), Platform::reg_arg1());
+
+    let a_address = witness.reg_ops[0].value.after;
+    assert_eq!(a_address, witness.reg_ops[0].value.before);
+    let a_address: WordAddr = a_address.into();
+
+    let b_address = witness.reg_ops[1].value.after;
+    assert_eq!(b_address, witness.reg_ops[1].value.before);
+    let b_address: WordAddr = b_address.into();
+
+    const A_MUL_B: [u8; 65] = [
+        4, 188, 11, 115, 232, 35, 63, 79, 186, 163, 11, 207, 165, 64, 247, 109, 81, 125, 56, 83,
+        131, 221, 140, 154, 19, 186, 109, 173, 9, 127, 142, 169, 219, 108, 17, 216, 218, 125, 37,
+        30, 87, 86, 194, 151, 20, 122, 64, 118, 123, 210, 29, 60, 209, 138, 131, 11, 247, 157, 212,
+        209, 123, 162, 111, 197, 70,
+    ];
+    let expect = bytes_to_words(A_MUL_B);
+
+    assert_eq!(witness.mem_ops.len(), 3 * UINT256_WORDS_FIELD_ELEMENT);
+    // Expect first half to consist of read/writes on P
+    for (i, write_op) in witness
+        .mem_ops
+        .iter()
+        .take(UINT256_WORDS_FIELD_ELEMENT)
+        .enumerate()
+    {
+        assert_eq!(write_op.addr, a_address + i);
+        assert_eq!(write_op.value.after, expect[i]);
+    }
+
+    // Expect second half to consist of reads on Q
+    for (i, write_op) in witness
+        .mem_ops
+        .iter()
+        .skip(UINT256_WORDS_FIELD_ELEMENT)
+        .take(UINT256_WORDS_FIELD_ELEMENT * UINT256_WORDS_FIELD_ELEMENT)
+        .enumerate()
+    {
+        assert_eq!(write_op.addr, b_address + i);
+        assert_eq!(write_op.value.after, write_op.value.before);
+    }
 
     Ok(())
 }

ceno_host/tests/test_elf.rs

@@ -18,6 +18,7 @@ pub const BLS12381_DOUBLE: u32 = 0x00_00_01_1F;
 pub const SECP256R1_ADD: u32 = 0x00_01_01_2C;
 pub const SECP256R1_DOUBLE: u32 = 0x00_00_01_2D;
 pub const SECP256R1_DECOMPRESS: u32 = 0x00_00_01_2E;
+pub const UINT256_MUL: u32 = 0x00_01_01_1D;
 
 pub const KECCAK_STATE_WORDS: usize = 25;
 

ceno_syscall/src/lib.rs

@@ -3,11 +3,13 @@ mod keccak;
 mod weierstrass_add;
 mod weierstrass_decompress;
 mod weierstrass_double;
+mod uint256;
 
 pub use keccak::KeccakInstruction;
 pub use weierstrass_add::WeierstrassAddAssignInstruction;
 pub use weierstrass_decompress::WeierstrassDecompressInstruction;
 pub use weierstrass_double::WeierstrassDoubleAssignInstruction;
+pub use uint256::U256MulInstruction;
 
 use ceno_emul::InsnKind;
 pub use halt::HaltInstruction;