Add compression logic

Author: hratoanina

.github/workflows/ci.yml

@@ -97,7 +97,7 @@ jobs:
           # cargo build --verbose --target ${{ env.target }} -p p3-fri-air
           # cargo build --verbose --target ${{ env.target }} -p p3-interpolation-air
           cargo build --verbose --target ${{ env.target }} -p p3-mmcs-air
-          cargo build --verbose --target ${{ env.target }} -p p3-poseidon2-air
+          cargo build --verbose --target ${{ env.target }} -p p3-poseidon2-circuit-air
           cargo build --verbose --target ${{ env.target }} -p p3-symmetric-air
           cargo build --verbose --target ${{ env.target }} -p p3-recursion
 

poseidon2-circuit-air/src/air.rs

@@ -156,24 +156,30 @@ pub(crate) fn eval<
 ) {
     air.p3_poseidon2.eval(builder);
 
+    // SPONGE CONSTRAINTS
     let next_no_reset = AB::Expr::ONE - next.reset.clone();
     for i in 0..(CAPACITY * D) {
         // The first row has capacity zeroed.
         builder
+            .when(local.is_sponge.clone())
             .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.is_sponge.clone())
             .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(),
-        );
+        builder
+            .when(local.is_sponge.clone())
+            .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(),
+            );
     }
 
     let mut next_absorb = [AB::Expr::ZERO; RATE];
@@ -187,10 +193,14 @@ pub(crate) fn eval<
     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(),
-        );
+        builder
+            .when(local.is_sponge.clone())
+            .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];
@@ -201,11 +211,15 @@ pub(crate) fn eval<
     }
     let current_no_absorb =
         array::from_fn::<_, RATE, _>(|i| AB::Expr::ONE - current_absorb[i].clone());
+    builder.assert_eq(
+        local.is_sponge.clone() * local.reset.clone(),
+        local.sponge_reset.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())
+            .when(local.sponge_reset.clone() * col.clone())
             .assert_zeros(arr);
     }
 
@@ -215,6 +229,11 @@ pub(crate) fn eval<
     //      * local.rate[i] comes from input lookups.
     // - If is_squeeze = 1:
     //      * local.rate is sent to output lookups.
+
+    // COMPRESSION CONSTRAINTS
+    // TODO: Add all lookups:
+    // - local input state comes from input lookups.
+    // - send local output state to output lookups.
 }
 
 impl<

poseidon2-circuit-air/src/columns.rs

@@ -6,12 +6,17 @@ use p3_poseidon2_air::Poseidon2Cols;
 ///
 /// They extend the P3 columns with some circuit-specific columns.
 ///
+/// `is_sponge` (transparent): if `1`, this row performs a sponge operation (absorb or squeeze);
+/// otherwise, it performs a compression.
 /// `reset` (transparent): indicates whether the state is being reset this row.
+/// `sponge_reset`: auxiliary column to keep constraint degrees below three.
 /// `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).
+/// `input_indices` (transparent): for each input element, indicates the index in the witness table for the
+/// memory lookup. It's either received (for an absorb or a compression) or sent (for a squeeze).
+/// `output_indices` (transparent): for each output element, indicates the index in the witness table for the
+/// memory lookup. Only used by compressions to send the output.
 #[repr(C)]
 pub struct Poseidon2CircuitCols<
     T,
@@ -25,9 +30,12 @@ pub struct Poseidon2CircuitCols<
     pub poseidon2:
         Poseidon2Cols<T, WIDTH, SBOX_DEGREE, SBOX_REGISTERS, HALF_FULL_ROUNDS, PARTIAL_ROUNDS>,
 
+    pub is_sponge: T,
     pub reset: T,
+    pub sponge_reset: T,
     pub absorb_flags: [T; RATE],
-    pub input_indices: [T; RATE],
+    pub input_indices: [T; WIDTH],
+    pub output_indices: [T; RATE],
 }
 
 pub const fn num_cols<

poseidon2-circuit-air/src/lib.rs

@@ -1,4 +1,4 @@
-//! And AIR for the Poseidon2 permutation.
+//! An AIR for the Poseidon2 table for recursion. Handles sponge operations and compressions.
 
 #![no_std]