refactor(zkvm,gkr): simplify chip finalize flow

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};
@@ -291,9 +290,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(), 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};
@@ -209,9 +208,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(), 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};
@@ -207,9 +206,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(), out_evals);
-    chip.add_layer(layer);
+    let chip = layout.finalize("fp2_mul".to_string(), cb);
 
     Ok((
         EcallFp2MulConfig {

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

@@ -7,7 +7,7 @@ use ceno_emul::{
 use ff_ext::ExtensionField;
 use gkr_iop::{
     ProtocolBuilder, ProtocolWitnessGenerator,
-    gkr::{GKRCircuit, booleanhypercube::BooleanHypercube, layer::Layer},
+    gkr::{GKRCircuit, booleanhypercube::BooleanHypercube},
     utils::lk_multiplicity::Multiplicity,
 };
 use itertools::{Itertools, izip};
@@ -132,10 +132,7 @@ impl<E: ExtensionField> Instruction<E> for KeccakInstruction<E> {
             })
             .collect::<Result<Vec<WriteMEM>, _>>()?;
 
-        let (out_evals, mut chip) = layout.finalize(cb);
-
-        let layer = Layer::from_circuit_builder(cb, Self::name(), out_evals);
-        chip.add_layer(layer);
+        let chip = layout.finalize(Self::name(), cb);
 
         let circuit = chip.gkr_circuit();
 

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

@@ -5,8 +5,7 @@ use ceno_emul::{
 };
 use ff_ext::{ExtensionField, FieldInto};
 use gkr_iop::{
-    ProtocolBuilder, ProtocolWitnessGenerator,
-    gkr::{GKRCircuit, layer::Layer},
+    ProtocolBuilder, ProtocolWitnessGenerator, gkr::GKRCircuit,
     utils::lk_multiplicity::Multiplicity,
 };
 use itertools::{Itertools, izip};
@@ -130,10 +129,7 @@ impl<E: ExtensionField> Instruction<E> for ShaExtendInstruction<E> {
             vm_state.ts,
         )?);
 
-        let (out_evals, mut chip) = layout.finalize(cb);
-
-        let layer = Layer::from_circuit_builder(cb, Self::name(), out_evals);
-        chip.add_layer(layer);
+        let chip = layout.finalize(Self::name(), cb);
 
         let circuit = chip.gkr_circuit();
 

ceno_zkvm/src/instructions/riscv/ecall/uint256.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, chain, izip};
@@ -177,10 +176,7 @@ impl<E: ExtensionField> Instruction<E> for Uint256MulInstruction<E> {
             .collect::<Result<Vec<WriteMEM>, _>>()?,
         );
 
-        let (out_evals, mut chip) = layout.finalize(cb);
-
-        let layer = Layer::from_circuit_builder(cb, "uint256_mul".to_string(), out_evals);
-        chip.add_layer(layer);
+        let chip = layout.finalize("uint256_mul".to_string(), cb);
 
         let circuit = chip.gkr_circuit();
 
@@ -501,10 +497,7 @@ impl<E: ExtensionField, Spec: Uint256InvSpec> Instruction<E> for Uint256InvInstr
             })
             .collect::<Result<Vec<WriteMEM>, _>>()?;
 
-        let (out_evals, mut chip) = layout.finalize(cb);
-
-        let layer = Layer::from_circuit_builder(cb, Spec::name(), out_evals);
-        chip.add_layer(layer);
+        let chip = layout.finalize(Spec::name(), cb);
 
         let circuit = chip.gkr_circuit();
 

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

@@ -7,8 +7,7 @@ use ceno_emul::{
 use ff_ext::ExtensionField;
 use generic_array::{GenericArray, typenum::Unsigned};
 use gkr_iop::{
-    ProtocolBuilder, ProtocolWitnessGenerator,
-    gkr::{GKRCircuit, layer::Layer},
+    ProtocolBuilder, ProtocolWitnessGenerator, gkr::GKRCircuit,
     utils::lk_multiplicity::Multiplicity,
 };
 use itertools::{Itertools, izip};
@@ -177,10 +176,7 @@ impl<E: ExtensionField, EC: EllipticCurve> Instruction<E>
                 .collect::<Result<Vec<WriteMEM>, _>>()?,
         );
 
-        let (out_evals, mut chip) = layout.finalize(cb);
-
-        let layer = Layer::from_circuit_builder(cb, "weierstrass_add".to_string(), out_evals);
-        chip.add_layer(layer);
+        let chip = layout.finalize("weierstrass_add".to_string(), cb);
 
         let circuit = chip.gkr_circuit();
 

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

@@ -7,8 +7,7 @@ use ceno_emul::{
 use ff_ext::ExtensionField;
 use generic_array::{GenericArray, typenum::Unsigned};
 use gkr_iop::{
-    ProtocolBuilder, ProtocolWitnessGenerator,
-    gkr::{GKRCircuit, layer::Layer},
+    ProtocolBuilder, ProtocolWitnessGenerator, gkr::GKRCircuit,
     utils::lk_multiplicity::Multiplicity,
 };
 use itertools::{Itertools, izip};
@@ -178,11 +177,7 @@ impl<E: ExtensionField, EC: EllipticCurve + WeierstrassParameters> Instruction<E
             .collect::<Result<Vec<WriteMEM>, _>>()?,
         );
 
-        let (out_evals, mut chip) = layout.finalize(cb);
-
-        let layer =
-            Layer::from_circuit_builder(cb, "weierstrass_decompress".to_string(), out_evals);
-        chip.add_layer(layer);
+        let chip = layout.finalize("weierstrass_decompress".to_string(), cb);
 
         let circuit = chip.gkr_circuit();
 

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

@@ -7,8 +7,7 @@ use ceno_emul::{
 use ff_ext::ExtensionField;
 use generic_array::{GenericArray, typenum::Unsigned};
 use gkr_iop::{
-    ProtocolBuilder, ProtocolWitnessGenerator,
-    gkr::{GKRCircuit, layer::Layer},
+    ProtocolBuilder, ProtocolWitnessGenerator, gkr::GKRCircuit,
     utils::lk_multiplicity::Multiplicity,
 };
 use itertools::{Itertools, izip};
@@ -150,10 +149,7 @@ impl<E: ExtensionField, EC: EllipticCurve + WeierstrassParameters> Instruction<E
             })
             .collect::<Result<Vec<WriteMEM>, _>>()?;
 
-        let (out_evals, mut chip) = layout.finalize(cb);
-
-        let layer = Layer::from_circuit_builder(cb, "weierstrass_double".to_string(), out_evals);
-        chip.add_layer(layer);
+        let chip = layout.finalize("weierstrass_double".to_string(), cb);
 
         let circuit = chip.gkr_circuit();
 

ceno_zkvm/src/precompiles/bitwise_keccakf.rs

@@ -18,7 +18,7 @@ use transcript::{BasicTranscript, Transcript};
 use witness::{InstancePaddingStrategy, RowMajorMatrix};
 
 use gkr_iop::{
-    OutEvalGroups, ProtocolBuilder, ProtocolWitnessGenerator,
+    ProtocolBuilder, ProtocolWitnessGenerator,
     chip::Chip,
     circuit_builder::{CircuitBuilder, ConstraintSystem},
     cpu::{CpuBackend, CpuProver},
@@ -776,7 +776,7 @@ impl<E: ExtensionField> KeccakLayout<E> {
 impl<E: ExtensionField> ProtocolBuilder<E> for KeccakLayout<E> {
     type Params = KeccakParams;
 
-    fn finalize(&mut self, _cb: &mut CircuitBuilder<E>) -> (OutEvalGroups, Chip<E>) {
+    fn finalize(&mut self, _name: String, _cb: &mut CircuitBuilder<E>) -> Chip<E> {
         unimplemented!()
     }
 

ceno_zkvm/src/precompiles/fptower/fp.rs

@@ -28,8 +28,8 @@ use derive::AlignedBorrow;
 use ff_ext::ExtensionField;
 use generic_array::{GenericArray, sequence::GenericSequence};
 use gkr_iop::{
-    OutEvalGroups, ProtocolBuilder, ProtocolWitnessGenerator, chip::Chip,
-    circuit_builder::CircuitBuilder, error::CircuitBuilderError, selector::SelectorType,
+    ProtocolBuilder, ProtocolWitnessGenerator, chip::Chip, circuit_builder::CircuitBuilder,
+    default_out_eval_groups, error::CircuitBuilderError, gkr::layer::Layer, selector::SelectorType,
 };
 use itertools::Itertools;
 use multilinear_extensions::{Expression, ToExpr, WitIn, util::max_usable_threads};
@@ -224,7 +224,7 @@ impl<E: ExtensionField, P: FpOpField> ProtocolBuilder<E> for FpOpLayout<E, P> {
         Ok(layout)
     }
 
-    fn finalize(&mut self, cb: &mut CircuitBuilder<E>) -> (OutEvalGroups, Chip<E>) {
+    fn finalize(&mut self, name: String, cb: &mut CircuitBuilder<E>) -> Chip<E> {
         self.n_fixed = cb.cs.num_fixed;
         self.n_committed = cb.cs.num_witin as usize;
         self.n_structural_witin = cb.cs.num_structural_witin as usize;
@@ -234,20 +234,11 @@ impl<E: ExtensionField, P: FpOpField> ProtocolBuilder<E> for FpOpLayout<E, P> {
         cb.cs.lk_selector = Some(self.selector_type_layout.sel_all.clone());
         cb.cs.zero_selector = Some(self.selector_type_layout.sel_all.clone());
 
-        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();
-        (
-            [
-                (0..r_len).collect_vec(),
-                (r_len..r_len + w_len).collect_vec(),
-                (r_len + w_len..r_len + w_len + lk_len).collect_vec(),
-                (0..zero_len).collect_vec(),
-            ],
-            Chip::new_from_cb(cb),
-        )
+        let out_evals = default_out_eval_groups(cb);
+        let mut chip = Chip::new_from_cb(cb);
+        let layer = Layer::from_circuit_builder(cb, name, out_evals);
+        chip.add_layer(layer);
+        chip
     }
 }
 
@@ -299,7 +290,7 @@ mod tests {
     use gkr_iop::{
         circuit_builder::{CircuitBuilder, ConstraintSystem},
         cpu::{CpuBackend, CpuProver},
-        gkr::{GKRProverOutput, layer::Layer},
+        gkr::GKRProverOutput,
         selector::SelectorContext,
     };
     use itertools::Itertools;
@@ -330,9 +321,7 @@ mod tests {
         let mut cb = CircuitBuilder::<E>::new(&mut cs);
         let mut layout =
             FpOpLayout::<E, P>::build_layer_logic(&mut cb, ()).expect("build_layer_logic failed");
-        let (out_evals, mut chip) = layout.finalize(&mut cb);
-        let layer = Layer::from_circuit_builder(&cb, "fp_op".to_string(), out_evals);
-        chip.add_layer(layer);
+        let chip = layout.finalize("fp_op".to_string(), &mut cb);
         let gkr_circuit = chip.gkr_circuit();
 
         let instances = (0..count)

ceno_zkvm/src/precompiles/fptower/fp2_addsub.rs

@@ -28,8 +28,8 @@ use derive::AlignedBorrow;
 use ff_ext::ExtensionField;
 use generic_array::{GenericArray, sequence::GenericSequence};
 use gkr_iop::{
-    OutEvalGroups, ProtocolBuilder, ProtocolWitnessGenerator, chip::Chip,
-    circuit_builder::CircuitBuilder, error::CircuitBuilderError, selector::SelectorType,
+    ProtocolBuilder, ProtocolWitnessGenerator, chip::Chip, circuit_builder::CircuitBuilder,
+    default_out_eval_groups, error::CircuitBuilderError, gkr::layer::Layer, selector::SelectorType,
 };
 use itertools::Itertools;
 use multilinear_extensions::{Expression, ToExpr, WitIn, util::max_usable_threads};
@@ -248,7 +248,7 @@ impl<E: ExtensionField, P: FpOpField> ProtocolBuilder<E> for Fp2AddSubAssignLayo
         Ok(layout)
     }
 
-    fn finalize(&mut self, cb: &mut CircuitBuilder<E>) -> (OutEvalGroups, Chip<E>) {
+    fn finalize(&mut self, name: String, cb: &mut CircuitBuilder<E>) -> Chip<E> {
         self.n_fixed = cb.cs.num_fixed;
         self.n_committed = cb.cs.num_witin as usize;
         self.n_structural_witin = cb.cs.num_structural_witin as usize;
@@ -258,20 +258,11 @@ impl<E: ExtensionField, P: FpOpField> ProtocolBuilder<E> for Fp2AddSubAssignLayo
         cb.cs.lk_selector = Some(self.selector_type_layout.sel_all.clone());
         cb.cs.zero_selector = Some(self.selector_type_layout.sel_all.clone());
 
-        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();
-        (
-            [
-                (0..r_len).collect_vec(),
-                (r_len..r_len + w_len).collect_vec(),
-                (r_len + w_len..r_len + w_len + lk_len).collect_vec(),
-                (0..zero_len).collect_vec(),
-            ],
-            Chip::new_from_cb(cb),
-        )
+        let out_evals = default_out_eval_groups(cb);
+        let mut chip = Chip::new_from_cb(cb);
+        let layer = Layer::from_circuit_builder(cb, name, out_evals);
+        chip.add_layer(layer);
+        chip
     }
 }
 
@@ -326,7 +317,7 @@ mod tests {
     use gkr_iop::{
         circuit_builder::{CircuitBuilder, ConstraintSystem},
         cpu::{CpuBackend, CpuProver},
-        gkr::{GKRProverOutput, layer::Layer},
+        gkr::GKRProverOutput,
         selector::SelectorContext,
     };
     use itertools::Itertools;
@@ -357,9 +348,7 @@ mod tests {
         let mut cb = CircuitBuilder::<E>::new(&mut cs);
         let mut layout = Fp2AddSubAssignLayout::<E, P>::build_layer_logic(&mut cb, ())
             .expect("build_layer_logic failed");
-        let (out_evals, mut chip) = layout.finalize(&mut cb);
-        let layer = Layer::from_circuit_builder(&cb, "fp2_addsub".to_string(), out_evals);
-        chip.add_layer(layer);
+        let chip = layout.finalize("fp2_addsub".to_string(), &mut cb);
         let gkr_circuit = chip.gkr_circuit();
 
         let instances = (0..count)