feat: add vector equality chip

chips/src/is_equal/air.rs

@@ -0,0 +1,44 @@
+use std::borrow::Borrow;
+
+use super::columns::{IsEqualCols, IsEqualIOCols, NUM_COLS};
+use super::IsEqualChip;
+use crate::sub_chip::{AirConfig, SubAir};
+use afs_stark_backend::interaction::Chip;
+use p3_air::{Air, AirBuilder, BaseAir};
+use p3_field::AbstractField;
+use p3_field::Field;
+use p3_matrix::Matrix;
+
+impl<F: Field> BaseAir<F> for IsEqualChip {
+    fn width(&self) -> usize {
+        NUM_COLS
+    }
+}
+
+impl<AB: AirBuilder> Air<AB> for IsEqualChip {
+    fn eval(&self, builder: &mut AB) {
+        let main = builder.main();
+
+        let local = main.row_slice(0);
+        let is_equal_cols: &IsEqualCols<_> = (*local).borrow();
+
+        SubAir::<AB>::eval(self, builder, is_equal_cols.io, is_equal_cols.inv);
+    }
+}
+
+impl AirConfig for IsEqualChip {
+    type Cols<T> = IsEqualCols<T>;
+}
+
+// No interactions
+impl<F: Field> Chip<F> for IsEqualChip {}
+
+impl<AB: AirBuilder> SubAir<AB> for IsEqualChip {
+    type IoView = IsEqualIOCols<AB::Var>;
+    type AuxView = AB::Var;
+
+    fn eval(&self, builder: &mut AB, io: Self::IoView, inv: Self::AuxView) {
+        builder.assert_eq((io.x - io.y) * inv + io.is_equal, AB::F::one());
+        builder.assert_eq((io.x - io.y) * io.is_equal, AB::F::zero());
+    }
+}

chips/src/is_equal/columns.rs

@@ -0,0 +1,30 @@
+use afs_derive::AlignedBorrow;
+
+pub const NUM_COLS: usize = 4;
+
+#[repr(C)]
+#[derive(AlignedBorrow)]
+pub struct IsEqualCols<F> {
+    pub io: IsEqualIOCols<F>,
+    pub inv: F,
+}
+
+#[derive(Clone, Copy)]
+pub struct IsEqualIOCols<T> {
+    pub x: T,
+    pub y: T,
+    pub is_equal: T,
+}
+
+impl<T: Clone> IsEqualCols<T> {
+    pub const fn new(x: T, y: T, is_equal: T, inv: T) -> IsEqualCols<T> {
+        IsEqualCols {
+            io: IsEqualIOCols { x, y, is_equal },
+            inv,
+        }
+    }
+
+    pub fn get_width() -> usize {
+        NUM_COLS
+    }
+}

chips/src/is_equal/mod.rs

@@ -0,0 +1,17 @@
+#[cfg(test)]
+pub mod tests;
+
+pub mod air;
+pub mod columns;
+pub mod trace;
+
+use p3_field::Field;
+
+#[derive(Default)]
+pub struct IsEqualChip {}
+
+impl IsEqualChip {
+    pub fn request<F: Field>(&self, x: F, y: F) -> bool {
+        x == y
+    }
+}

chips/src/is_equal/tests.rs

@@ -0,0 +1,69 @@
+use super::IsEqualChip;
+use p3_field::AbstractField;
+
+use afs_stark_backend::{prover::USE_DEBUG_BUILDER, verifier::VerificationError};
+use afs_test_utils::config::baby_bear_poseidon2::run_simple_test_no_pis;
+
+#[test]
+fn test_single_is_equal() {
+    let x = AbstractField::from_canonical_u32(97);
+    let y = AbstractField::from_canonical_u32(97);
+
+    let chip = IsEqualChip {};
+
+    let trace = chip.generate_trace(vec![x], vec![y]);
+
+    run_simple_test_no_pis(vec![&chip], vec![trace]).expect("Verification failed");
+}
+
+#[test]
+fn test_single_is_equal2() {
+    let x = AbstractField::from_canonical_u32(127);
+    let y = AbstractField::from_canonical_u32(74);
+
+    let chip = IsEqualChip {};
+
+    let trace = chip.generate_trace(vec![x], vec![y]);
+
+    run_simple_test_no_pis(vec![&chip], vec![trace]).expect("Verification failed");
+}
+
+#[test]
+fn test_single_is_zero_fail() {
+    let x = AbstractField::from_canonical_u32(187);
+    let y = AbstractField::from_canonical_u32(123);
+
+    let chip = IsEqualChip {};
+
+    let mut trace = chip.generate_trace(vec![x], vec![y]);
+    trace.values[2] = AbstractField::from_canonical_u32(1);
+
+    USE_DEBUG_BUILDER.with(|debug| {
+        *debug.lock().unwrap() = false;
+    });
+    assert_eq!(
+        run_simple_test_no_pis(vec![&chip], vec![trace]),
+        Err(VerificationError::OodEvaluationMismatch),
+        "Expected constraint to fail"
+    );
+}
+
+#[test]
+fn test_single_is_zero_fail2() {
+    let x = AbstractField::from_canonical_u32(123);
+    let y = AbstractField::from_canonical_u32(123);
+
+    let chip = IsEqualChip {};
+
+    let mut trace = chip.generate_trace(vec![x], vec![y]);
+    trace.values[2] = AbstractField::from_canonical_u32(0);
+
+    USE_DEBUG_BUILDER.with(|debug| {
+        *debug.lock().unwrap() = false;
+    });
+    assert_eq!(
+        run_simple_test_no_pis(vec![&chip], vec![trace]),
+        Err(VerificationError::OodEvaluationMismatch),
+        "Expected constraint to fail"
+    );
+}

chips/src/is_equal/trace.rs

@@ -0,0 +1,37 @@
+use itertools::Itertools;
+use p3_field::Field;
+use p3_matrix::dense::RowMajorMatrix;
+
+use crate::sub_chip::LocalTraceInstructions;
+
+use super::{columns::IsEqualCols, IsEqualChip};
+
+impl IsEqualChip {
+    pub fn generate_trace<F: Field>(&self, x: Vec<F>, y: Vec<F>) -> RowMajorMatrix<F> {
+        let rows = x
+            .into_iter()
+            .zip_eq(y)
+            .flat_map(|(x, y)| {
+                let is_equal_cols = self.generate_trace_row((x, y));
+                [
+                    is_equal_cols.io.x,
+                    is_equal_cols.io.y,
+                    is_equal_cols.io.is_equal,
+                    is_equal_cols.inv,
+                ]
+            })
+            .collect::<Vec<_>>();
+
+        RowMajorMatrix::new(rows, IsEqualCols::<F>::get_width())
+    }
+}
+
+impl<F: Field> LocalTraceInstructions<F> for IsEqualChip {
+    type LocalInput = (F, F);
+
+    fn generate_trace_row(&self, local_input: Self::LocalInput) -> Self::Cols<F> {
+        let is_equal = self.request(local_input.0, local_input.1);
+        let inv = (local_input.0 - local_input.1 + F::from_bool(is_equal)).inverse();
+        IsEqualCols::new(local_input.0, local_input.1, F::from_bool(is_equal), inv)
+    }
+}

chips/src/is_equal_vec/air.rs

@@ -0,0 +1,77 @@
+use std::borrow::Borrow;
+
+use afs_stark_backend::interaction::Chip;
+use p3_air::AirBuilder;
+use p3_air::{Air, BaseAir};
+use p3_field::AbstractField;
+use p3_field::Field;
+use p3_matrix::Matrix;
+
+use crate::sub_chip::{AirConfig, SubAir};
+
+use super::{
+    columns::{IsEqualVecAuxCols, IsEqualVecCols, IsEqualVecIOCols},
+    IsEqualVecChip,
+};
+
+// No interactions
+impl<F: Field> Chip<F> for IsEqualVecChip {}
+
+impl AirConfig for IsEqualVecChip {
+    type Cols<T> = IsEqualVecCols<T>;
+}
+
+impl<F: Field> BaseAir<F> for IsEqualVecChip {
+    fn width(&self) -> usize {
+        self.get_width()
+    }
+}
+
+/// Imposes AIR constaints within each row
+/// Indices are as follows:
+/// 0..2*vec_len: vector input
+/// 2*vec_len..3*vec_len: cumulative equality AND (answer in index 3*vec_len-1)
+/// 3*vec_len..4*vec_len: inverse used to constrain nonzero when equality holds
+///
+/// At first index naively implements is_equal constraints
+/// At every index constrains cumulative NAND difference
+/// At every transition index prohibits 0 followed by 1, and constrains
+/// 1 with equality must be followed by 1
+/// When product does not change, inv is 0, when product changes, inverse is inverse of difference
+impl<AB: AirBuilder> Air<AB> for IsEqualVecChip {
+    fn eval(&self, builder: &mut AB) {
+        let main = builder.main();
+
+        let local = main.row_slice(0);
+        let local: &[AB::Var] = (*local).borrow();
+
+        let is_equal_vec_cols = IsEqualVecCols::from_slice(local, self.vec_len);
+
+        SubAir::<AB>::eval(self, builder, is_equal_vec_cols.io, is_equal_vec_cols.aux);
+    }
+}
+
+impl<AB: AirBuilder> SubAir<AB> for IsEqualVecChip {
+    type IoView = IsEqualVecIOCols<AB::Var>;
+    type AuxView = IsEqualVecAuxCols<AB::Var>;
+
+    fn eval(&self, builder: &mut AB, io: Self::IoView, aux: Self::AuxView) {
+        let IsEqualVecIOCols { x, y, prod: _ } = io;
+        let IsEqualVecAuxCols { prods, invs } = aux;
+        let vec_len = self.vec_len;
+        // initialize prods[0] = is_equal(x[0], y[0])
+        builder.assert_eq(prods[0] + (x[0] - y[0]) * invs[0], AB::F::one());
+
+        for i in 0..vec_len {
+            // constrain prods[i] = 0 if x[i] != y[i]
+            builder.assert_zero(prods[i] * (x[i] - y[i]));
+        }
+
+        for i in 0..vec_len - 1 {
+            // if prod[i] == 0 all after are 0
+            builder.assert_eq(prods[i] * prods[i + 1], prods[i + 1]);
+            // prods[i] == 1 forces prods[i+1] == is_equal(x[i+1], y[i+1])
+            builder.assert_eq(prods[i + 1] + (x[i + 1] - y[i + 1]) * invs[i + 1], prods[i]);
+        }
+    }
+}

chips/src/is_equal_vec/columns.rs

@@ -0,0 +1,63 @@
+#[derive(Default)]
+pub struct IsEqualVecIOCols<T> {
+    pub x: Vec<T>,
+    pub y: Vec<T>,
+    pub prod: T,
+}
+
+#[derive(Default)]
+pub struct IsEqualVecAuxCols<T> {
+    pub prods: Vec<T>,
+    pub invs: Vec<T>,
+}
+
+#[derive(Default)]
+pub struct IsEqualVecCols<T> {
+    pub io: IsEqualVecIOCols<T>,
+    pub aux: IsEqualVecAuxCols<T>,
+}
+
+impl<T: Clone> IsEqualVecCols<T> {
+    pub fn new(x: Vec<T>, y: Vec<T>, prods: Vec<T>, invs: Vec<T>) -> Self {
+        Self {
+            io: IsEqualVecIOCols {
+                x,
+                y,
+                prod: prods[prods.len() - 1].clone(),
+            },
+            aux: IsEqualVecAuxCols { prods, invs },
+        }
+    }
+
+    pub fn from_slice(slc: &[T], vec_len: usize) -> Self {
+        let x = slc[0..vec_len].to_vec();
+        let y = slc[vec_len..2 * vec_len].to_vec();
+        let prod = slc[3 * vec_len - 1].clone();
+        let prods = slc[2 * vec_len..3 * vec_len].to_vec();
+        let invs = slc[3 * vec_len..4 * vec_len].to_vec();
+
+        Self {
+            io: IsEqualVecIOCols { x, y, prod },
+            aux: IsEqualVecAuxCols { prods, invs },
+        }
+    }
+
+    pub fn flatten(&self) -> Vec<T> {
+        self.io
+            .x
+            .iter()
+            .chain(self.io.y.iter())
+            .chain(self.aux.prods.iter())
+            .chain(self.aux.invs.iter())
+            .cloned()
+            .collect()
+    }
+
+    pub fn get_width(&self) -> usize {
+        4 * self.vec_len()
+    }
+
+    pub fn vec_len(&self) -> usize {
+        self.io.x.len()
+    }
+}

chips/src/is_equal_vec/mod.rs

@@ -0,0 +1,21 @@
+#[cfg(test)]
+pub mod tests;
+
+pub mod air;
+pub mod columns;
+pub mod trace;
+
+#[derive(Default)]
+pub struct IsEqualVecChip {
+    vec_len: usize,
+}
+
+impl IsEqualVecChip {
+    pub fn request<F: Clone + PartialEq>(&self, x: &[F], y: &[F]) -> bool {
+        x == y
+    }
+
+    pub fn get_width(&self) -> usize {
+        4 * self.vec_len
+    }
+}