test: added negative tests

chips/src/is_equal/air.rs

@@ -4,7 +4,7 @@ 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, FilteredAirBuilder};
+use p3_air::{Air, AirBuilder, BaseAir};
 use p3_field::AbstractField;
 use p3_field::Field;
 use p3_matrix::Matrix;
@@ -22,12 +22,7 @@ impl<AB: AirBuilder> Air<AB> for IsEqualChip {
         let local = main.row_slice(0);
         let is_equal_cols: &IsEqualCols<_> = (*local).borrow();
 
-        SubAir::<AB>::eval(
-            self,
-            &mut builder.when(AB::Expr::one()),
-            is_equal_cols.io,
-            is_equal_cols.inv,
-        );
+        SubAir::<AB>::eval(self, builder, is_equal_cols.io, is_equal_cols.inv);
     }
 }
 
@@ -42,7 +37,7 @@ impl<AB: AirBuilder> SubAir<AB> for IsEqualChip {
     type IoView = IsEqualIOCols<AB::Var>;
     type AuxView = AB::Var;
 
-    fn eval(&self, builder: &mut FilteredAirBuilder<AB>, io: Self::IoView, inv: Self::AuxView) {
+    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_vec/air.rs

@@ -1,7 +1,7 @@
 use std::borrow::Borrow;
 
 use afs_stark_backend::interaction::Chip;
-use p3_air::{Air, AirBuilder, BaseAir, FilteredAirBuilder};
+use p3_air::{Air, AirBuilder, BaseAir};
 use p3_field::AbstractField;
 use p3_field::Field;
 use p3_matrix::Matrix;
@@ -46,20 +46,15 @@ impl<AB: AirBuilder> Air<AB> for IsEqualVecChip {
 
         let is_equal_vec_cols = IsEqualVecCols::from_slice(local, self.vec_len);
 
-        SubAir::<AB>::eval(
-            self,
-            &mut builder.when(AB::Expr::one()),
-            is_equal_vec_cols.io,
-            is_equal_vec_cols.aux,
-        );
+        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 FilteredAirBuilder<AB>, io: Self::IoView, aux: Self::AuxView) {
+    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;

chips/src/is_zero/air.rs

@@ -1,7 +1,7 @@
 use std::borrow::Borrow;
 
 use afs_stark_backend::interaction::Chip;
-use p3_air::{Air, AirBuilder, BaseAir, FilteredAirBuilder};
+use p3_air::{Air, AirBuilder, BaseAir};
 use p3_field::AbstractField;
 use p3_field::Field;
 use p3_matrix::Matrix;
@@ -33,20 +33,15 @@ impl<AB: AirBuilder> Air<AB> for IsZeroChip {
         let local = main.row_slice(0);
         let is_zero_cols: &IsZeroCols<_> = (*local).borrow();
 
-        SubAir::<AB>::eval(
-            self,
-            &mut builder.when(AB::Expr::one()),
-            is_zero_cols.io,
-            is_zero_cols.inv,
-        );
+        SubAir::<AB>::eval(self, builder, is_zero_cols.io, is_zero_cols.inv);
     }
 }
 
 impl<AB: AirBuilder> SubAir<AB> for IsZeroChip {
     type IoView = IsZeroIOCols<AB::Var>;
     type AuxView = AB::Var;
 
-    fn eval(&self, builder: &mut FilteredAirBuilder<AB>, io: Self::IoView, inv: Self::AuxView) {
+    fn eval(&self, builder: &mut AB, io: Self::IoView, inv: Self::AuxView) {
         builder.assert_eq(io.x * io.is_zero, AB::F::zero());
         builder.assert_eq(io.is_zero + io.x * inv, AB::F::one());
     }

chips/src/page_read_write/init_final_chip/trace.rs

@@ -1,5 +1,3 @@
-use std::iter;
-
 use p3_field::AbstractField;
 use p3_matrix::dense::RowMajorMatrix;
 use p3_uni_stark::{StarkGenericConfig, Val};

chips/src/page_read_write/middle_chip/air.rs

@@ -7,11 +7,7 @@ use p3_matrix::Matrix;
 
 use super::{columns::MiddleChipCols, MiddleChip};
 use crate::{
-    is_equal::columns::IsEqualIOCols,
-    is_equal_vec::{
-        columns::{IsEqualVecCols, IsEqualVecIOCols},
-        IsEqualVecChip,
-    },
+    is_equal_vec::{columns::IsEqualVecCols, IsEqualVecChip},
     sub_chip::{AirConfig, SubAir},
 };
 
@@ -41,8 +37,6 @@ where
     fn eval(&self, builder: &mut AB) {
         let main = &builder.partitioned_main()[0].clone();
 
-        // let main = builder.main();
-
         let (local, next) = (main.row_slice(0), main.row_slice(1));
         let local: &[AB::Var] = (*local).borrow();
         let next: &[AB::Var] = (*next).borrow();
@@ -68,23 +62,38 @@ where
         builder.when_first_row().assert_zero(local_cols.same_key);
         builder.when_first_row().assert_zero(local_cols.same_val);
 
-        let is_equal_cols_vec = local_cols.page_row[1..self.key_len + 1]
+        // Making sure same_key is correct across rows
+        let is_equal_keys_vec = local_cols.page_row[1..self.key_len + 1]
             .to_vec()
             .into_iter()
             .chain(next_cols.page_row[1..self.key_len + 1].to_vec())
             .chain(next_cols.is_equal_key_aux.flatten())
             .collect::<Vec<AB::Var>>();
+        let is_equal_keys = IsEqualVecCols::from_slice(&is_equal_keys_vec, self.key_len);
+        let is_equal_keys_chip = IsEqualVecChip::new(self.key_len);
 
-        let is_equal_cols = IsEqualVecCols::from_slice(&is_equal_cols_vec, self.key_len);
+        SubAir::eval(
+            &is_equal_keys_chip,
+            &mut builder.when_transition(),
+            is_equal_keys.io,
+            is_equal_keys.aux,
+        );
 
-        let is_equal_vec_chip = IsEqualVecChip::new(self.key_len);
+        // Making sure same_val is correct across rows
+        let is_equal_vals_vec = local_cols.page_row[self.key_len + 1..]
+            .to_vec()
+            .into_iter()
+            .chain(next_cols.page_row[self.key_len + 1..].to_vec())
+            .chain(next_cols.is_equal_val_aux.flatten())
+            .collect::<Vec<AB::Var>>();
+        let is_equal_vals = IsEqualVecCols::from_slice(&is_equal_vals_vec, self.val_len);
+        let is_equal_vals_chip = IsEqualVecChip::new(self.val_len);
 
-        // TODO: make sure same_key and same_val are correct for the rest of the rows
         SubAir::eval(
-            &is_equal_vec_chip,
+            &is_equal_vals_chip,
             &mut builder.when_transition(),
-            is_equal_cols.io,
-            is_equal_cols.aux,
+            is_equal_vals.io,
+            is_equal_vals.aux,
         );
 
         // TODO: make sure all rows are sorted
@@ -127,7 +136,6 @@ where
 
         // Making sure that every read uses the same value as the last operation
         // read => same_val
-        println!("local_cols: {:?}", local_cols);
         builder.assert_one(or(
             local_cols.is_extra.into(),
             or(local_cols.op_type.into(), local_cols.same_val.into()),

chips/src/page_read_write/middle_chip/chip.rs

@@ -42,7 +42,7 @@ impl<F: PrimeField64> Chip<F> for MiddleChip {
         let num_cols = self.air_width();
         let all_cols = (0..num_cols).collect::<Vec<usize>>();
 
-        let cols_to_receive = MiddleChipCols::<F>::cols_numbered(
+        let cols_to_receive = MiddleChipCols::<usize>::from_slice(
             &all_cols,
             self.page_width(),
             self.key_len,
@@ -55,7 +55,7 @@ impl<F: PrimeField64> Chip<F> for MiddleChip {
         let num_cols = self.air_width();
         let all_cols = (0..num_cols).collect::<Vec<usize>>();
 
-        let cols_to_send = MiddleChipCols::<F>::cols_numbered(
+        let cols_to_send = MiddleChipCols::<usize>::from_slice(
             &all_cols,
             self.page_width(),
             self.key_len,

chips/src/page_read_write/middle_chip/columns.rs

@@ -1,22 +1,22 @@
 use crate::is_equal_vec::columns::IsEqualVecAuxCols;
 
 // TODO: maybe separate into two structs, io and aux
-// TODO: add comments on what those bits do
 #[derive(Debug)]
 pub struct MiddleChipCols<T> {
-    pub is_initial: T,
-    pub is_final: T,
-    pub clk: T,
-    pub page_row: Vec<T>,
-    pub op_type: T, // 0 for read, 1 for write
+    pub is_initial: T, // this bit indicates if this row comes from the initial page
+    pub is_final: T, // this bit indicates if this row should go to the final page (last row for the key)
 
-    pub same_key: T,
-    pub same_val: T,
+    pub clk: T,           // timestamp for the operation
+    pub page_row: Vec<T>, // the row of the page: starts with is_alloc bit, then key, then val
+    pub op_type: T,       // 0 for read, 1 for write
+
+    pub same_key: T, // this bit indicates if the key matches the one in the previous row (should be 0 in first row)
+    pub same_val: T, // this bit indicates if the val matches the one in the previous row (should be 0 in first row)
 
     pub is_extra: T, // a bit to indicate if this is an extra row that should be ignored
 
-    pub is_equal_key_aux: IsEqualVecAuxCols<T>,
-    pub is_equal_val_aux: IsEqualVecAuxCols<T>,
+    pub is_equal_key_aux: IsEqualVecAuxCols<T>, // auxiliary columns used for same_key
+    pub is_equal_val_aux: IsEqualVecAuxCols<T>, // auxiliary columns used for same_val
 }
 
 impl<T> MiddleChipCols<T>
@@ -74,7 +74,6 @@ where
         flattened
     }
 
-    // TODO: is there a way to merge from_slice and cols_numbered?
     pub fn from_slice(slc: &[T], page_width: usize, key_len: usize, val_len: usize) -> Self {
         Self {
             is_initial: slc[0].clone(),
@@ -95,35 +94,4 @@ where
             ),
         }
     }
-
-    pub fn cols_numbered(
-        cols: &[usize],
-        page_width: usize,
-        key_len: usize,
-        val_len: usize,
-    ) -> MiddleChipCols<usize> {
-        println!("cols numbered: {} {} {}", page_width, key_len, val_len);
-        println!("cols_len: {}", cols.len());
-
-        MiddleChipCols {
-            is_initial: cols[0],
-            is_final: cols[1],
-            clk: cols[2],
-            page_row: cols[3..3 + page_width].to_vec(),
-            op_type: cols[3 + page_width],
-            same_key: cols[4 + page_width],
-            same_val: cols[5 + page_width],
-            is_extra: cols[6 + page_width],
-            is_equal_key_aux: IsEqualVecAuxCols::new(
-                cols[7 + page_width..7 + page_width + key_len].to_vec(),
-                cols[7 + page_width + key_len..7 + page_width + 2 * key_len].to_vec(),
-            ),
-            is_equal_val_aux: IsEqualVecAuxCols::new(
-                cols[7 + page_width + 2 * key_len..7 + page_width + 2 * key_len + val_len].to_vec(),
-                cols[7 + page_width + 2 * key_len + val_len
-                    ..7 + page_width + 2 * key_len + 2 * val_len]
-                    .to_vec(),
-            ),
-        }
-    }
 }

chips/src/page_read_write/middle_chip/mod.rs

@@ -1,5 +1,3 @@
-use super::page_controller::Operation;
-
 mod air;
 mod chip;
 mod columns;

chips/src/page_read_write/middle_chip/trace.rs

@@ -53,7 +53,8 @@ impl MiddleChip {
                 .collect::<Vec<Val<SC>>>()
         };
 
-        let gen_row = |page: &mut Vec<Vec<u32>>,
+        let gen_row = |is_first_row: &mut bool,
+                       page: &mut Vec<Vec<u32>>,
                        index: usize,
                        is_initial: u8,
                        is_final: u8,
@@ -72,6 +73,17 @@ impl MiddleChip {
             let cur_key = page[index][1..self.key_len + 1].to_vec();
             let cur_val = page[index][self.key_len + 1..].to_vec();
 
+            if *is_first_row {
+                // Making sure the last_key and last_val are different from current when its the first row
+                *last_key = cur_key.clone();
+                *last_val = cur_val.clone();
+
+                last_key[0] += 1;
+                last_val[0] += 1;
+
+                *is_first_row = false;
+            }
+
             let my_last_key = last_key.clone();
             let my_last_val = last_val.clone();
 
@@ -118,9 +130,9 @@ impl MiddleChip {
 
         let mut rows = vec![];
 
-        // TODO: do I need to make sure this is different than first actual row?
-        let mut last_key = vec![4234233; self.key_len];
-        let mut last_val = vec![3534543; self.val_len];
+        let mut last_key = vec![0; self.key_len];
+        let mut last_val = vec![0; self.val_len];
+        let mut is_first_row = true;
 
         let mut i = 0;
         while i < ops.len() {
@@ -137,6 +149,7 @@ impl MiddleChip {
                 idx = *key_index_map.get(&cur_key).unwrap();
 
                 rows.push(gen_row(
+                    &mut is_first_row,
                     page,
                     idx,
                     1,
@@ -161,6 +174,7 @@ impl MiddleChip {
                     .collect();
 
                 rows.push(gen_row(
+                    &mut is_first_row,
                     page,
                     idx,
                     0,
@@ -175,6 +189,7 @@ impl MiddleChip {
 
             // Adding the is_final row to the trace
             rows.push(gen_row(
+                &mut is_first_row,
                 page,
                 idx,
                 0,
@@ -195,6 +210,7 @@ impl MiddleChip {
         // Adding rows to the trace to make the height trace_degree
         while rows.len() < trace_degree {
             rows.push(gen_row(
+                &mut is_first_row,
                 page,
                 0,
                 0,

chips/src/page_read_write/mod.rs

@@ -1,6 +1,6 @@
-mod init_final_chip;
-mod middle_chip;
-mod page_controller;
+pub mod init_final_chip;
+pub mod middle_chip;
+pub mod page_controller;
 
 #[cfg(test)]
 pub mod tests;

chips/src/page_read_write/page_controller/mod.rs

@@ -4,7 +4,6 @@ use p3_field::AbstractField;
 use p3_matrix::dense::{DenseMatrix, RowMajorMatrix};
 use p3_matrix::Matrix;
 use p3_uni_stark::{StarkGenericConfig, Val};
-use std::collections::HashMap;
 
 use super::init_final_chip::InitFinalChip;
 use super::middle_chip::MiddleChip;
@@ -52,11 +51,11 @@ where
 }
 
 impl<SC: StarkGenericConfig> PageController<SC> {
-    pub fn new(bus_index: usize) -> Self {
+    pub fn new(bus_index: usize, key_len: usize, val_len: usize) -> Self {
         Self {
-            init_chip: InitFinalChip::new(bus_index, 1, true),
-            middle_chip: MiddleChip::new(bus_index, 1, 1),
-            final_chip: InitFinalChip::new(bus_index, 1, false),
+            init_chip: InitFinalChip::new(bus_index, 1 + key_len + val_len, true),
+            middle_chip: MiddleChip::new(bus_index, key_len, val_len),
+            final_chip: InitFinalChip::new(bus_index, 1 + key_len + val_len, false),
 
             init_chip_trace: None,
             middle_chip_trace: None,