feat: memory chip handler (#271)

Chip handler for memory, identical to that for registers. 

Once the logic is verified, the next step will be to deduplicate.

ceno_rt/src/params.rs

@@ -1,3 +1,3 @@
 pub const WORD_SIZE: usize = 4;
 
-pub const INFO_OUT_ADDR: u32 = 0xC000_0000;
+pub const INFO_OUT_ADDR: u32 = 0xC000_0000;

ceno_zkvm/src/chip_handler.rs

@@ -8,6 +8,7 @@ use crate::{
 
 pub mod general;
 pub mod global_state;
+pub mod memory;
 pub mod register;
 pub mod utils;
 
@@ -42,3 +43,31 @@ pub trait RegisterChipOperations<E: ExtensionField, NR: Into<String>, N: FnOnce(
         value: RegisterExpr<E>,
     ) -> Result<(Expression<E>, ExprLtConfig), ZKVMError>;
 }
+
+/// The common representation of a memory value.
+/// Format: `[u16; 2]`, least-significant-first.
+pub type MemoryExpr<E> = [Expression<E>; 2];
+
+pub trait MemoryChipOperations<E: ExtensionField, NR: Into<String>, N: FnOnce() -> NR> {
+    #[allow(dead_code)]
+    fn memory_read(
+        &mut self,
+        name_fn: N,
+        memory_addr: &WitIn,
+        prev_ts: Expression<E>,
+        ts: Expression<E>,
+        value: crate::chip_handler::MemoryExpr<E>,
+    ) -> Result<(Expression<E>, ExprLtConfig), ZKVMError>;
+
+    #[allow(clippy::too_many_arguments)]
+    #[allow(dead_code)]
+    fn memory_write(
+        &mut self,
+        name_fn: N,
+        memory_addr: &WitIn,
+        prev_ts: Expression<E>,
+        ts: Expression<E>,
+        prev_values: crate::chip_handler::MemoryExpr<E>,
+        value: crate::chip_handler::MemoryExpr<E>,
+    ) -> Result<(Expression<E>, ExprLtConfig), ZKVMError>;
+}

ceno_zkvm/src/chip_handler/memory.rs

@@ -0,0 +1,105 @@
+use crate::{
+    chip_handler::{MemoryChipOperations, MemoryExpr},
+    circuit_builder::CircuitBuilder,
+    error::ZKVMError,
+    expression::{Expression, ToExpr, WitIn},
+    instructions::riscv::config::ExprLtConfig,
+    structs::RAMType,
+};
+use ff_ext::ExtensionField;
+
+impl<'a, E: ExtensionField, NR: Into<String>, N: FnOnce() -> NR> MemoryChipOperations<E, NR, N>
+    for CircuitBuilder<'a, E>
+{
+    #[allow(dead_code)]
+    fn memory_read(
+        &mut self,
+        name_fn: N,
+        memory_addr: &WitIn,
+        prev_ts: Expression<E>,
+        ts: Expression<E>,
+        value: MemoryExpr<E>,
+    ) -> Result<(Expression<E>, ExprLtConfig), ZKVMError> {
+        self.namespace(name_fn, |cb| {
+            // READ (a, v, t)
+            let read_record = cb.rlc_chip_record(
+                [
+                    vec![Expression::<E>::Constant(E::BaseField::from(
+                        RAMType::Memory as u64,
+                    ))],
+                    vec![memory_addr.expr()],
+                    value.to_vec(),
+                    vec![prev_ts.clone()],
+                ]
+                .concat(),
+            );
+            // Write (a, v, t)
+            let write_record = cb.rlc_chip_record(
+                [
+                    vec![Expression::<E>::Constant(E::BaseField::from(
+                        RAMType::Memory as u64,
+                    ))],
+                    vec![memory_addr.expr()],
+                    value.to_vec(),
+                    vec![ts.clone()],
+                ]
+                .concat(),
+            );
+            cb.read_record(|| "read_record", read_record)?;
+            cb.write_record(|| "write_record", write_record)?;
+
+            // assert prev_ts < current_ts
+            let lt_cfg = cb.less_than(|| "prev_ts < ts", prev_ts, ts.clone(), Some(true))?;
+
+            let next_ts = ts + 1.into();
+
+            Ok((next_ts, lt_cfg))
+        })
+    }
+
+    #[allow(dead_code)]
+    fn memory_write(
+        &mut self,
+        name_fn: N,
+        memory_addr: &WitIn,
+        prev_ts: Expression<E>,
+        ts: Expression<E>,
+        prev_values: MemoryExpr<E>,
+        value: MemoryExpr<E>,
+    ) -> Result<(Expression<E>, ExprLtConfig), ZKVMError> {
+        self.namespace(name_fn, |cb| {
+            // READ (a, v, t)
+            let read_record = cb.rlc_chip_record(
+                [
+                    vec![Expression::<E>::Constant(E::BaseField::from(
+                        RAMType::Memory as u64,
+                    ))],
+                    vec![memory_addr.expr()],
+                    prev_values.to_vec(),
+                    vec![prev_ts.clone()],
+                ]
+                .concat(),
+            );
+            // Write (a, v, t)
+            let write_record = cb.rlc_chip_record(
+                [
+                    vec![Expression::<E>::Constant(E::BaseField::from(
+                        RAMType::Memory as u64,
+                    ))],
+                    vec![memory_addr.expr()],
+                    value.to_vec(),
+                    vec![ts.clone()],
+                ]
+                .concat(),
+            );
+            cb.read_record(|| "read_record", read_record)?;
+            cb.write_record(|| "write_record", write_record)?;
+
+            let lt_cfg = cb.less_than(|| "prev_ts < ts", prev_ts, ts.clone(), Some(true))?;
+
+            let next_ts = ts + 1.into();
+
+            Ok((next_ts, lt_cfg))
+        })
+    }
+}

ceno_zkvm/src/structs.rs

@@ -56,6 +56,7 @@ pub enum ROMType {
 pub enum RAMType {
     GlobalState,
     Register,
+    Memory,
 }
 
 /// A point is a vector of num_var length

mpcs/benches/basecode.rs

@@ -9,7 +9,8 @@ use mpcs::{
     util::{
         arithmetic::interpolate_field_type_over_boolean_hypercube,
         plonky2_util::reverse_index_bits_in_place_field_type,
-    }, Basefold, BasefoldBasecodeParams, BasefoldSpec, EncodingScheme, PolynomialCommitmentScheme
+    },
+    Basefold, BasefoldBasecodeParams, BasefoldSpec, EncodingScheme, PolynomialCommitmentScheme,
 };
 
 use multilinear_extensions::mle::{DenseMultilinearExtension, FieldType};
@@ -55,7 +56,6 @@ fn bench_encoding(c: &mut Criterion, is_base: bool) {
                 })
                 .collect_vec();
 
-
             group.bench_function(
                 BenchmarkId::new("batch_encode", format!("{}-{}", num_vars, batch_size)),
                 |b| {

mpcs/benches/commit_open_verify_rs.rs

@@ -10,11 +10,13 @@ use mpcs::{
     PolynomialCommitmentScheme,
 };
 
-use multilinear_extensions::mle::{DenseMultilinearExtension, MultilinearExtension};
+use multilinear_extensions::{
+    mle::{DenseMultilinearExtension, MultilinearExtension},
+    virtual_poly_v2::ArcMultilinearExtension,
+};
 use rand::{rngs::OsRng, SeedableRng};
 use rand_chacha::ChaCha8Rng;
 use rayon::iter::{IntoParallelIterator, ParallelIterator};
-use multilinear_extensions::virtual_poly_v2::ArcMultilinearExtension;
 use transcript::Transcript;
 
 type Pcs = Basefold<GoldilocksExt2, BasefoldRSParams, ChaCha8Rng>;
@@ -293,9 +295,8 @@ fn bench_simple_batch_commit_open_verify_goldilocks(c: &mut Criterion, is_base:
                 },
             );
 
-            let polys: Vec<ArcMultilinearExtension<GoldilocksExt2>> = polys.into_iter()
-                .map(|poly| poly.into())
-                .collect_vec();
+            let polys: Vec<ArcMultilinearExtension<GoldilocksExt2>> =
+                polys.into_iter().map(|poly| poly.into()).collect_vec();
 
             let point = (0..num_vars)
                 .map(|_| transcript.get_and_append_challenge(b"Point").elements)

mpcs/benches/hashing.rs

@@ -1,8 +1,8 @@
 use ark_std::test_rng;
-use criterion::{Criterion, criterion_group, criterion_main};
+use criterion::{criterion_group, criterion_main, Criterion};
 use ff::Field;
 use goldilocks::Goldilocks;
-use mpcs::util::hash::{Digest, DIGEST_WIDTH, hash_two_digests, new_hasher};
+use mpcs::util::hash::{hash_two_digests, new_hasher, Digest, DIGEST_WIDTH};
 
 fn random_ceno_goldy() -> Goldilocks {
     Goldilocks::random(&mut test_rng())
@@ -36,4 +36,4 @@ pub fn criterion_benchmark(c: &mut Criterion) {
 }
 
 criterion_group!(benches, criterion_benchmark);
-criterion_main!(benches);
+criterion_main!(benches);

mpcs/benches/rscode.rs

@@ -10,8 +10,7 @@ use mpcs::{
         arithmetic::interpolate_field_type_over_boolean_hypercube,
         plonky2_util::reverse_index_bits_in_place_field_type,
     },
-    Basefold, BasefoldRSParams, BasefoldSpec, EncodingScheme,
-    PolynomialCommitmentScheme,
+    Basefold, BasefoldRSParams, BasefoldSpec, EncodingScheme, PolynomialCommitmentScheme,
 };
 
 use multilinear_extensions::mle::{DenseMultilinearExtension, FieldType};
@@ -57,7 +56,6 @@ fn bench_encoding(c: &mut Criterion, is_base: bool) {
                 })
                 .collect_vec();
 
-
             group.bench_function(
                 BenchmarkId::new("batch_encode", format!("{}-{}", num_vars, batch_size)),
                 |b| {