rw record ramtype expression (#1098)

Author: hero78119

ceno_zkvm/src/instructions/global.rs

@@ -188,7 +188,7 @@ impl<E: ExtensionField, P> GlobalConfig<E, P> {
 
         let mut record = vec![];
         record.push(addr.expr());
-        record.push(ram_type);
+        record.push(ram_type.clone());
         record.extend(value.memory_expr());
         record.push(local_clk.expr());
 
@@ -205,16 +205,17 @@ impl<E: ExtensionField, P> GlobalConfig<E, P> {
             local_clk.expr(),
         )?;
         // TODO: enforce shard = shard_id in the public values
-
-        cb.read_record(
+        cb.read_rlc_record(
             || "r_record",
-            gkr_iop::RAMType::Memory, // FIXME: should be dynamic, either Register or Memory
+            ram_type.clone(),
             record.clone(),
+            cb.rlc_chip_record(record.clone()),
         )?;
-        cb.write_record(
+        cb.write_rlc_record(
             || "w_record",
-            gkr_iop::RAMType::Memory, // FIXME: should be dynamic, either Register or Memory
+            ram_type,
             record.clone(),
+            cb.rlc_chip_record(record),
         )?;
 
         // enforces final_sum = \sum_i (x_i, y_i) using ecc quark protocol
@@ -441,7 +442,7 @@ impl<E: ExtensionField, P: Permutation<[E::BaseField; POSEIDON2_BABYBEAR_WIDTH]>
         // this is workaround, as call `construct_circuit` will not initialized selector
         // we can remove this one all opcode unittest migrate to call `build_gkr_iop_circuit`
 
-        assert!(num_structural_witin == 3);
+        assert_eq!(num_structural_witin, 3);
         let selector_r_witin = WitIn { id: 0 };
         let selector_w_witin = WitIn { id: 1 };
         let selector_zero_witin = WitIn { id: 2 };

gkr_iop/src/circuit_builder.rs

@@ -419,12 +419,22 @@ impl<E: ExtensionField> ConstraintSystem<E> {
         record: Vec<Expression<E>>,
     ) -> Result<(), CircuitBuilderError> {
         let rlc_record = self.rlc_chip_record(record.clone());
+        self.read_rlc_record(name_fn, (ram_type as u64).into(), record, rlc_record)
+    }
+
+    pub fn read_rlc_record<NR: Into<String>, N: FnOnce() -> NR>(
+        &mut self,
+        name_fn: N,
+        ram_type: Expression<E>,
+        record: Vec<Expression<E>>,
+        rlc_record: Expression<E>,
+    ) -> Result<(), CircuitBuilderError> {
         self.r_expressions.push(rlc_record);
         let path = self.ns.compute_path(name_fn().into());
         self.r_expressions_namespace_map.push(path);
         // Since r_expression is RLC(record) and when we're debugging
         // it's helpful to recover the value of record itself.
-        self.r_ram_types.push(((ram_type as u64).into(), record));
+        self.r_ram_types.push((ram_type, record));
         Ok(())
     }
 
@@ -435,10 +445,22 @@ impl<E: ExtensionField> ConstraintSystem<E> {
         record: Vec<Expression<E>>,
     ) -> Result<(), CircuitBuilderError> {
         let rlc_record = self.rlc_chip_record(record.clone());
+        self.write_rlc_record(name_fn, (ram_type as u64).into(), record, rlc_record)
+    }
+
+    pub fn write_rlc_record<NR: Into<String>, N: FnOnce() -> NR>(
+        &mut self,
+        name_fn: N,
+        ram_type: Expression<E>,
+        record: Vec<Expression<E>>,
+        rlc_record: Expression<E>,
+    ) -> Result<(), CircuitBuilderError> {
         self.w_expressions.push(rlc_record);
         let path = self.ns.compute_path(name_fn().into());
         self.w_expressions_namespace_map.push(path);
-        self.w_ram_types.push(((ram_type as u64).into(), record));
+        // Since w_expression is RLC(record) and when we're debugging
+        // it's helpful to recover the value of record itself.
+        self.w_ram_types.push((ram_type, record));
         Ok(())
     }
 
@@ -696,6 +718,21 @@ impl<'a, E: ExtensionField> CircuitBuilder<'a, E> {
         self.cs.read_record(name_fn, ram_type, record)
     }
 
+    pub fn read_rlc_record<NR, N>(
+        &mut self,
+        name_fn: N,
+        ram_type: Expression<E>,
+        record: Vec<Expression<E>>,
+        rlc_record: Expression<E>,
+    ) -> Result<(), CircuitBuilderError>
+    where
+        NR: Into<String>,
+        N: FnOnce() -> NR,
+    {
+        self.cs
+            .read_rlc_record(name_fn, ram_type, record, rlc_record)
+    }
+
     pub fn write_record<NR, N>(
         &mut self,
         name_fn: N,
@@ -709,6 +746,21 @@ impl<'a, E: ExtensionField> CircuitBuilder<'a, E> {
         self.cs.write_record(name_fn, ram_type, record)
     }
 
+    pub fn write_rlc_record<NR, N>(
+        &mut self,
+        name_fn: N,
+        ram_type: Expression<E>,
+        record: Vec<Expression<E>>,
+        rlc_record: Expression<E>,
+    ) -> Result<(), CircuitBuilderError>
+    where
+        NR: Into<String>,
+        N: FnOnce() -> NR,
+    {
+        self.cs
+            .write_rlc_record(name_fn, ram_type, record, rlc_record)
+    }
+
     pub fn rlc_chip_record(&self, records: Vec<Expression<E>>) -> Expression<E> {
         self.cs.rlc_chip_record(records)
     }