Use a sane type for Relation in the MIPS constraints

o1vm/src/interpreters/keccak/column.rs

@@ -1,7 +1,10 @@
 //! This module defines the custom columns used in the Keccak witness, which
 //! are aliases for the actual Keccak witness columns also defined here.
 use self::{Absorbs::*, Sponges::*, Steps::*};
-use crate::interpreters::keccak::{ZKVM_KECCAK_COLS_CURR, ZKVM_KECCAK_COLS_NEXT};
+use crate::{
+    interpreters::keccak::{ZKVM_KECCAK_COLS_CURR, ZKVM_KECCAK_COLS_NEXT},
+    RelationColumnType,
+};
 use kimchi::circuits::polynomials::keccak::constants::{
     CHI_SHIFTS_B_LEN, CHI_SHIFTS_B_OFF, CHI_SHIFTS_SUM_LEN, CHI_SHIFTS_SUM_OFF, PIRHO_DENSE_E_LEN,
     PIRHO_DENSE_E_OFF, PIRHO_DENSE_ROT_E_LEN, PIRHO_DENSE_ROT_E_OFF, PIRHO_EXPAND_ROT_E_LEN,
@@ -373,10 +376,10 @@ impl<T: Clone> IndexMut<ColumnAlias> for KeccakWitness<T> {
     }
 }
 
-impl ColumnIndexer<usize> for ColumnAlias {
+impl ColumnIndexer<RelationColumnType> for ColumnAlias {
     const N_COL: usize = N_ZKVM_KECCAK_REL_COLS + N_ZKVM_KECCAK_SEL_COLS;
-    fn to_column(self) -> Column<usize> {
-        Column::Relation(usize::from(self))
+    fn to_column(self) -> Column<RelationColumnType> {
+        Column::Relation(RelationColumnType::Scratch(usize::from(self)))
     }
 }
 

o1vm/src/interpreters/mips/column.rs

@@ -1,4 +1,7 @@
-use crate::interpreters::mips::Instruction::{self, IType, JType, RType};
+use crate::{
+    interpreters::mips::Instruction::{self, IType, JType, RType},
+    RelationColumnType,
+};
 use kimchi_msm::{
     columns::{Column, ColumnIndexer},
     witness::Witness,
@@ -143,10 +146,10 @@ impl<T: Clone> IndexMut<ColumnAlias> for MIPSWitness<T> {
     }
 }
 
-impl ColumnIndexer<usize> for ColumnAlias {
+impl ColumnIndexer<RelationColumnType> for ColumnAlias {
     const N_COL: usize = N_MIPS_COLS;
 
-    fn to_column(self) -> Column<usize> {
+    fn to_column(self) -> Column<RelationColumnType> {
         match self {
             Self::ScratchState(ss) => {
                 assert!(
@@ -155,7 +158,7 @@ impl ColumnIndexer<usize> for ColumnAlias {
                     SCRATCH_SIZE,
                     ss
                 );
-                Column::Relation(ss)
+                Column::Relation(RelationColumnType::Scratch(ss))
             }
             Self::ScratchStateInverse(ss) => {
                 assert!(
@@ -164,9 +167,9 @@ impl ColumnIndexer<usize> for ColumnAlias {
                     SCRATCH_SIZE_INVERSE,
                     ss
                 );
-                Column::Relation(SCRATCH_SIZE + ss)
+                Column::Relation(RelationColumnType::ScratchInverse(ss))
             }
-            Self::InstructionCounter => Column::Relation(SCRATCH_SIZE + SCRATCH_SIZE_INVERSE),
+            Self::InstructionCounter => Column::Relation(RelationColumnType::InstructionCounter),
             // TODO: what happens with error? It does not have a corresponding alias
             Self::Selector(s) => {
                 assert!(

o1vm/src/lib.rs

@@ -24,6 +24,7 @@ pub mod utils;
 
 pub mod test_preimage_read;
 
+use crate::pickles::column_env::RelationColumnType;
 use kimchi::circuits::{
     berkeley_columns::BerkeleyChallengeTerm,
     expr::{ConstantExpr, Expr},
@@ -43,4 +44,4 @@ pub use ramlookup::{LookupMode as RAMLookupMode, RAMLookup};
 /// `P(X, Y, Z) = q_x X + q_y Y + q_m X Y + q_o Z + q_c`
 /// To represent this multi-variate polynomial using the expression framework,
 /// we would use 3 different columns.
-pub(crate) type E<F> = Expr<ConstantExpr<F, BerkeleyChallengeTerm>, Column<usize>>;
+pub(crate) type E<F> = Expr<ConstantExpr<F, BerkeleyChallengeTerm>, Column<RelationColumnType>>;

o1vm/src/pickles/column_env.rs

@@ -14,6 +14,14 @@ use kimchi::circuits::{
 
 type Evals<F> = Evaluations<F, Radix2EvaluationDomain<F>>;
 
+#[derive(PartialEq, Eq, Clone, Copy, Debug, Hash)]
+pub enum RelationColumnType {
+    Scratch(usize),
+    ScratchInverse(usize),
+    InstructionCounter,
+    Error,
+}
+
 /// The collection of polynomials (all in evaluation form) and constants
 /// required to evaluate an expression as a polynomial.
 ///
@@ -35,39 +43,33 @@ pub struct ColumnEnvironment<'a, F: FftField> {
     pub domain: EvaluationDomains<F>,
 }
 
-pub fn get_all_columns() -> Vec<Column<usize>> {
-    let mut cols = Vec::<Column<usize>>::with_capacity(
+pub fn get_all_columns() -> Vec<Column<RelationColumnType>> {
+    let mut cols = Vec::<Column<RelationColumnType>>::with_capacity(
         SCRATCH_SIZE + SCRATCH_SIZE_INVERSE + 2 + N_MIPS_SEL_COLS,
     );
-    for i in 0..SCRATCH_SIZE + SCRATCH_SIZE_INVERSE + 2 {
-        cols.push(Column::Relation(i));
+    for i in 0..SCRATCH_SIZE {
+        cols.push(Column::Relation(RelationColumnType::Scratch(i)));
     }
+    for i in 0..SCRATCH_SIZE_INVERSE {
+        cols.push(Column::Relation(RelationColumnType::ScratchInverse(i)));
+    }
+    cols.push(Column::Relation(RelationColumnType::InstructionCounter));
+    cols.push(Column::Relation(RelationColumnType::Error));
     for i in 0..N_MIPS_SEL_COLS {
         cols.push(Column::DynamicSelector(i));
     }
     cols
 }
 
 impl<G> WitnessColumns<G, [G; N_MIPS_SEL_COLS]> {
-    pub fn get_column(&self, col: &Column<usize>) -> Option<&G> {
+    pub fn get_column(&self, col: &Column<RelationColumnType>) -> Option<&G> {
         match *col {
-            Column::Relation(i) => {
-                if i < SCRATCH_SIZE {
-                    let res = &self.scratch[i];
-                    Some(res)
-                } else if i < SCRATCH_SIZE + SCRATCH_SIZE_INVERSE {
-                    let res = &self.scratch_inverse[i - SCRATCH_SIZE];
-                    Some(res)
-                } else if i == SCRATCH_SIZE + SCRATCH_SIZE_INVERSE {
-                    let res = &self.instruction_counter;
-                    Some(res)
-                } else if i == SCRATCH_SIZE + SCRATCH_SIZE_INVERSE + 1 {
-                    let res = &self.error;
-                    Some(res)
-                } else {
-                    panic!("We should not have that many relation columns. We have {} columns and index {} was given", SCRATCH_SIZE + SCRATCH_SIZE_INVERSE + 2, i);
-                }
-            }
+            Column::Relation(i) => match i {
+                RelationColumnType::Scratch(i) => Some(&self.scratch[i]),
+                RelationColumnType::ScratchInverse(i) => Some(&self.scratch_inverse[i]),
+                RelationColumnType::InstructionCounter => Some(&self.instruction_counter),
+                RelationColumnType::Error => Some(&self.error),
+            },
             Column::DynamicSelector(i) => {
                 assert!(
                     i < N_MIPS_SEL_COLS,
@@ -91,9 +93,7 @@ impl<G> WitnessColumns<G, [G; N_MIPS_SEL_COLS]> {
 impl<'a, F: FftField> TColumnEnvironment<'a, F, BerkeleyChallengeTerm, BerkeleyChallenges<F>>
     for ColumnEnvironment<'a, F>
 {
-    // FIXME: do we change to the MIPS column type?
-    // We do not want to keep kimchi_msm/generic prover
-    type Column = Column<usize>;
+    type Column = Column<RelationColumnType>;
 
     fn get_column(&self, col: &Self::Column) -> Option<&'a Evals<F>> {
         self.witness.get_column(col)

o1vm/src/pickles/tests.rs

@@ -12,11 +12,15 @@ use crate::{
         interpreter::{self, InterpreterEnv},
         Instruction,
     },
-    pickles::{verifier::verify, MAXIMUM_DEGREE_CONSTRAINTS, TOTAL_NUMBER_OF_CONSTRAINTS},
+    pickles::{
+        column_env::RelationColumnType, verifier::verify, MAXIMUM_DEGREE_CONSTRAINTS,
+        TOTAL_NUMBER_OF_CONSTRAINTS,
+    },
+    E,
 };
 use ark_ff::{Field, One, UniformRand, Zero};
 use kimchi::circuits::{domains::EvaluationDomains, expr::Expr, gate::CurrOrNext};
-use kimchi_msm::{columns::Column, expr::E};
+use kimchi_msm::columns::Column;
 use log::debug;
 use mina_curves::pasta::{Fp, Fq, Pallas, PallasParameters};
 use mina_poseidon::{
@@ -78,9 +82,26 @@ fn test_small_circuit() {
         },
     };
     let mut expr = Expr::zero();
-    for i in 0..SCRATCH_SIZE + SCRATCH_SIZE_INVERSE + 2 {
-        expr += Expr::cell(Column::Relation(i), CurrOrNext::Curr);
+    for i in 0..SCRATCH_SIZE {
+        expr += Expr::cell(
+            Column::Relation(RelationColumnType::Scratch(i)),
+            CurrOrNext::Curr,
+        );
+    }
+    for i in 0..SCRATCH_SIZE_INVERSE {
+        expr += Expr::cell(
+            Column::Relation(RelationColumnType::ScratchInverse(i)),
+            CurrOrNext::Curr,
+        );
     }
+    expr += Expr::cell(
+        Column::Relation(RelationColumnType::InstructionCounter),
+        CurrOrNext::Curr,
+    );
+    expr += Expr::cell(
+        Column::Relation(RelationColumnType::Error),
+        CurrOrNext::Curr,
+    );
     let mut rng = make_test_rng(None);
 
     type BaseSponge = DefaultFqSponge<PallasParameters, PlonkSpongeConstantsKimchi>;

o1vm/src/pickles/verifier.rs

@@ -24,7 +24,7 @@ use poly_commitment::{
 };
 
 use super::{
-    column_env::get_all_columns,
+    column_env::{get_all_columns, RelationColumnType},
     proof::{Proof, WitnessColumns},
 };
 use crate::{interpreters::mips::column::N_MIPS_SEL_COLS, E};
@@ -40,7 +40,7 @@ struct ColumnEval<'a, G: AffineRepr> {
 }
 
 impl<G: AffineRepr> ColumnEvaluations<G::ScalarField> for ColumnEval<'_, G> {
-    type Column = Column<usize>;
+    type Column = Column<RelationColumnType>;
     fn evaluate(
         &self,
         col: Self::Column,