Refactor columns

Author: hratoanina

poseidon2-circuit-air/src/air.rs

@@ -9,11 +9,19 @@ use p3_poseidon2_air::{Poseidon2Air, RoundConstants};
 
 use crate::{Poseidon2CircuitCols, num_cols};
 
+/// Extends the Poseidon2 AIR with recursion circuit-specific columns and constraints.
 /// Assumes the field size is at least 16 bits.
+///
+/// SPECIFIC ASSUMPTIONS:
+/// - Memory elements from the witness table are extension elements of degree D.
+/// - RATE and CAPACITY are the number of extension elements in the rate/capacity.
+/// - WIDTH is the number of field elements in the state, i.e., (RATE + CAPACITY) * D.
+/// - `reset` can only be set during an absorb.
 #[derive(Debug)]
 pub struct Poseidon2CircuitAir<
     F: PrimeCharacteristicRing,
     LinearLayers,
+    const D: usize,
     const RATE: usize,
     const CAPACITY: usize,
     const WIDTH: usize,
@@ -36,6 +44,7 @@ pub struct Poseidon2CircuitAir<
 impl<
     F: PrimeCharacteristicRing,
     LinearLayers,
+    const D: usize,
     const RATE: usize,
     const CAPACITY: usize,
     const WIDTH: usize,
@@ -47,6 +56,7 @@ impl<
     Poseidon2CircuitAir<
         F,
         LinearLayers,
+        D,
         RATE,
         CAPACITY,
         WIDTH,
@@ -59,6 +69,11 @@ impl<
     pub const fn new(
         constants: RoundConstants<F, WIDTH, HALF_FULL_ROUNDS, PARTIAL_ROUNDS>,
     ) -> Self {
+        assert!((CAPACITY + RATE) * D == WIDTH);
+        assert!(RATE.is_multiple_of(D));
+        assert!(CAPACITY.is_multiple_of(D));
+        assert!(WIDTH.is_multiple_of(D));
+
         Self {
             p3_poseidon2: Poseidon2Air::new(constants),
         }
@@ -68,6 +83,7 @@ impl<
 impl<
     F: PrimeCharacteristicRing + Sync,
     LinearLayers: Sync,
+    const D: usize,
     const RATE: usize,
     const CAPACITY: usize,
     const WIDTH: usize,
@@ -79,6 +95,7 @@ impl<
     for Poseidon2CircuitAir<
         F,
         LinearLayers,
+        D,
         RATE,
         CAPACITY,
         WIDTH,
@@ -89,13 +106,14 @@ impl<
     >
 {
     fn width(&self) -> usize {
-        num_cols::<WIDTH, SBOX_DEGREE, SBOX_REGISTERS, HALF_FULL_ROUNDS, PARTIAL_ROUNDS>()
+        num_cols::<WIDTH, RATE, SBOX_DEGREE, SBOX_REGISTERS, HALF_FULL_ROUNDS, PARTIAL_ROUNDS>()
     }
 }
 
 pub(crate) fn eval<
     AB: AirBuilder,
     LinearLayers: GenericPoseidon2LinearLayers<WIDTH>,
+    const D: usize,
     const RATE: usize,
     const CAPACITY: usize,
     const WIDTH: usize,
@@ -107,6 +125,7 @@ pub(crate) fn eval<
     air: &Poseidon2CircuitAir<
         AB::F,
         LinearLayers,
+        D,
         RATE,
         CAPACITY,
         WIDTH,
@@ -119,6 +138,7 @@ pub(crate) fn eval<
     local: &Poseidon2CircuitCols<
         AB::Var,
         WIDTH,
+        RATE,
         SBOX_DEGREE,
         SBOX_REGISTERS,
         HALF_FULL_ROUNDS,
@@ -127,6 +147,7 @@ pub(crate) fn eval<
     next: &Poseidon2CircuitCols<
         AB::Var,
         WIDTH,
+        RATE,
         SBOX_DEGREE,
         SBOX_REGISTERS,
         HALF_FULL_ROUNDS,
@@ -135,39 +156,71 @@ pub(crate) fn eval<
 ) {
     air.p3_poseidon2.eval(builder);
 
-    // When resetting the state, we just have to clear the capacity. The rate will be overwritten by the input.
-    builder
-        .when(local.reset.clone())
-        .assert_zeros::<CAPACITY, _>(array::from_fn(|i| local.poseidon2.inputs[i + RATE].clone()));
-
-    // If the next row doesn't reset, propagate the capacity.
     let next_no_reset = AB::Expr::ONE - next.reset.clone();
-    builder
-        .when(next_no_reset)
-        .assert_zeros::<CAPACITY, _>(array::from_fn(|i| {
-            next.poseidon2.inputs[i + RATE].clone()
-                - local.poseidon2.ending_full_rounds[HALF_FULL_ROUNDS - 1].post[i + RATE].clone()
-        }));
+    for i in 0..(CAPACITY * D) {
+        // The first row has capacity zeroed.
+        builder
+            .when_first_row()
+            .assert_zero(local.poseidon2.inputs[RATE * D + i].clone());
+
+        // When resetting the state, we just have to clear the capacity. The rate will be overwritten by the input.
+        builder
+            .when(local.reset.clone())
+            .assert_zero(local.poseidon2.inputs[RATE * D + i].clone());
+
+        // If the next row doesn't reset, propagate the capacity.
+        builder.when(next_no_reset.clone()).assert_zero(
+            next.poseidon2.inputs[RATE * D + i].clone()
+                - local.poseidon2.ending_full_rounds[HALF_FULL_ROUNDS - 1].post[RATE * D + i]
+                    .clone(),
+        );
+    }
 
-    // If the next row is squeezing, then we propagate the entire output state forward.
-    let next_squeeze = AB::Expr::ONE - next.absorb.clone();
-    builder
-        .when(next_squeeze)
-        .assert_zeros::<WIDTH, _>(array::from_fn(|i| {
-            next.poseidon2.inputs[i].clone()
-                - local.poseidon2.ending_full_rounds[HALF_FULL_ROUNDS - 1].post[i].clone()
-        }));
+    let mut next_absorb = [AB::Expr::ZERO; RATE];
+    for i in 0..RATE {
+        for col in next_absorb.iter_mut().take(RATE).skip(i) {
+            *col += next.absorb_flags[i].clone();
+        }
+    }
+    let next_no_absorb = array::from_fn::<_, RATE, _>(|i| AB::Expr::ONE - next_absorb[i].clone());
+    // In the next row, each rate element not being absorbed comes from the current row.
+    for index in 0..(RATE * D) {
+        let i = index / D;
+        let j = index % D;
+        builder.when(next_no_absorb[i].clone()).assert_zero(
+            next.poseidon2.inputs[i * D + j].clone()
+                - local.poseidon2.ending_full_rounds[HALF_FULL_ROUNDS - 1].post[i * D + j].clone(),
+        );
+    }
+
+    let mut current_absorb = [AB::Expr::ZERO; RATE];
+    for i in 0..RATE {
+        for col in current_absorb.iter_mut().take(RATE).skip(i) {
+            *col += local.absorb_flags[i].clone();
+        }
+    }
+    let current_no_absorb =
+        array::from_fn::<_, RATE, _>(|i| AB::Expr::ONE - current_absorb[i].clone());
+    // During a reset, the rate elements not being absorbed are zeroed.
+    for (i, col) in current_no_absorb.iter().enumerate().take(RATE) {
+        let arr = array::from_fn::<_, D, _>(|j| local.poseidon2.inputs[i * D + j].clone().into());
+        builder
+            .when(local.reset.clone() * col.clone())
+            .assert_zeros(arr);
+    }
 
+    let _is_squeeze = AB::Expr::ONE - current_absorb[0].clone();
     // TODO: Add all lookups:
-    // - If local.absorb = 1:
-    //      * local.rate comes from input lookups.
-    // - If local.absorb = 0:
+    // - If current_absorb[i] = 1:
+    //      * local.rate[i] comes from input lookups.
+    // - If is_squeeze = 1:
     //      * local.rate is sent to output lookups.
 }
 
 impl<
     AB: AirBuilder,
     LinearLayers: GenericPoseidon2LinearLayers<WIDTH>,
+    const D: usize,
     const RATE: usize,
     const CAPACITY: usize,
     const WIDTH: usize,
@@ -179,6 +232,7 @@ impl<
     for Poseidon2CircuitAir<
         AB::F,
         LinearLayers,
+        D,
         RATE,
         CAPACITY,
         WIDTH,
@@ -199,6 +253,7 @@ impl<
         eval::<
             _,
             _,
+            D,
             RATE,
             CAPACITY,
             WIDTH,

poseidon2-circuit-air/src/columns.rs

@@ -5,10 +5,18 @@ use p3_poseidon2_air::Poseidon2Cols;
 /// Columns for a Poseidon2 AIR which computes one permutation per row.
 ///
 /// They extend the P3 columns with some circuit-specific columns.
+///
+/// `reset` (transparent): indicates whether the state is being reset this row.
+/// `absorb_flags` (transparent): for each rate element, indicates if it is being absorbed this row.
+/// At most one flag is set to 1 per row: if `absorb_flags[i]` is 1, then all elements up to the `i`-th
+/// are absorbed; the rest are propagated from the previous row.
+/// `input_indices` (transparent): for each rate element, indicates the index in the witness table for the
+/// memory lookup. It's either received (for an absorb) or sent (for a squeeze).
 #[repr(C)]
 pub struct Poseidon2CircuitCols<
     T,
     const WIDTH: usize,
+    const RATE: usize,
     const SBOX_DEGREE: u64,
     const SBOX_REGISTERS: usize,
     const HALF_FULL_ROUNDS: usize,
@@ -18,13 +26,13 @@ pub struct Poseidon2CircuitCols<
         Poseidon2Cols<T, WIDTH, SBOX_DEGREE, SBOX_REGISTERS, HALF_FULL_ROUNDS, PARTIAL_ROUNDS>,
 
     pub reset: T,
-    pub absorb: T,
-    pub input_indices: [T; WIDTH],
-    pub output_indices: [T; WIDTH],
+    pub absorb_flags: [T; RATE],
+    pub input_indices: [T; RATE],
 }
 
 pub const fn num_cols<
     const WIDTH: usize,
+    const RATE: usize,
     const SBOX_DEGREE: u64,
     const SBOX_REGISTERS: usize,
     const HALF_FULL_ROUNDS: usize,
@@ -34,6 +42,7 @@ pub const fn num_cols<
         Poseidon2CircuitCols<
             u8,
             WIDTH,
+            RATE,
             SBOX_DEGREE,
             SBOX_REGISTERS,
             HALF_FULL_ROUNDS,
@@ -45,6 +54,7 @@ pub const fn num_cols<
 impl<
     T,
     const WIDTH: usize,
+    const RATE: usize,
     const SBOX_DEGREE: u64,
     const SBOX_REGISTERS: usize,
     const HALF_FULL_ROUNDS: usize,
@@ -54,6 +64,7 @@ impl<
         Poseidon2CircuitCols<
             T,
             WIDTH,
+            RATE,
             SBOX_DEGREE,
             SBOX_REGISTERS,
             HALF_FULL_ROUNDS,
@@ -66,6 +77,7 @@ impl<
     ) -> &Poseidon2CircuitCols<
         T,
         WIDTH,
+        RATE,
         SBOX_DEGREE,
         SBOX_REGISTERS,
         HALF_FULL_ROUNDS,
@@ -75,6 +87,7 @@ impl<
             self.align_to::<Poseidon2CircuitCols<
                 T,
                 WIDTH,
+                RATE,
                 SBOX_DEGREE,
                 SBOX_REGISTERS,
                 HALF_FULL_ROUNDS,
@@ -91,6 +104,7 @@ impl<
 impl<
     T,
     const WIDTH: usize,
+    const RATE: usize,
     const SBOX_DEGREE: u64,
     const SBOX_REGISTERS: usize,
     const HALF_FULL_ROUNDS: usize,
@@ -100,6 +114,7 @@ impl<
         Poseidon2CircuitCols<
             T,
             WIDTH,
+            RATE,
             SBOX_DEGREE,
             SBOX_REGISTERS,
             HALF_FULL_ROUNDS,
@@ -112,6 +127,7 @@ impl<
     ) -> &mut Poseidon2CircuitCols<
         T,
         WIDTH,
+        RATE,
         SBOX_DEGREE,
         SBOX_REGISTERS,
         HALF_FULL_ROUNDS,
@@ -121,6 +137,7 @@ impl<
             self.align_to_mut::<Poseidon2CircuitCols<
                 T,
                 WIDTH,
+                RATE,
                 SBOX_DEGREE,
                 SBOX_REGISTERS,
                 HALF_FULL_ROUNDS,