Merge pull request #2948 from o1-labs/feature/add-parameter-to-column

Parameterise `Column`

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

@@ -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/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

@@ -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_7/columns.rs

@@ -25,7 +25,7 @@ impl<const STATE_SIZE: usize, const NB_FULL_ROUND: usize> ColumnIndexer
 {
     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_8_56_5_3_2/columns.rs

@@ -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/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

@@ -70,7 +70,7 @@ pub enum AdditionColumn {
 impl ColumnIndexer 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),
@@ -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/constraints.rs

@@ -15,7 +15,7 @@ use crate::{
 
 pub struct ConstraintBuilderEnv<F: PrimeField, LT: LookupTableID> {
     /// An indexed set of constraints.
-    pub constraints: Vec<Expr<ConstantExpr<F, BerkeleyChallengeTerm>, Column>>,
+    pub constraints: Vec<Expr<ConstantExpr<F, BerkeleyChallengeTerm>, Column<usize>>>,
     /// Aggregated lookups or "reads".
     pub lookup_reads: BTreeMap<LT, Vec<Vec<E<F>>>>,
     /// Aggregated "write" lookups, for runtime tables.

msm/src/circuit_design/witness.rs

@@ -284,7 +284,7 @@ impl<
         LT: LookupTableID,
     > WitnessBuilderEnv<F, CIx, N_WIT, N_REL, N_DSEL, N_FSEL, LT>
 {
-    pub fn write_column_raw(&mut self, position: Column, value: F) {
+    pub fn write_column_raw(&mut self, position: Column<usize>, value: F) {
         match position {
             Column::Relation(i) => self.witness.last_mut().unwrap().cols[i] = value,
             Column::FixedSelector(_) => {

msm/src/column_env.rs

@@ -54,7 +54,7 @@ impl<
     > TColumnEnvironment<'a, F, BerkeleyChallengeTerm, BerkeleyChallenges<F>>
     for ColumnEnvironment<'a, N_WIT, N_REL, N_DSEL, N_FSEL, F, ID>
 {
-    type Column = crate::columns::Column;
+    type Column = crate::columns::Column<usize>;
 
     fn get_column(
         &self,

msm/src/columns.rs

@@ -5,9 +5,9 @@ use kimchi::circuits::expr::{CacheId, FormattedOutput};
 
 /// Describe a generic indexed variable X_{i}.
 #[derive(PartialEq, Eq, Clone, Copy, Debug, Hash)]
-pub enum Column {
+pub enum Column<T> {
     /// Columns related to the relation encoded in the circuit
-    Relation(usize),
+    Relation(T),
     /// Columns related to dynamic selectors to indicate gate type
     DynamicSelector(usize),
     /// Constant column that is /always/ fixed for a given circuit.
@@ -25,17 +25,17 @@ pub enum Column {
     LookupFixedTable(u32),
 }
 
-impl Column {
+impl Column<usize> {
     /// Adds offset if the column is `Relation`. Fails otherwise.
-    pub fn add_rel_offset(self, offset: usize) -> Column {
+    pub fn add_rel_offset(self, offset: usize) -> Column<usize> {
         let Column::Relation(i) = self else {
             todo!("add_rel_offset is only implemented for the relation columns")
         };
         Column::Relation(offset + i)
     }
 }
 
-impl FormattedOutput for Column {
+impl FormattedOutput for Column<usize> {
     fn latex(&self, _cache: &mut HashMap<CacheId, Self>) -> String {
         match self {
             Column::Relation(i) => format!("x_{{{i}}}"),
@@ -78,12 +78,12 @@ pub trait ColumnIndexer: core::fmt::Debug + Copy + Eq + Ord {
     const N_COL: usize;
 
     /// Flatten the column "alias" into the integer-like column.
-    fn to_column(self) -> Column;
+    fn to_column(self) -> Column<usize>;
 }
 
 // Implementation to be compatible with folding if we use generic column
 // constraints
-impl FoldingColumnTrait for Column {
+impl<T: Copy> FoldingColumnTrait for Column<T> {
     fn is_witness(&self) -> bool {
         match self {
             // Witness

msm/src/expr.rs

@@ -23,9 +23,10 @@ use crate::columns::Column;
 /// use kimchi::circuits::expr::{ConstantExprInner, ExprInner, Operations, Variable};
 /// use kimchi::circuits::gate::CurrOrNext;
 /// use kimchi::circuits::berkeley_columns::BerkeleyChallengeTerm;
-/// use kimchi_msm::columns::Column;
+/// use kimchi_msm::columns::{Column as GenericColumn};
 /// use kimchi_msm::expr::E;
 /// pub type Fp = ark_bn254::Fr;
+/// pub type Column = GenericColumn<usize>;
 /// let x1 = E::<Fp>::Atom(
 ///     ExprInner::<Operations<ConstantExprInner<Fp, BerkeleyChallengeTerm>>, Column>::Cell(Variable {
 ///         col: Column::Relation(1),
@@ -48,16 +49,16 @@ use crate::columns::Column;
 /// ```
 /// A list of such constraints is used to represent the entire circuit and will
 /// be used to build the quotient polynomial.
-pub type E<F> = Expr<ConstantExpr<F, BerkeleyChallengeTerm>, Column>;
+pub type E<F> = Expr<ConstantExpr<F, BerkeleyChallengeTerm>, Column<usize>>;
 
-pub fn curr_cell<F: Field>(col: Column) -> E<F> {
+pub fn curr_cell<F: Field>(col: Column<usize>) -> E<F> {
     E::Atom(ExprInner::Cell(Variable {
         col,
         row: CurrOrNext::Curr,
     }))
 }
 
-pub fn next_cell<F: Field>(col: Column) -> E<F> {
+pub fn next_cell<F: Field>(col: Column<usize>) -> E<F> {
     E::Atom(ExprInner::Cell(Variable {
         col,
         row: CurrOrNext::Next,

msm/src/fec/columns.rs

@@ -53,7 +53,7 @@ pub enum FECColumn {
 
 impl ColumnIndexer for FECColumnInput {
     const N_COL: usize = 4 * N_LIMBS_LARGE;
-    fn to_column(self) -> Column {
+    fn to_column(self) -> Column<usize> {
         match self {
             FECColumnInput::XP(i) => {
                 assert!(i < N_LIMBS_LARGE);
@@ -77,7 +77,7 @@ impl ColumnIndexer for FECColumnInput {
 
 impl ColumnIndexer for FECColumnOutput {
     const N_COL: usize = 2 * N_LIMBS_SMALL;
-    fn to_column(self) -> Column {
+    fn to_column(self) -> Column<usize> {
         match self {
             FECColumnOutput::XR(i) => {
                 assert!(i < N_LIMBS_SMALL);
@@ -93,7 +93,7 @@ impl ColumnIndexer for FECColumnOutput {
 
 impl ColumnIndexer for FECColumnInter {
     const N_COL: usize = 4 * N_LIMBS_LARGE + 10 * N_LIMBS_SMALL + 9;
-    fn to_column(self) -> Column {
+    fn to_column(self) -> Column<usize> {
         match self {
             FECColumnInter::F(i) => {
                 assert!(i < N_LIMBS_LARGE);
@@ -148,7 +148,7 @@ impl ColumnIndexer for FECColumnInter {
 
 impl ColumnIndexer for FECColumn {
     const N_COL: usize = FEC_N_COLUMNS;
-    fn to_column(self) -> Column {
+    fn to_column(self) -> Column<usize> {
         match self {
             FECColumn::Input(input) => input.to_column(),
             FECColumn::Inter(inter) => inter.to_column().add_rel_offset(FECColumnInput::N_COL),

msm/src/ffa/columns.rs

@@ -22,7 +22,7 @@ pub enum FFAColumn {
 
 impl ColumnIndexer for FFAColumn {
     const N_COL: usize = FFA_N_COLUMNS;
-    fn to_column(self) -> Column {
+    fn to_column(self) -> Column<usize> {
         let to_column_inner = |offset, i| {
             assert!(i < N_LIMBS);
             Column::Relation(N_LIMBS * offset + i)

msm/src/logup.rs

@@ -349,7 +349,7 @@ impl<'lt, G, ID: LookupTableID> IntoIterator for &'lt LookupProof<G, ID> {
 /// h(X) * (β + t(X)) * (β + f(X)) = (β + t(X)) + m(X) * (β + f(X))
 /// ```
 pub fn combine_lookups<F: PrimeField, ID: LookupTableID>(
-    column: Column,
+    column: Column<usize>,
     lookups: Vec<Logup<E<F>, ID>>,
 ) -> E<F> {
     let joint_combiner = {
@@ -419,7 +419,7 @@ pub fn constraint_lookups<F: PrimeField, ID: LookupTableID>(
     lookup_writes: &BTreeMap<ID, Vec<Vec<E<F>>>>,
 ) -> Vec<E<F>> {
     let mut constraints: Vec<E<F>> = vec![];
-    let mut lookup_terms_cols: Vec<Column> = vec![];
+    let mut lookup_terms_cols: Vec<Column<usize>> = vec![];
     lookup_reads.iter().for_each(|(table_id, reads)| {
         let mut idx_partial_sum = 0;
         let table_id_u32 = table_id.to_u32();

msm/src/proof.rs

@@ -95,7 +95,7 @@ impl<
         ID: LookupTableID,
     > ColumnEvaluations<F> for ProofEvaluations<N_WIT, N_REL, N_DSEL, N_FSEL, F, ID>
 {
-    type Column = crate::columns::Column;
+    type Column = crate::columns::Column<usize>;
 
     fn evaluate(&self, col: Self::Column) -> Result<PointEvaluations<F>, ExprError<Self::Column>> {
         // TODO: substitute when non-literal generic constants are available

msm/src/serialization/column.rs

@@ -42,7 +42,7 @@ pub enum SerializationColumn {
 
 impl ColumnIndexer for SerializationColumn {
     const N_COL: usize = N_COL_SER;
-    fn to_column(self) -> Column {
+    fn to_column(self) -> Column<usize> {
         match self {
             Self::CurrentRow => Column::FixedSelector(0),
             Self::PreviousCoeffRow => Column::FixedSelector(1),

msm/src/test/test_circuit/columns.rs

@@ -23,7 +23,7 @@ pub enum TestColumn {
 
 impl ColumnIndexer for TestColumn {
     const N_COL: usize = N_COL_TEST;
-    fn to_column(self) -> Column {
+    fn to_column(self) -> Column<usize> {
         let to_column_inner = |offset, i| {
             assert!(i < N_LIMBS);
             Column::Relation(N_LIMBS * offset + i)

o1vm/src/interpreters/keccak/column.rs

@@ -375,7 +375,7 @@ impl<T: Clone> IndexMut<ColumnAlias> for KeccakWitness<T> {
 
 impl ColumnIndexer for ColumnAlias {
     const N_COL: usize = N_ZKVM_KECCAK_REL_COLS + N_ZKVM_KECCAK_SEL_COLS;
-    fn to_column(self) -> Column {
+    fn to_column(self) -> Column<usize> {
         Column::Relation(usize::from(self))
     }
 }
@@ -403,7 +403,7 @@ impl<T: Clone> IndexMut<Steps> for KeccakWitness<T> {
 
 impl ColumnIndexer for Steps {
     const N_COL: usize = N_ZKVM_KECCAK_REL_COLS + N_ZKVM_KECCAK_SEL_COLS;
-    fn to_column(self) -> Column {
+    fn to_column(self) -> Column<usize> {
         Column::DynamicSelector(usize::from(self) - N_ZKVM_KECCAK_REL_COLS)
     }
 }

o1vm/src/interpreters/mips/column.rs

@@ -146,7 +146,7 @@ impl<T: Clone> IndexMut<ColumnAlias> for MIPSWitness<T> {
 impl ColumnIndexer for ColumnAlias {
     const N_COL: usize = N_MIPS_COLS;
 
-    fn to_column(self) -> Column {
+    fn to_column(self) -> Column<usize> {
         match self {
             Self::ScratchState(ss) => {
                 assert!(
@@ -200,7 +200,7 @@ impl<T: Clone> IndexMut<Instruction> for MIPSWitness<T> {
 
 impl ColumnIndexer for Instruction {
     const N_COL: usize = N_MIPS_REL_COLS + N_MIPS_SEL_COLS;
-    fn to_column(self) -> Column {
+    fn to_column(self) -> Column<usize> {
         Column::DynamicSelector(usize::from(self) - N_MIPS_REL_COLS)
     }
 }