fix/circuit-order: Reliable order of circuits in proofs (#522)

_Extracted from #368_

---------

Co-authored-by: Aurélien Nicolas <[email protected]>
Co-authored-by: Matthias Görgens <[email protected]>

ceno_zkvm/src/scheme/prover.rs

@@ -86,9 +86,9 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMProver<E, PCS> {
         // commit to main traces
         let mut commitments = BTreeMap::new();
         let mut wits = BTreeMap::new();
-        // sort by circuit name, and we rely on an assumption that
-        // table circuit witnesses come after opcode circuit witnesses
-        for (circuit_name, witness) in witnesses.witnesses {
+
+        // commit to opcode circuits first and then commit to table circuits, sorted by name
+        for (circuit_name, witness) in witnesses.into_iter_sorted() {
             let commit_dur = std::time::Instant::now();
             let num_instances = witness.num_instances();
             let witness = match num_instances {

ceno_zkvm/src/scheme/tests.rs

@@ -119,7 +119,12 @@ fn test_rw_lk_expression_combination() {
         // get proof
         let prover = ZKVMProver::new(pk);
         let mut transcript = Transcript::new(b"test");
-        let wits_in = zkvm_witness.witnesses.remove(&name).unwrap().into_mles();
+        let wits_in = zkvm_witness
+            .into_iter_sorted()
+            .next()
+            .unwrap()
+            .1
+            .into_mles();
         // commit to main traces
         let commit = Pcs::batch_commit_and_write(&prover.pk.pp, &wits_in, &mut transcript).unwrap();
         let wits_in = wits_in.into_iter().map(|v| v.into()).collect_vec();

ceno_zkvm/src/structs.rs

@@ -9,7 +9,7 @@ use crate::{
 };
 use ceno_emul::StepRecord;
 use ff_ext::ExtensionField;
-use itertools::Itertools;
+use itertools::{Itertools, chain};
 use mpcs::PolynomialCommitmentScheme;
 use multilinear_extensions::{
     mle::DenseMultilinearExtension, virtual_poly_v2::ArcMultilinearExtension,
@@ -205,7 +205,8 @@ impl<E: ExtensionField> ZKVMFixedTraces<E> {
 
 #[derive(Default)]
 pub struct ZKVMWitnesses<E: ExtensionField> {
-    pub witnesses: BTreeMap<String, RowMajorMatrix<E::BaseField>>,
+    witnesses_opcodes: BTreeMap<String, RowMajorMatrix<E::BaseField>>,
+    witnesses_tables: BTreeMap<String, RowMajorMatrix<E::BaseField>>,
     lk_mlts: BTreeMap<String, LkMultiplicity>,
     combined_lk_mlt: Option<Vec<HashMap<u64, usize>>>,
 }
@@ -222,7 +223,8 @@ impl<E: ExtensionField> ZKVMWitnesses<E> {
         let cs = cs.get_cs(&OC::name()).unwrap();
         let (witness, logup_multiplicity) =
             OC::assign_instances(config, cs.num_witin as usize, records)?;
-        assert!(self.witnesses.insert(OC::name(), witness).is_none());
+        assert!(self.witnesses_opcodes.insert(OC::name(), witness).is_none());
+        assert!(!self.witnesses_tables.contains_key(&OC::name()));
         assert!(
             self.lk_mlts
                 .insert(OC::name(), logup_multiplicity)
@@ -273,10 +275,16 @@ impl<E: ExtensionField> ZKVMWitnesses<E> {
             self.combined_lk_mlt.as_ref().unwrap(),
             input,
         )?;
-        assert!(self.witnesses.insert(TC::name(), witness).is_none());
+        assert!(self.witnesses_tables.insert(TC::name(), witness).is_none());
+        assert!(!self.witnesses_opcodes.contains_key(&TC::name()));
 
         Ok(())
     }
+
+    /// Iterate opcode circuits, then table circuits, sorted by name.
+    pub fn into_iter_sorted(self) -> impl Iterator<Item = (String, RowMajorMatrix<E::BaseField>)> {
+        chain(self.witnesses_opcodes, self.witnesses_tables)
+    }
 }
 
 #[derive(Debug)]