Track lookup values in interpreter and proof

ivc/src/expr_eval.rs

@@ -23,10 +23,10 @@ use strum::EnumCount;
 /// Generic structure containing column vectors.
 pub struct GenericVecStructure<G: KimchiCurve>(pub Vec<Vec<G::ScalarField>>);
 
-impl<G: KimchiCurve> Index<GenericColumn> for GenericVecStructure<G> {
+impl<G: KimchiCurve> Index<GenericColumn<usize>> for GenericVecStructure<G> {
     type Output = [G::ScalarField];
 
-    fn index(&self, index: GenericColumn) -> &Self::Output {
+    fn index(&self, index: GenericColumn<usize>) -> &Self::Output {
         match index {
             GenericColumn::FixedSelector(i) => &self.0[i],
             _ => panic!("should not happen"),
@@ -73,7 +73,7 @@ impl<
     }
 
     pub fn process_extended_folding_column<
-        FC: FoldingConfig<Column = GenericColumn, Curve = Curve, Challenge = PlonkishChallenge>,
+        FC: FoldingConfig<Column = GenericColumn<usize>, Curve = Curve, Challenge = PlonkishChallenge>,
     >(
         &self,
         col: &ExtendedFoldingColumn<FC>,
@@ -105,7 +105,7 @@ impl<
     /// Evaluates the expression in the provided side
     pub fn eval_naive_fexpr<
         'a,
-        FC: FoldingConfig<Column = GenericColumn, Curve = Curve, Challenge = PlonkishChallenge>,
+        FC: FoldingConfig<Column = GenericColumn<usize>, Curve = Curve, Challenge = PlonkishChallenge>,
     >(
         &'a self,
         exp: &FoldingExp<FC>,
@@ -136,7 +136,7 @@ impl<
     /// For FoldingCompatibleExp
     pub fn eval_naive_fcompat<
         'a,
-        FC: FoldingConfig<Column = GenericColumn, Curve = Curve, Challenge = PlonkishChallenge>,
+        FC: FoldingConfig<Column = GenericColumn<usize>, Curve = Curve, Challenge = PlonkishChallenge>,
     >(
         &'a self,
         exp: &FoldingCompatibleExpr<FC>,

ivc/src/ivc/columns.rs

@@ -378,7 +378,7 @@ pub enum IVCColumn {
     Block6UOutput,
 }
 
-impl ColumnIndexer for IVCColumn {
+impl ColumnIndexer<usize> for IVCColumn {
     /// Number of columns used by the IVC circuit
     /// It contains at least the columns used for Poseidon.
     /// It does not include the additional columns that might be required
@@ -391,7 +391,7 @@ impl ColumnIndexer for IVCColumn {
     // We also add 1 for the FoldIteration column.
     const N_COL: usize = IVCPoseidonColumn::N_COL + 1 + N_BLOCKS;
 
-    fn to_column(self) -> Column {
+    fn to_column(self) -> Column<usize> {
         match self {
             // We keep a column that will be used for the folding iteration.
             // Question: do we need it for all the rows or does it appear only

ivc/src/ivc/helpers.rs

@@ -10,7 +10,11 @@ use kimchi_msm::{
 use super::{LIMB_BITSIZE_XLARGE, N_LIMBS_XLARGE};
 
 /// Helper. Combines small limbs into big limbs.
-pub fn combine_large_to_xlarge<F: PrimeField, CIx: ColumnIndexer, Env: ColAccessCap<F, CIx>>(
+pub fn combine_large_to_xlarge<
+    F: PrimeField,
+    CIx: ColumnIndexer<usize>,
+    Env: ColAccessCap<F, CIx>,
+>(
     x: [Env::Variable; N_LIMBS_LARGE],
 ) -> [Env::Variable; N_LIMBS_XLARGE] {
     combine_limbs_m_to_n::<
@@ -25,7 +29,11 @@ pub fn combine_large_to_xlarge<F: PrimeField, CIx: ColumnIndexer, Env: ColAccess
 }
 
 /// Helper. Combines 17x15bit limbs into 1 native field element.
-pub fn combine_small_to_full<F: PrimeField, CIx: ColumnIndexer, Env: ColAccessCap<F, CIx>>(
+pub fn combine_small_to_full<
+    F: PrimeField,
+    CIx: ColumnIndexer<usize>,
+    Env: ColAccessCap<F, CIx>,
+>(
     x: [Env::Variable; N_LIMBS_SMALL],
 ) -> Env::Variable {
     let [res] =

ivc/src/ivc/mod.rs

@@ -74,8 +74,8 @@ mod tests {
     type IVCWitnessBuilderEnvRaw<LT> = WitnessBuilderEnv<
         Fp,
         IVCColumn,
-        { <IVCColumn as ColumnIndexer>::N_COL - N_BLOCKS },
-        { <IVCColumn as ColumnIndexer>::N_COL - N_BLOCKS },
+        { <IVCColumn as ColumnIndexer<usize>>::N_COL - N_BLOCKS },
+        { <IVCColumn as ColumnIndexer<usize>>::N_COL - N_BLOCKS },
         0,
         N_FSEL_IVC,
         LT,

ivc/src/plonkish_lang.rs

@@ -75,13 +75,13 @@ impl<
 {
 }
 
-impl<const N_COL: usize, const N_FSEL: usize, F: FftField, Evals: CombinableEvals<F>> Index<Column>
-    for PlonkishWitnessGeneric<N_COL, N_FSEL, F, Evals>
+impl<const N_COL: usize, const N_FSEL: usize, F: FftField, Evals: CombinableEvals<F>>
+    Index<Column<usize>> for PlonkishWitnessGeneric<N_COL, N_FSEL, F, Evals>
 {
     type Output = [F];
 
     /// Map a column alias to the corresponding witness column.
-    fn index(&self, index: Column) -> &Self::Output {
+    fn index(&self, index: Column<usize>) -> &Self::Output {
         match index {
             Column::Relation(i) => self.witness.cols[i].e_as_slice(),
             Column::FixedSelector(i) => self.fixed_selectors[i].e_as_slice(),

ivc/src/poseidon_55_0_7_3_2/columns.rs

@@ -29,7 +29,7 @@ pub enum PoseidonColumn<const STATE_SIZE: usize, const NB_FULL_ROUND: usize> {
     RoundConstant(usize, usize),
 }
 
-impl<const STATE_SIZE: usize, const NB_FULL_ROUND: usize> ColumnIndexer
+impl<const STATE_SIZE: usize, const NB_FULL_ROUND: usize> ColumnIndexer<usize>
     for PoseidonColumn<STATE_SIZE, NB_FULL_ROUND>
 {
     // - STATE_SIZE input columns
@@ -42,7 +42,7 @@ impl<const STATE_SIZE: usize, const NB_FULL_ROUND: usize> ColumnIndexer
     // - STATE_SIZE * NB_FULL_ROUND constants
     const N_COL: usize = STATE_SIZE + 5 * NB_FULL_ROUND * STATE_SIZE;
 
-    fn to_column(self) -> Column {
+    fn to_column(self) -> Column<usize> {
         match self {
             PoseidonColumn::Input(i) => {
                 assert!(i < STATE_SIZE);

ivc/src/poseidon_55_0_7_3_2/mod.rs

@@ -43,8 +43,8 @@ mod tests {
     type PoseidonWitnessBuilderEnv = WitnessBuilderEnv<
         Fp,
         TestPoseidonColumn,
-        { <TestPoseidonColumn as ColumnIndexer>::N_COL },
-        { <TestPoseidonColumn as ColumnIndexer>::N_COL },
+        { <TestPoseidonColumn as ColumnIndexer<usize>>::N_COL },
+        { <TestPoseidonColumn as ColumnIndexer<usize>>::N_COL },
         N_DSEL,
         N_FSEL,
         DummyLookupTable,

ivc/src/poseidon_55_0_7_3_7/columns.rs

@@ -20,12 +20,12 @@ pub enum PoseidonColumn<const STATE_SIZE: usize, const NB_FULL_ROUND: usize> {
     RoundConstant(usize, usize),
 }
 
-impl<const STATE_SIZE: usize, const NB_FULL_ROUND: usize> ColumnIndexer
+impl<const STATE_SIZE: usize, const NB_FULL_ROUND: usize> ColumnIndexer<usize>
     for PoseidonColumn<STATE_SIZE, NB_FULL_ROUND>
 {
     const N_COL: usize = STATE_SIZE + NB_FULL_ROUND * STATE_SIZE;
 
-    fn to_column(self) -> Column {
+    fn to_column(self) -> Column<usize> {
         match self {
             PoseidonColumn::Input(i) => {
                 assert!(i < STATE_SIZE);

ivc/src/poseidon_55_0_7_3_7/mod.rs

@@ -41,8 +41,8 @@ mod tests {
     type PoseidonWitnessBuilderEnv = WitnessBuilderEnv<
         Fp,
         TestPoseidonColumn,
-        { <TestPoseidonColumn as ColumnIndexer>::N_COL },
-        { <TestPoseidonColumn as ColumnIndexer>::N_COL },
+        { <TestPoseidonColumn as ColumnIndexer<usize>>::N_COL },
+        { <TestPoseidonColumn as ColumnIndexer<usize>>::N_COL },
         N_DSEL,
         N_FSEL,
         DummyLookupTable,

ivc/src/poseidon_8_56_5_3_2/columns.rs

@@ -30,7 +30,7 @@ pub enum PoseidonColumn<
 }
 
 impl<const STATE_SIZE: usize, const NB_FULL_ROUND: usize, const NB_PARTIAL_ROUND: usize>
-    ColumnIndexer for PoseidonColumn<STATE_SIZE, NB_FULL_ROUND, NB_PARTIAL_ROUND>
+    ColumnIndexer<usize> for PoseidonColumn<STATE_SIZE, NB_FULL_ROUND, NB_PARTIAL_ROUND>
 {
     // - STATE_SIZE input columns
     // - for each partial round:
@@ -54,7 +54,7 @@ impl<const STATE_SIZE: usize, const NB_FULL_ROUND: usize, const NB_PARTIAL_ROUND
             + (4 + STATE_SIZE - 1) * NB_PARTIAL_ROUND // partial round
             + STATE_SIZE * (NB_PARTIAL_ROUND + NB_FULL_ROUND); // fixed selectors
 
-    fn to_column(self) -> Column {
+    fn to_column(self) -> Column<usize> {
         // number of reductions for
         // x -> x^2 -> x^4 -> x^5 -> x^5 * MDS
         let nb_red = 4;

ivc/src/poseidon_8_56_5_3_2/mod.rs

@@ -32,8 +32,8 @@ mod tests {
     type PoseidonWitnessBuilderEnv = WitnessBuilderEnv<
         Fp,
         Column,
-        { <Column as ColumnIndexer>::N_COL },
-        { <Column as ColumnIndexer>::N_COL },
+        { <Column as ColumnIndexer<usize>>::N_COL },
+        { <Column as ColumnIndexer<usize>>::N_COL },
         N_DSEL,
         N_FSEL,
         DummyLookupTable,

ivc/src/prover.rs

@@ -91,7 +91,7 @@ impl<
         F: Clone,
     > ColumnEvaluations<F> for ProofEvaluations<N_WIT, N_REL, N_DSEL, N_FSEL, F>
 {
-    type Column = kimchi_msm::columns::Column;
+    type Column = kimchi_msm::columns::Column<usize>;
 
     fn evaluate(&self, col: Self::Column) -> Result<PointEvaluations<F>, ExprError<Self::Column>> {
         // TODO: substitute when non-literal generic constants are available
@@ -147,7 +147,7 @@ pub struct Proof<
 pub fn prove<
     EFqSponge: Clone + FqSponge<Fq, G, Fp>,
     EFrSponge: FrSponge<Fp>,
-    FC: FoldingConfig<Column = GenericColumn, Curve = G, Challenge = PlonkishChallenge>,
+    FC: FoldingConfig<Column = GenericColumn<usize>, Curve = G, Challenge = PlonkishChallenge>,
     RNG,
     const N_WIT: usize,
     const N_WIT_QUAD: usize, // witness columns + quad columns

ivc/src/verifier.rs

@@ -40,7 +40,7 @@ pub type Fq = ark_bn254::Fq;
 pub fn verify<
     EFqSponge: Clone + FqSponge<Fq, G, Fp>,
     EFrSponge: FrSponge<Fp>,
-    FC: FoldingConfig<Column = GenericColumn, Curve = G, Challenge = PlonkishChallenge>,
+    FC: FoldingConfig<Column = GenericColumn<usize>, Curve = G, Challenge = PlonkishChallenge>,
     const N_WIT: usize,
     const N_REL: usize,
     const N_DSEL: usize,

ivc/tests/folding_ivc.rs

@@ -71,7 +71,7 @@ fn test_regression_additional_columns_reduction_to_degree_2() {
     impl Witness<Curve> for TestWitness {}
 
     impl FoldingConfig for TestConfig {
-        type Column = Column;
+        type Column = Column<usize>;
 
         type Selector = ();
 
@@ -92,7 +92,7 @@ fn test_regression_additional_columns_reduction_to_degree_2() {
 
     struct Env;
 
-    impl FoldingEnv<Fp, TestInstance, TestWitness, Column, Challenge, ()> for Env {
+    impl FoldingEnv<Fp, TestInstance, TestWitness, Column<usize>, Challenge, ()> for Env {
         type Structure = ();
 
         fn new(
@@ -103,7 +103,7 @@ fn test_regression_additional_columns_reduction_to_degree_2() {
             todo!()
         }
 
-        fn col(&self, _col: Column, _curr_or_next: CurrOrNext, _side: Side) -> &[Fp] {
+        fn col(&self, _col: Column<usize>, _curr_or_next: CurrOrNext, _side: Side) -> &[Fp] {
             todo!()
         }
 

ivc/tests/simple.rs

@@ -67,10 +67,10 @@ pub enum AdditionColumn {
     C,
 }
 
-impl ColumnIndexer for AdditionColumn {
+impl ColumnIndexer<usize> for AdditionColumn {
     const N_COL: usize = 3;
 
-    fn to_column(self) -> Column {
+    fn to_column(self) -> Column<usize> {
         match self {
             AdditionColumn::A => Column::Relation(0),
             AdditionColumn::B => Column::Relation(1),
@@ -119,7 +119,7 @@ pub fn heavy_test_simple_add() {
     const N_FSEL_TOTAL: usize = N_FSEL_IVC;
 
     // Total number of witness columns in IVC. The blocks are public selectors.
-    const N_WIT_IVC: usize = <IVCColumn as ColumnIndexer>::N_COL - N_FSEL_IVC;
+    const N_WIT_IVC: usize = <IVCColumn as ColumnIndexer<usize>>::N_COL - N_FSEL_IVC;
 
     // Number of witness columns in the circuit.
     // It consists of the columns of the inner circuit and the columns for the
@@ -189,7 +189,7 @@ pub fn heavy_test_simple_add() {
         const N_ALPHAS: usize,
     > = StandardConfig<
         Curve,
-        Column,
+        Column<usize>,
         PlonkishChallenge,
         PlonkishInstance<Curve, N_COL_TOTAL, N_CHALS, N_ALPHAS>, // TODO check if it's quad or not
         PlonkishWitness<N_COL_TOTAL, N_FSEL, Fp>,

msm/src/circuit_design/capabilities.rs

@@ -10,7 +10,7 @@ use ark_ff::PrimeField;
 
 /// Environment capability for accessing and reading columns. This is necessary for
 /// building constraints.
-pub trait ColAccessCap<F: PrimeField, CIx: ColumnIndexer> {
+pub trait ColAccessCap<F: PrimeField, CIx: ColumnIndexer<usize>> {
     // NB: 'static here means that `Variable` does not contain any
     // references with a lifetime less than 'static. Which is true in
     // our case. Necessary for `set_assert_mapper`
@@ -39,15 +39,15 @@ pub trait ColAccessCap<F: PrimeField, CIx: ColumnIndexer> {
 
 /// Environment capability similar to `ColAccessCap` but for /also
 /// writing/ columns. Used on the witness side.
-pub trait ColWriteCap<F: PrimeField, CIx: ColumnIndexer>
+pub trait ColWriteCap<F: PrimeField, CIx: ColumnIndexer<usize>>
 where
     Self: ColAccessCap<F, CIx>,
 {
     fn write_column(&mut self, col: CIx, value: &Self::Variable);
 }
 
 /// Capability for invoking table lookups.
-pub trait LookupCap<F: PrimeField, CIx: ColumnIndexer, LT: LookupTableID>
+pub trait LookupCap<F: PrimeField, CIx: ColumnIndexer<usize>, LT: LookupTableID>
 where
     Self: ColAccessCap<F, CIx>,
 {
@@ -62,7 +62,7 @@ where
 /// Holds a "current" row that can be moved forward with `next_row`.
 /// The `ColWriteCap` and `ColAccessCap` reason in terms of current
 /// row. The two other methods can be used to read/write previous.
-pub trait MultiRowReadCap<F: PrimeField, CIx: ColumnIndexer>
+pub trait MultiRowReadCap<F: PrimeField, CIx: ColumnIndexer<usize>>
 where
     Self: ColWriteCap<F, CIx>,
 {
@@ -84,7 +84,7 @@ where
 // F-typed inputs to a function.
 /// A direct field access capability modelling an abstract witness
 /// builder. Not for constraint building.
-pub trait DirectWitnessCap<F: PrimeField, CIx: ColumnIndexer>
+pub trait DirectWitnessCap<F: PrimeField, CIx: ColumnIndexer<usize>>
 where
     Self: MultiRowReadCap<F, CIx>,
 {
@@ -106,7 +106,7 @@ where
 /// partially) in the constraint builder case. For example, "hcopy",
 /// despite its name, does not do any "write", so hcopy !=>
 /// write_column.
-pub trait HybridCopyCap<F: PrimeField, CIx: ColumnIndexer>
+pub trait HybridCopyCap<F: PrimeField, CIx: ColumnIndexer<usize>>
 where
     Self: ColAccessCap<F, CIx>,
 {
@@ -120,7 +120,7 @@ where
 ////////////////////////////////////////////////////////////////////////////
 
 /// Write an array of values simultaneously.
-pub fn read_column_array<F, Env, const ARR_N: usize, CIx: ColumnIndexer, ColMap>(
+pub fn read_column_array<F, Env, const ARR_N: usize, CIx: ColumnIndexer<usize>, ColMap>(
     env: &mut Env,
     column_map: ColMap,
 ) -> [Env::Variable; ARR_N]
@@ -133,7 +133,7 @@ where
 }
 
 /// Write a field element directly as a constant.
-pub fn write_column_const<F, Env, CIx: ColumnIndexer>(env: &mut Env, col: CIx, var: &F)
+pub fn write_column_const<F, Env, CIx: ColumnIndexer<usize>>(env: &mut Env, col: CIx, var: &F)
 where
     F: PrimeField,
     Env: ColWriteCap<F, CIx>,
@@ -142,7 +142,7 @@ where
 }
 
 /// Write an array of values simultaneously.
-pub fn write_column_array<F, Env, const ARR_N: usize, CIx: ColumnIndexer, ColMap>(
+pub fn write_column_array<F, Env, const ARR_N: usize, CIx: ColumnIndexer<usize>, ColMap>(
     env: &mut Env,
     input: [Env::Variable; ARR_N],
     column_map: ColMap,
@@ -157,7 +157,7 @@ pub fn write_column_array<F, Env, const ARR_N: usize, CIx: ColumnIndexer, ColMap
 }
 
 /// Write an array of /field/ values simultaneously.
-pub fn write_column_array_const<F, Env, const ARR_N: usize, CIx: ColumnIndexer, ColMap>(
+pub fn write_column_array_const<F, Env, const ARR_N: usize, CIx: ColumnIndexer<usize>, ColMap>(
     env: &mut Env,
     input: &[F; ARR_N],
     column_map: ColMap,