feat: group expressions by selectors

ceno_recursion/src/zkvm_verifier/verifier.rs

@@ -540,9 +540,9 @@ pub fn verify_chip_proof<C: Config>(
     } = &composed_cs;
     let one: Ext<C::F, C::EF> = builder.constant(C::EF::ONE);
 
-    let r_len = cs.r_expressions.len() + cs.r_table_expressions.len();
-    let w_len = cs.w_expressions.len() + cs.w_table_expressions.len();
-    let lk_len = cs.lk_expressions.len() + cs.lk_table_expressions.len();
+    let r_len = cs.r_expressions_len() + cs.r_table_expressions_len();
+    let w_len = cs.w_expressions_len() + cs.w_table_expressions_len();
+    let lk_len = cs.lk_expressions_len() + cs.lk_table_expressions_len();
     let num_batched = r_len + w_len + lk_len;
 
     let r_counts_per_instance: Usize<C::N> = Usize::from(r_len);
@@ -609,7 +609,7 @@ pub fn verify_chip_proof<C: Config>(
     );
     builder.cycle_tracker_end(format!("verify tower proof for opcode {circuit_name}",).as_str());
 
-    if cs.lk_table_expressions.is_empty() {
+    if cs.lk_table_expressions_len() == 0 {
         builder
             .range(0, logup_p_evals.len())
             .for_each(|idx_vec, builder| {
@@ -624,7 +624,7 @@ pub fn verify_chip_proof<C: Config>(
     builder.assert_usize_eq(logup_q_evals.len(), lk_counts_per_instance.clone());
 
     // GKR circuit
-    let out_evals_len: Usize<C::N> = if cs.lk_table_expressions.is_empty() {
+    let out_evals_len: Usize<C::N> = if cs.lk_table_expressions_len() == 0 {
         builder.eval(record_evals.len() + logup_q_evals.len())
     } else {
         builder.eval(record_evals.len() + logup_p_evals.len() + logup_q_evals.len())
@@ -639,7 +639,7 @@ pub fn verify_chip_proof<C: Config>(
         });
 
     let end: Usize<C::N> = Usize::uninit(builder);
-    if !cs.lk_table_expressions.is_empty() {
+    if cs.lk_table_expressions_len() > 0 {
         builder.assign(&end, record_evals.len() + logup_p_evals.len());
         let p_slice = out_evals.slice(builder, record_evals.len(), end.clone());
 
@@ -688,7 +688,7 @@ pub fn verify_chip_proof<C: Config>(
             gkr_circuit
                 .layers
                 .first()
-                .map(|layer| layer.out_sel_and_eval_exprs.len())
+                .map(|layer| layer.selector_ctxs_len())
                 .unwrap_or(0)
         ]
     } else {

ceno_zkvm/src/chip_handler.rs

@@ -1,5 +1,5 @@
 use ff_ext::ExtensionField;
-use gkr_iop::{error::CircuitBuilderError, gadgets::AssertLtConfig};
+use gkr_iop::{error::CircuitBuilderError, gadgets::AssertLtConfig, selector::SelectorType};
 
 use crate::instructions::riscv::constants::UINT_LIMBS;
 use multilinear_extensions::{Expression, ToExpr};
@@ -53,6 +53,7 @@ pub type AddressExpr<E> = Expression<E>;
 /// Format: `[u16; UINT_LIMBS]`, least-significant-first.
 pub type MemoryExpr<E> = [Expression<E>; UINT_LIMBS];
 
+#[allow(clippy::type_complexity)]
 pub trait MemoryChipOperations<E: ExtensionField, NR: Into<String>, N: FnOnce() -> NR> {
     fn memory_read(
         &mut self,
@@ -73,4 +74,16 @@ pub trait MemoryChipOperations<E: ExtensionField, NR: Into<String>, N: FnOnce()
         prev_values: MemoryExpr<E>,
         value: MemoryExpr<E>,
     ) -> Result<(Expression<E>, AssertLtConfig), CircuitBuilderError>;
+
+    fn memory_write_with_rw_selectors(
+        &mut self,
+        name_fn: N,
+        memory_addr: &AddressExpr<E>,
+        prev_ts: Expression<E>,
+        ts: Expression<E>,
+        prev_values: MemoryExpr<E>,
+        value: MemoryExpr<E>,
+        r_selector: &SelectorType<E>,
+        w_selector: &SelectorType<E>,
+    ) -> Result<(Expression<E>, AssertLtConfig), CircuitBuilderError>;
 }

ceno_zkvm/src/chip_handler/memory.rs

@@ -5,7 +5,7 @@ use crate::{
     structs::RAMType,
 };
 use ff_ext::ExtensionField;
-use gkr_iop::error::CircuitBuilderError;
+use gkr_iop::{error::CircuitBuilderError, selector::SelectorType};
 use multilinear_extensions::Expression;
 
 impl<E: ExtensionField, NR: Into<String>, N: FnOnce() -> NR> MemoryChipOperations<E, NR, N>
@@ -48,4 +48,28 @@ impl<E: ExtensionField, NR: Into<String>, N: FnOnce() -> NR> MemoryChipOperation
             value,
         )
     }
+
+    fn memory_write_with_rw_selectors(
+        &mut self,
+        name_fn: N,
+        memory_addr: &AddressExpr<E>,
+        prev_ts: Expression<E>,
+        ts: Expression<E>,
+        prev_values: MemoryExpr<E>,
+        value: MemoryExpr<E>,
+        r_selector: &SelectorType<E>,
+        w_selector: &SelectorType<E>,
+    ) -> Result<(Expression<E>, AssertLtConfig), CircuitBuilderError> {
+        self.ram_type_write_with_rw_selectors(
+            name_fn,
+            RAMType::Memory,
+            memory_addr.clone(),
+            prev_ts,
+            ts,
+            prev_values,
+            value,
+            r_selector,
+            w_selector,
+        )
+    }
 }

ceno_zkvm/src/instructions.rs

@@ -6,6 +6,7 @@ use ceno_emul::StepRecord;
 use ff_ext::ExtensionField;
 use gkr_iop::{
     chip::Chip,
+    default_out_eval_groups,
     gkr::{GKRCircuit, layer::Layer},
     selector::SelectorType,
     utils::lk_multiplicity::Multiplicity,
@@ -44,35 +45,16 @@ pub trait Instruction<E: ExtensionField> {
         param: &ProgramParams,
     ) -> Result<(Self::InstructionConfig, GKRCircuit<E>), ZKVMError> {
         let config = Self::construct_circuit(cb, param)?;
-        let w_len = cb.cs.w_expressions.len();
-        let r_len = cb.cs.r_expressions.len();
-        let lk_len = cb.cs.lk_expressions.len();
-        let zero_len =
-            cb.cs.assert_zero_expressions.len() + cb.cs.assert_zero_sumcheck_expressions.len();
 
         let selector = cb.create_placeholder_structural_witin(|| "selector");
         let selector_type = SelectorType::Prefix(selector.expr());
+        cb.cs.set_default_read_selector(selector_type.clone());
+        cb.cs.set_default_write_selector(selector_type.clone());
+        cb.cs.set_default_lookup_selector(selector_type.clone());
+        cb.cs.set_default_zero_selector(selector_type.clone());
 
-        // all shared the same selector
-        let (out_evals, mut chip) = (
-            [
-                // r_record
-                (0..r_len).collect_vec(),
-                // w_record
-                (r_len..r_len + w_len).collect_vec(),
-                // lk_record
-                (r_len + w_len..r_len + w_len + lk_len).collect_vec(),
-                // zero_record
-                (0..zero_len).collect_vec(),
-            ],
-            Chip::new_from_cb(cb, 0),
-        );
-
-        // register selector to legacy constrain system
-        cb.cs.r_selector = Some(selector_type.clone());
-        cb.cs.w_selector = Some(selector_type.clone());
-        cb.cs.lk_selector = Some(selector_type.clone());
-        cb.cs.zero_selector = Some(selector_type.clone());
+        let out_evals = default_out_eval_groups(cb);
+        let mut chip = Chip::new_from_cb(cb, 0);
 
         let layer = Layer::from_circuit_builder(cb, format!("{}_main", Self::name()), 0, out_evals);
         chip.add_layer(layer);

ceno_zkvm/src/instructions/riscv/arith.rs

@@ -218,6 +218,6 @@ mod test {
         )
         .unwrap();
 
-        MockProver::assert_satisfied_raw(&cb, raw_witin, &[insn_code], None, Some(lkm));
+        MockProver::assert_satisfied_raw(&cb, raw_witin, &[insn_code], 1, None, Some(lkm));
     }
 }

ceno_zkvm/src/instructions/riscv/arith_imm.rs

@@ -89,6 +89,6 @@ mod test {
         )
         .unwrap();
 
-        MockProver::assert_satisfied_raw(&cb, raw_witin, &[insn_code], None, Some(lkm));
+        MockProver::assert_satisfied_raw(&cb, raw_witin, &[insn_code], 1, None, Some(lkm));
     }
 }

ceno_zkvm/src/instructions/riscv/auipc.rs

@@ -272,6 +272,6 @@ mod tests {
         )
         .unwrap();
 
-        MockProver::assert_satisfied_raw(&cb, raw_witin, &[insn_code], None, Some(lkm));
+        MockProver::assert_satisfied_raw(&cb, raw_witin, &[insn_code], 1, None, Some(lkm));
     }
 }

ceno_zkvm/src/instructions/riscv/branch/test.rs

@@ -52,7 +52,7 @@ fn impl_opcode_beq(take_branch: bool, a: u32, b: u32) {
     )
     .unwrap();
 
-    MockProver::assert_satisfied_raw(&cb, raw_witin, &[insn_code], None, Some(lkm));
+    MockProver::assert_satisfied_raw(&cb, raw_witin, &[insn_code], 1, None, Some(lkm));
 }
 
 #[test]
@@ -94,7 +94,7 @@ fn impl_opcode_bne(take_branch: bool, a: u32, b: u32) {
     )
     .unwrap();
 
-    MockProver::assert_satisfied_raw(&cb, raw_witin, &[insn_code], None, Some(lkm));
+    MockProver::assert_satisfied_raw(&cb, raw_witin, &[insn_code], 1, None, Some(lkm));
 }
 
 #[test]
@@ -138,7 +138,14 @@ fn impl_bltu_circuit(taken: bool, a: u32, b: u32) -> Result<(), ZKVMError> {
     )
     .unwrap();
 
-    MockProver::assert_satisfied_raw(&circuit_builder, raw_witin, &[insn_code], None, Some(lkm));
+    MockProver::assert_satisfied_raw(
+        &circuit_builder,
+        raw_witin,
+        &[insn_code],
+        1,
+        None,
+        Some(lkm),
+    );
     Ok(())
 }
 
@@ -183,7 +190,14 @@ fn impl_bgeu_circuit(taken: bool, a: u32, b: u32) -> Result<(), ZKVMError> {
     )
     .unwrap();
 
-    MockProver::assert_satisfied_raw(&circuit_builder, raw_witin, &[insn_code], None, Some(lkm));
+    MockProver::assert_satisfied_raw(
+        &circuit_builder,
+        raw_witin,
+        &[insn_code],
+        1,
+        None,
+        Some(lkm),
+    );
     Ok(())
 }
 
@@ -235,7 +249,14 @@ fn impl_blt_circuit<E: ExtensionField>(taken: bool, a: i32, b: i32) -> Result<()
     )
     .unwrap();
 
-    MockProver::assert_satisfied_raw(&circuit_builder, raw_witin, &[insn_code], None, Some(lkm));
+    MockProver::assert_satisfied_raw(
+        &circuit_builder,
+        raw_witin,
+        &[insn_code],
+        1,
+        None,
+        Some(lkm),
+    );
     Ok(())
 }
 
@@ -287,6 +308,13 @@ fn impl_bge_circuit<E: ExtensionField>(taken: bool, a: i32, b: i32) -> Result<()
     )
     .unwrap();
 
-    MockProver::assert_satisfied_raw(&circuit_builder, raw_witin, &[insn_code], None, Some(lkm));
+    MockProver::assert_satisfied_raw(
+        &circuit_builder,
+        raw_witin,
+        &[insn_code],
+        1,
+        None,
+        Some(lkm),
+    );
     Ok(())
 }

ceno_zkvm/src/instructions/riscv/div.rs

@@ -180,7 +180,7 @@ mod test {
             .expect("instruction must declare at least one InsnKind");
         let insn_code = encode_rv32(insn_kind, 2, 3, 4, 0);
         // values assignment
-        let ([raw_witin, _], lkm) = Insn::assign_instances(
+        let ([raw_witin, raw_structural_witin], lkm) = Insn::assign_instances(
             &config,
             &mut ShardContext::default(),
             cb.cs.num_witin as usize,
@@ -215,9 +215,14 @@ mod test {
                 .into_iter()
                 .map(|v| v.into())
                 .collect_vec(),
-            &[],
+            &raw_structural_witin
+                .to_mles()
+                .into_iter()
+                .map(|v| v.into())
+                .collect_vec(),
             &[insn_code],
             expected_errors,
+            1,
             None,
             Some(lkm),
         );

ceno_zkvm/src/instructions/riscv/dummy/test.rs

@@ -28,5 +28,5 @@ fn test_large_ecall_dummy_keccak() {
     )
     .unwrap();
 
-    MockProver::assert_satisfied_raw(&cb, raw_witin, &program, None, Some(lkm));
+    MockProver::assert_satisfied_raw(&cb, raw_witin, &program, 1, None, Some(lkm));
 }

ceno_zkvm/src/instructions/riscv/ecall/fptower_fp.rs

@@ -7,8 +7,7 @@ use ceno_emul::{
 use ff_ext::ExtensionField;
 use generic_array::typenum::Unsigned;
 use gkr_iop::{
-    ProtocolBuilder, ProtocolWitnessGenerator,
-    gkr::{GKRCircuit, layer::Layer},
+    ProtocolBuilder, ProtocolWitnessGenerator, gkr::GKRCircuit,
     utils::lk_multiplicity::Multiplicity,
 };
 use itertools::{Itertools, izip};
@@ -250,7 +249,7 @@ fn build_fp_op_circuit<E: ExtensionField, P: FpOpField + NumWords>(
         0.into(),
     ))?;
 
-    let mut layout = <FpOpLayout<E, P> as ProtocolBuilder<E>>::build_layer_logic(cb, ())?;
+    let layout = <FpOpLayout<E, P> as ProtocolBuilder<E>>::build_layer_logic(cb, ())?;
 
     let mut mem_rw = izip!(&layout.input32_exprs[0], &layout.output32_exprs)
         .enumerate()
@@ -287,10 +286,7 @@ fn build_fp_op_circuit<E: ExtensionField, P: FpOpField + NumWords>(
             .collect::<Result<Vec<WriteMEM>, _>>()?,
     );
 
-    let (out_evals, mut chip) = layout.finalize(cb);
-    let layer =
-        Layer::from_circuit_builder(cb, layer_name.to_string(), layout.n_challenges, out_evals);
-    chip.add_layer(layer);
+    let chip = layout.finalize(layer_name.to_string(), cb);
 
     Ok((
         EcallFpOpConfig {

ceno_zkvm/src/instructions/riscv/ecall/fptower_fp2_add.rs

@@ -6,8 +6,7 @@ use ceno_emul::{
 use ff_ext::ExtensionField;
 use generic_array::typenum::Unsigned;
 use gkr_iop::{
-    ProtocolBuilder, ProtocolWitnessGenerator,
-    gkr::{GKRCircuit, layer::Layer},
+    ProtocolBuilder, ProtocolWitnessGenerator, gkr::GKRCircuit,
     utils::lk_multiplicity::Multiplicity,
 };
 use itertools::{Itertools, izip};
@@ -163,8 +162,7 @@ fn build_fp2_add_circuit<E: ExtensionField, P: FpOpField + Fp2AddSpec + NumWords
         0.into(),
     ))?;
 
-    let mut layout =
-        <Fp2AddSubAssignLayout<E, P> as ProtocolBuilder<E>>::build_layer_logic(cb, ())?;
+    let layout = <Fp2AddSubAssignLayout<E, P> as ProtocolBuilder<E>>::build_layer_logic(cb, ())?;
 
     let mut mem_rw = izip!(&layout.input32_exprs[0], &layout.output32_exprs)
         .enumerate()
@@ -201,10 +199,7 @@ fn build_fp2_add_circuit<E: ExtensionField, P: FpOpField + Fp2AddSpec + NumWords
             .collect::<Result<Vec<WriteMEM>, _>>()?,
     );
 
-    let (out_evals, mut chip) = layout.finalize(cb);
-    let layer =
-        Layer::from_circuit_builder(cb, "fp2_add".to_string(), layout.n_challenges, out_evals);
-    chip.add_layer(layer);
+    let chip = layout.finalize("fp2_add".to_string(), cb);
 
     Ok((
         EcallFp2AddConfig {

ceno_zkvm/src/instructions/riscv/ecall/fptower_fp2_mul.rs

@@ -6,8 +6,7 @@ use ceno_emul::{
 use ff_ext::ExtensionField;
 use generic_array::typenum::Unsigned;
 use gkr_iop::{
-    ProtocolBuilder, ProtocolWitnessGenerator,
-    gkr::{GKRCircuit, layer::Layer},
+    ProtocolBuilder, ProtocolWitnessGenerator, gkr::GKRCircuit,
     utils::lk_multiplicity::Multiplicity,
 };
 use itertools::{Itertools, izip};
@@ -162,7 +161,7 @@ fn build_fp2_mul_circuit<E: ExtensionField, P: FpOpField + Fp2MulSpec + NumWords
         0.into(),
     ))?;
 
-    let mut layout = <Fp2MulAssignLayout<E, P> as ProtocolBuilder<E>>::build_layer_logic(cb, ())?;
+    let layout = <Fp2MulAssignLayout<E, P> as ProtocolBuilder<E>>::build_layer_logic(cb, ())?;
 
     let mut mem_rw = izip!(&layout.input32_exprs[0], &layout.output32_exprs)
         .enumerate()
@@ -199,10 +198,7 @@ fn build_fp2_mul_circuit<E: ExtensionField, P: FpOpField + Fp2MulSpec + NumWords
             .collect::<Result<Vec<WriteMEM>, _>>()?,
     );
 
-    let (out_evals, mut chip) = layout.finalize(cb);
-    let layer =
-        Layer::from_circuit_builder(cb, "fp2_mul".to_string(), layout.n_challenges, out_evals);
-    chip.add_layer(layer);
+    let chip = layout.finalize("fp2_mul".to_string(), cb);
 
     Ok((
         EcallFp2MulConfig {