Merge pull request #1287 from o1-labs/witness/index-cell

Create `IndexCell` to initialize witness from vectors

kimchi/src/circuits/polynomials/foreign_field_add/witness.rs

@@ -185,9 +185,9 @@ fn init_ffadd_row<F: PrimeField>(
     overflow: F,
     carry: F,
 ) {
-    let witness_shape: Vec<[Box<dyn WitnessCell<F>>; COLUMNS]> = vec![
+    let layout: [Vec<Box<dyn WitnessCell<F>>>; 1] = [
         // ForeignFieldAdd row
-        [
+        vec![
             VariableCell::create("left_lo"),
             VariableCell::create("left_mi"),
             VariableCell::create("left_hi"),
@@ -209,7 +209,7 @@ fn init_ffadd_row<F: PrimeField>(
     witness::init(
         witness,
         offset,
-        &witness_shape,
+        &layout,
         &variable_map!["left_lo" => left[LO], "left_mi" => left[MI], "left_hi" => left[HI], "right_lo" => right[LO], "right_mi" => right[MI], "right_hi" => right[HI], "overflow" => overflow, "carry" => carry],
     );
 }
@@ -221,8 +221,8 @@ fn init_bound_rows<F: PrimeField>(
     bound: &[F; 3],
     carry: &F,
 ) {
-    let witness_shape: Vec<[Box<dyn WitnessCell<F>>; COLUMNS]> = vec![
-        [
+    let layout: [Vec<Box<dyn WitnessCell<F>>>; 2] = [
+        vec![
             // ForeignFieldAdd row
             VariableCell::create("result_lo"),
             VariableCell::create("result_mi"),
@@ -240,7 +240,7 @@ fn init_bound_rows<F: PrimeField>(
             ConstantCell::create(F::zero()),
             ConstantCell::create(F::zero()),
         ],
-        [
+        vec![
             // Zero Row
             VariableCell::create("bound_lo"),
             VariableCell::create("bound_mi"),
@@ -263,7 +263,7 @@ fn init_bound_rows<F: PrimeField>(
     witness::init(
         witness,
         offset,
-        &witness_shape,
+        &layout,
         &variable_map!["carry" => *carry, "result_lo" => result[LO], "result_mi" => result[MI], "result_hi" => result[HI], "bound_lo" => bound[LO], "bound_mi" => bound[MI], "bound_hi" => bound[HI]],
     );
 }

kimchi/src/circuits/polynomials/foreign_field_mul/witness.rs

@@ -41,10 +41,10 @@ use super::circuitgates;
 //
 //     so that most significant limb, q2, is in W[2][0].
 //
-fn create_layout<F: PrimeField>() -> [[Box<dyn WitnessCell<F>>; COLUMNS]; 2] {
+fn create_layout<F: PrimeField>() -> [Vec<Box<dyn WitnessCell<F>>>; 2] {
     [
         // ForeignFieldMul row
-        [
+        vec![
             // Copied for multi-range-check
             VariableCell::create("left_input0"),
             VariableCell::create("left_input1"),
@@ -64,7 +64,7 @@ fn create_layout<F: PrimeField>() -> [[Box<dyn WitnessCell<F>>; COLUMNS]; 2] {
             VariableBitsCell::create("carry1", 90, None),
         ],
         // Zero row
-        [
+        vec![
             // Copied for multi-range-check
             VariableCell::create("remainder01"),
             VariableCell::create("remainder2"),

kimchi/src/circuits/polynomials/range_check/witness.rs

@@ -29,14 +29,14 @@ use o1_utils::foreign_field::BigUintForeignFieldHelpers;
 ///     For example, we can convert the `RangeCheck0` circuit gate into
 ///     a 64-bit lookup by adding two copy constraints to constrain
 ///     columns 1 and 2 to zero.
-fn layout<F: PrimeField>() -> Vec<[Box<dyn WitnessCell<F>>; COLUMNS]> {
-    vec![
+fn layout<F: PrimeField>() -> [Vec<Box<dyn WitnessCell<F>>>; 4] {
+    [
         /* row 1, RangeCheck0 row */
         range_check_0_row("v0", 0),
         /* row 2, RangeCheck0 row */
         range_check_0_row("v1", 1),
         /* row 3, RangeCheck1 row */
-        [
+        vec![
             VariableCell::create("v2"),
             VariableCell::create("v12"), // optional
             /* 2-bit crumbs (placed here to keep lookup pattern */
@@ -58,7 +58,7 @@ fn layout<F: PrimeField>() -> Vec<[Box<dyn WitnessCell<F>>; COLUMNS]> {
             CopyBitsCell::create(2, 0, 22, 24),
         ],
         /* row 4, Zero row */
-        [
+        vec![
             CopyBitsCell::create(2, 0, 20, 22),
             /* 2-bit crumbs (placed here to keep lookup pattern */
             /*               the same as RangeCheck0) */
@@ -86,8 +86,8 @@ fn layout<F: PrimeField>() -> Vec<[Box<dyn WitnessCell<F>>; COLUMNS]> {
 pub fn range_check_0_row<F: PrimeField>(
     limb_name: &'static str,
     row: usize,
-) -> [Box<dyn WitnessCell<F>>; COLUMNS] {
-    [
+) -> Vec<Box<dyn WitnessCell<F>>> {
+    vec![
         VariableCell::create(limb_name),
         /* 12-bit copies */
         // Copy cells are required because we have a limit

kimchi/src/circuits/polynomials/rot.rs

@@ -266,16 +266,16 @@ where
 
 // ROTATION WITNESS COMPUTATION
 
-fn layout_rot64<F: PrimeField>(curr_row: usize) -> [[Box<dyn WitnessCell<F>>; COLUMNS]; 3] {
+fn layout_rot64<F: PrimeField>(curr_row: usize) -> [Vec<Box<dyn WitnessCell<F>>>; 3] {
     [
         rot_row(),
         range_check_0_row("shifted", curr_row + 1),
         range_check_0_row("excess", curr_row + 2),
     ]
 }
 
-fn rot_row<F: PrimeField>() -> [Box<dyn WitnessCell<F>>; COLUMNS] {
-    [
+fn rot_row<F: PrimeField>() -> Vec<Box<dyn WitnessCell<F>>> {
+    vec![
         VariableCell::create("word"),
         VariableCell::create("rotated"),
         VariableCell::create("excess"),

kimchi/src/circuits/polynomials/xor.rs

@@ -169,7 +169,7 @@ where
 }
 
 // Witness layout
-fn layout<F: PrimeField>(curr_row: usize, bits: usize) -> Vec<[Box<dyn WitnessCell<F>>; COLUMNS]> {
+fn layout<F: PrimeField>(curr_row: usize, bits: usize) -> Vec<Vec<Box<dyn WitnessCell<F>>>> {
     let num_xor = num_xors(bits);
     let mut layout = (0..num_xor)
         .map(|i| xor_row(i, curr_row + i))
@@ -178,9 +178,9 @@ fn layout<F: PrimeField>(curr_row: usize, bits: usize) -> Vec<[Box<dyn WitnessCe
     layout
 }
 
-fn xor_row<F: PrimeField>(nybble: usize, curr_row: usize) -> [Box<dyn WitnessCell<F>>; COLUMNS] {
+fn xor_row<F: PrimeField>(nybble: usize, curr_row: usize) -> Vec<Box<dyn WitnessCell<F>>> {
     let start = nybble * 16;
-    [
+    vec![
         VariableBitsCell::create("in1", start, None),
         VariableBitsCell::create("in2", start, None),
         VariableBitsCell::create("out", start, None),
@@ -199,8 +199,8 @@ fn xor_row<F: PrimeField>(nybble: usize, curr_row: usize) -> [Box<dyn WitnessCel
     ]
 }
 
-fn zero_row<F: PrimeField>() -> [Box<dyn WitnessCell<F>>; COLUMNS] {
-    [
+fn zero_row<F: PrimeField>() -> Vec<Box<dyn WitnessCell<F>>> {
+    vec![
         ConstantCell::create(F::zero()),
         ConstantCell::create(F::zero()),
         ConstantCell::create(F::zero()),

kimchi/src/circuits/witness/constant_cell.rs

@@ -1,5 +1,4 @@
 use super::{variables::Variables, WitnessCell};
-use crate::circuits::polynomial::COLUMNS;
 use ark_ff::Field;
 
 /// Witness cell with constant value
@@ -14,8 +13,8 @@ impl<F: Field> ConstantCell<F> {
     }
 }
 
-impl<F: Field> WitnessCell<F> for ConstantCell<F> {
-    fn value(&self, _witness: &mut [Vec<F>; COLUMNS], _variables: &Variables<F>) -> F {
+impl<F: Field, const W: usize> WitnessCell<F, F, W> for ConstantCell<F> {
+    fn value(&self, _witness: &mut [Vec<F>; W], _variables: &Variables<F>, _index: usize) -> F {
         self.value
     }
 }

kimchi/src/circuits/witness/copy_bits_cell.rs

@@ -2,7 +2,6 @@ use ark_ff::Field;
 use o1_utils::FieldHelpers;
 
 use super::{variables::Variables, WitnessCell};
-use crate::circuits::polynomial::COLUMNS;
 
 /// Witness cell copied from bits of another witness cell
 pub struct CopyBitsCell {
@@ -24,8 +23,8 @@ impl CopyBitsCell {
     }
 }
 
-impl<F: Field> WitnessCell<F> for CopyBitsCell {
-    fn value(&self, witness: &mut [Vec<F>; COLUMNS], _variables: &Variables<F>) -> F {
+impl<F: Field, const W: usize> WitnessCell<F, F, W> for CopyBitsCell {
+    fn value(&self, witness: &mut [Vec<F>; W], _variables: &Variables<F>, _index: usize) -> F {
         F::from_bits(&witness[self.col][self.row].to_bits()[self.start..self.end])
             .expect("failed to deserialize field bits for copy bits cell")
     }

kimchi/src/circuits/witness/copy_cell.rs

@@ -1,7 +1,6 @@
 use ark_ff::Field;
 
 use super::{variables::Variables, WitnessCell};
-use crate::circuits::polynomial::COLUMNS;
 
 /// Witness cell copied from another witness cell
 pub struct CopyCell {
@@ -16,8 +15,8 @@ impl CopyCell {
     }
 }
 
-impl<F: Field> WitnessCell<F> for CopyCell {
-    fn value(&self, witness: &mut [Vec<F>; COLUMNS], _variables: &Variables<F>) -> F {
+impl<F: Field, const W: usize> WitnessCell<F, F, W> for CopyCell {
+    fn value(&self, witness: &mut [Vec<F>; W], _variables: &Variables<F>, _index: usize) -> F {
         witness[self.col][self.row]
     }
 }

kimchi/src/circuits/witness/copy_shift_cell.rs

@@ -1,5 +1,4 @@
 use super::{variables::Variables, WitnessCell};
-use crate::circuits::polynomial::COLUMNS;
 use ark_ff::Field;
 
 /// Witness cell copied from another cell and shifted
@@ -16,8 +15,8 @@ impl CopyShiftCell {
     }
 }
 
-impl<F: Field> WitnessCell<F> for CopyShiftCell {
-    fn value(&self, witness: &mut [Vec<F>; COLUMNS], _variables: &Variables<F>) -> F {
+impl<F: Field, const W: usize> WitnessCell<F, F, W> for CopyShiftCell {
+    fn value(&self, witness: &mut [Vec<F>; W], _variables: &Variables<F>, _index: usize) -> F {
         F::from(2u32).pow([self.shift]) * witness[self.col][self.row]
     }
 }

kimchi/src/circuits/witness/index_cell.rs

@@ -0,0 +1,29 @@
+use super::{variables::Variables, WitnessCell};
+use ark_ff::Field;
+
+/// Witness cell assigned from an indexable variable
+/// See [Variables] for more details
+pub struct IndexCell<'a> {
+    name: &'a str,
+    length: usize,
+}
+
+impl<'a> IndexCell<'a> {
+    /// Create witness cell assigned from a variable name a length
+    pub fn create(name: &'a str, from: usize, to: usize) -> Box<IndexCell<'a>> {
+        Box::new(IndexCell {
+            name,
+            length: to - from,
+        })
+    }
+}
+
+impl<'a, F: Field, const W: usize> WitnessCell<F, Vec<F>, W> for IndexCell<'a> {
+    fn value(&self, _witness: &mut [Vec<F>; W], variables: &Variables<Vec<F>>, index: usize) -> F {
+        assert!(index < self.length, "index out of bounds of `IndexCell`");
+        variables[self.name][index]
+    }
+    fn length(&self) -> usize {
+        self.length
+    }
+}