XorBitsAir

chips/src/xor_bits/air.rs

@@ -9,20 +9,20 @@ use crate::sub_chip::{AirConfig, SubAir};
 
 use super::{
     columns::{XorBitCols, XorCols, XorIOCols},
-    XorBitsAir, XorBitsChip,
+    XorBitsAir,
 };
 
-impl<const N: usize> AirConfig for XorBitsChip<N> {
+impl<const N: usize> AirConfig for XorBitsAir<N> {
     type Cols<T> = XorCols<N, T>;
 }
 
-impl<F: Field, const N: usize> BaseAir<F> for XorBitsChip<N> {
+impl<F: Field, const N: usize> BaseAir<F> for XorBitsAir<N> {
     fn width(&self) -> usize {
         XorCols::<N, F>::get_width()
     }
 }
 
-impl<AB: AirBuilder, const N: usize> Air<AB> for XorBitsChip<N> {
+impl<AB: AirBuilder, const N: usize> Air<AB> for XorBitsAir<N> {
     fn eval(&self, builder: &mut AB) {
         let main = builder.main();
 
@@ -31,7 +31,7 @@ impl<AB: AirBuilder, const N: usize> Air<AB> for XorBitsChip<N> {
 
         let xor_cols = XorCols::<N, AB::Var>::from_slice(local);
 
-        SubAir::eval(&self.air, builder, xor_cols.io, xor_cols.bits);
+        SubAir::eval(self, builder, xor_cols.io, xor_cols.bits);
     }
 }
 

chips/src/xor_bits/chip.rs → chips/src/xor_bits/bridge.rs

@@ -4,9 +4,9 @@ use p3_field::{Field, PrimeField64};
 
 use crate::sub_chip::SubAirBridge;
 
-use super::{columns::XorCols, XorBitsChip};
+use super::{columns::XorCols, XorBitsAir};
 
-impl<F: PrimeField64, const N: usize> AirBridge<F> for XorBitsChip<N> {
+impl<F: PrimeField64, const N: usize> AirBridge<F> for XorBitsAir<N> {
     fn receives(&self) -> Vec<Interaction<F>> {
         let num_cols = XorCols::<N, F>::get_width();
         let indices = (0..num_cols).collect::<Vec<usize>>();
@@ -16,7 +16,7 @@ impl<F: PrimeField64, const N: usize> AirBridge<F> for XorBitsChip<N> {
     }
 }
 
-impl<F: Field, const N: usize> SubAirBridge<F> for XorBitsChip<N> {
+impl<F: Field, const N: usize> SubAirBridge<F> for XorBitsAir<N> {
     fn receives(&self, col_indices: XorCols<N, usize>) -> Vec<Interaction<F>> {
         let io_indices = col_indices.io;
         vec![Interaction {
@@ -26,7 +26,7 @@ impl<F: Field, const N: usize> SubAirBridge<F> for XorBitsChip<N> {
                 VirtualPairCol::single_main(io_indices.z),
             ],
             count: VirtualPairCol::constant(F::one()),
-            argument_index: self.bus_index(),
+            argument_index: self.bus_index,
         }]
     }
 }

chips/src/xor_bits/mod.rs

@@ -1,24 +1,28 @@
 use parking_lot::Mutex;
 
 pub mod air;
-pub mod chip;
+pub mod bridge;
 pub mod columns;
 pub mod trace;
 
-// We separate XorBitsAir and XorBitsChip because we want a struct that specifies just the constraints of XOR,
-// while the chip needs some additional fields to receive interactions.
-
 /// AIR that computes the xor of two numbers of at most N bits each.
 /// This struct only implements SubAir.
 #[derive(Default)]
-pub struct XorBitsAir<const N: usize>;
+pub struct XorBitsAir<const N: usize> {
+    pub bus_index: usize,
+}
+
+impl<const N: usize> XorBitsAir<N> {
+    pub fn calc_xor(&self, a: u32, b: u32) -> u32 {
+        a ^ b
+    }
+}
 
 #[derive(Default)]
 /// A chip that computes the xor of two numbers of at most N bits each.
 /// This chip consists of the AIR as well as a receiver to handle counting requests.
 pub struct XorBitsChip<const N: usize> {
-    air: XorBitsAir<N>,
-    bus_index: usize,
+    pub air: XorBitsAir<N>,
 
     /// List of all requests sent to the chip
     pairs: Mutex<Vec<(u32, u32)>>,
@@ -27,23 +31,14 @@ pub struct XorBitsChip<const N: usize> {
 impl<const N: usize> XorBitsChip<N> {
     pub fn new(bus_index: usize, pairs: Vec<(u32, u32)>) -> Self {
         Self {
-            air: XorBitsAir,
-            bus_index,
+            air: XorBitsAir { bus_index },
             pairs: Mutex::new(pairs),
         }
     }
 
-    pub fn bus_index(&self) -> usize {
-        self.bus_index
-    }
-
-    fn calc_xor(&self, a: u32, b: u32) -> u32 {
-        a ^ b
-    }
-
     pub fn request(&self, a: u32, b: u32) -> u32 {
         let mut pairs_locked = self.pairs.lock();
         pairs_locked.push((a, b));
-        self.calc_xor(a, b)
+        self.air.calc_xor(a, b)
     }
 }

chips/src/xor_bits/trace.rs

@@ -5,7 +5,7 @@ use crate::sub_chip::LocalTraceInstructions;
 
 use super::{
     columns::{XorBitCols, XorCols, XorIOCols},
-    XorBitsChip,
+    XorBitsAir, XorBitsChip,
 };
 
 impl<const N: usize> XorBitsChip<N> {
@@ -17,14 +17,14 @@ impl<const N: usize> XorBitsChip<N> {
 
         let rows = pairs_locked
             .iter()
-            .flat_map(|(x, y)| self.generate_trace_row((*x, *y)).flatten())
+            .flat_map(|(x, y)| self.air.generate_trace_row((*x, *y)).flatten())
             .collect();
 
         RowMajorMatrix::new(rows, num_xor_cols)
     }
 }
 
-impl<const N: usize, F: AbstractField> LocalTraceInstructions<F> for XorBitsChip<N> {
+impl<const N: usize, F: AbstractField> LocalTraceInstructions<F> for XorBitsAir<N> {
     /// The input is (x, y) to be XOR-ed.
     type LocalInput = (u32, u32);
 

chips/tests/integration_test.rs

@@ -115,7 +115,7 @@ fn test_xor_bits_chip() {
     for requester in &requesters {
         all_chips.push(requester);
     }
-    all_chips.push(&*xor_chip);
+    all_chips.push(&xor_chip.air);
 
     let all_traces = requesters_traces
         .into_iter()
@@ -165,7 +165,7 @@ fn negative_test_xor_bits_chip() {
     );
 
     let result = run_simple_test_no_pis(
-        vec![&dummy_requester, &*xor_chip],
+        vec![&dummy_requester, &xor_chip.air],
         vec![dummy_trace, xor_chip_trace],
     );