feat: inner join

chips/src/inner_join/final_table/bridge.rs

@@ -0,0 +1,55 @@
+use afs_stark_backend::interaction::{AirBridge, Interaction};
+use p3_air::VirtualPairCol;
+use p3_field::PrimeField;
+
+use super::FinalTableAir;
+use crate::{indexed_output_page_air::columns::IndexedOutputPageCols, utils::to_vcols};
+
+impl<F: PrimeField> AirBridge<F> for FinalTableAir {
+    /// Sends the same thing as FinalPageAir
+    fn sends(&self) -> Vec<Interaction<F>> {
+        AirBridge::sends(&self.final_air)
+    }
+
+    /// Receives (idx, data) of T1 for every allocated row on t1_output_bus (sent by t1_chip)
+    /// Receives (idx, data) of T2 for every allocated row on t2_output_bus (sent by t2_chip)
+    fn receives(&self) -> Vec<Interaction<F>> {
+        let num_cols = self.air_width();
+        let all_cols = (0..num_cols).collect::<Vec<usize>>();
+
+        let table_cols = IndexedOutputPageCols::<usize>::from_slice(
+            &all_cols,
+            self.final_air.idx_len,
+            self.final_air.data_len,
+            self.final_air.idx_limb_bits,
+            self.final_air.idx_decomp,
+        );
+
+        let t1_cols = table_cols.page_cols.data[self.fkey_start..self.fkey_end]
+            .iter()
+            .chain(table_cols.page_cols.data[self.t2_data_len..].iter())
+            .copied()
+            .collect::<Vec<usize>>();
+
+        let t2_cols = table_cols
+            .page_cols
+            .idx
+            .iter()
+            .chain(table_cols.page_cols.data[..self.t2_data_len].iter())
+            .copied()
+            .collect::<Vec<usize>>();
+
+        vec![
+            Interaction {
+                fields: to_vcols(&t1_cols),
+                count: VirtualPairCol::single_main(table_cols.page_cols.is_alloc),
+                argument_index: self.t1_output_bus_index,
+            },
+            Interaction {
+                fields: to_vcols(&t2_cols),
+                count: VirtualPairCol::single_main(table_cols.page_cols.is_alloc),
+                argument_index: self.t2_output_bus_index,
+            },
+        ]
+    }
+}

chips/src/inner_join/final_table/mod.rs

@@ -0,0 +1,101 @@
+/// This chip is mostly just FinalPageAir, but with different interactions
+/// Most of the code for it is here (just calls the correspondsing functions
+/// from FinalPageAir), but the new interactions are in bridge.rs
+use std::sync::Arc;
+
+use afs_stark_backend::air_builders::PartitionedAirBuilder;
+use p3_air::{Air, BaseAir};
+use p3_field::{Field, PrimeField};
+use p3_matrix::dense::RowMajorMatrix;
+use p3_uni_stark::{StarkGenericConfig, Val};
+
+use crate::{
+    common::page::Page, indexed_output_page_air::IndexedOutputPageAir,
+    range_gate::RangeCheckerGateChip,
+};
+
+pub mod bridge;
+
+#[derive(Clone)]
+pub(super) struct FinalTableAir {
+    t1_output_bus_index: usize,
+    t2_output_bus_index: usize,
+
+    /// Foreign key indices within the data
+    fkey_start: usize,
+    fkey_end: usize,
+
+    t2_data_len: usize,
+
+    final_air: IndexedOutputPageAir,
+}
+
+impl FinalTableAir {
+    #[allow(clippy::too_many_arguments)]
+    pub fn new(
+        t1_output_bus_index: usize,
+        t2_output_bus_index: usize,
+        range_bus_index: usize,
+        idx_len: usize,
+        t1_data_len: usize,
+        t2_data_len: usize,
+        fkey_start: usize,
+        fkey_end: usize,
+        idx_limb_bits: usize,
+        decomp: usize,
+    ) -> Self {
+        Self {
+            t1_output_bus_index,
+            t2_output_bus_index,
+            fkey_start,
+            fkey_end,
+            t2_data_len,
+            final_air: IndexedOutputPageAir::new(
+                range_bus_index,
+                idx_len,
+                t1_data_len + t2_data_len,
+                idx_limb_bits,
+                decomp,
+            ),
+        }
+    }
+
+    pub fn table_width(&self) -> usize {
+        self.final_air.page_width()
+    }
+
+    pub fn aux_width(&self) -> usize {
+        self.final_air.aux_width()
+    }
+
+    pub fn air_width(&self) -> usize {
+        self.final_air.air_width()
+    }
+
+    pub fn gen_aux_trace<SC: StarkGenericConfig>(
+        &self,
+        page: &Page,
+        range_checker: Arc<RangeCheckerGateChip>,
+    ) -> RowMajorMatrix<Val<SC>>
+    where
+        Val<SC>: PrimeField,
+    {
+        self.final_air.gen_aux_trace::<SC>(page, range_checker)
+    }
+}
+
+impl<F: Field> BaseAir<F> for FinalTableAir {
+    fn width(&self) -> usize {
+        self.air_width()
+    }
+}
+
+impl<AB: PartitionedAirBuilder> Air<AB> for FinalTableAir
+where
+    AB::M: Clone,
+{
+    fn eval(&self, builder: &mut AB) {
+        // Making sure the page is in the proper format
+        Air::eval(&self.final_air, builder);
+    }
+}

chips/src/inner_join/initial_table/air.rs

@@ -0,0 +1,32 @@
+use afs_stark_backend::air_builders::PartitionedAirBuilder;
+use p3_air::{Air, AirBuilder, BaseAir};
+use p3_field::Field;
+use p3_matrix::Matrix;
+
+use super::InitialTableAir;
+
+impl<F: Field> BaseAir<F> for InitialTableAir {
+    fn width(&self) -> usize {
+        self.air_width()
+    }
+}
+
+impl<AB: PartitionedAirBuilder> Air<AB> for InitialTableAir
+where
+    AB::M: Clone,
+{
+    fn eval(&self, builder: &mut AB) {
+        let table_trace: &<AB as AirBuilder>::M = &builder.partitioned_main()[0].clone();
+        let aux_trace: &<AB as AirBuilder>::M = &builder.partitioned_main()[1].clone();
+
+        let (table_local, aux_local) = (table_trace.row_slice(0), aux_trace.row_slice(0));
+
+        let is_alloc = table_local[0];
+        let mult_cnt = aux_local[0];
+
+        // Ensuring that mult_cnt is zero if is_alloc is zero
+        // This is important because we never want to send/receive data if
+        // the row in unallocated
+        builder.assert_eq(mult_cnt, mult_cnt * is_alloc);
+    }
+}

chips/src/inner_join/initial_table/bridge.rs

@@ -0,0 +1,91 @@
+use afs_stark_backend::interaction::{AirBridge, Interaction};
+use p3_air::VirtualPairCol;
+use p3_field::PrimeField;
+
+use crate::utils::to_vcols;
+
+use super::{columns::TableCols, InitialTableAir, TableType};
+
+impl<F: PrimeField> AirBridge<F> for InitialTableAir {
+    /// For T1:
+    /// - Sends idx (primary key) with multiplicity is_alloc on t1_intersector_bus (received by intersector_chip)
+    /// - Sends (idx, data) with multiplicity out_mult on t1_output_bus (received by output_chip)
+    /// For T2:
+    /// - Sends foreign key with multiplicity is_alloc on t2_intersector_bus (received by intersector_chip)
+    /// - Sends (idx, data) with multiplicity out_mult on t2_output_bus (received by output_chip)
+    fn sends(&self) -> Vec<Interaction<F>> {
+        let num_cols = self.air_width();
+        let all_cols = (0..num_cols).collect::<Vec<usize>>();
+
+        let table_cols = TableCols::<usize>::from_slice(&all_cols, self.idx_len, self.data_len);
+
+        match self.table_type {
+            TableType::T1 {
+                t1_intersector_bus_index,
+                t1_output_bus_index,
+            } => {
+                vec![
+                    Interaction {
+                        fields: to_vcols(&table_cols.page_cols.idx),
+                        count: VirtualPairCol::single_main(table_cols.page_cols.is_alloc),
+                        argument_index: t1_intersector_bus_index,
+                    },
+                    Interaction {
+                        fields: to_vcols(
+                            &[table_cols.page_cols.idx, table_cols.page_cols.data].concat(),
+                        ),
+                        count: VirtualPairCol::single_main(table_cols.out_mult),
+                        argument_index: t1_output_bus_index,
+                    },
+                ]
+            }
+            TableType::T2 {
+                t2_intersector_bus_index,
+                t2_output_bus_index,
+                fkey_start,
+                fkey_end,
+                ..
+            } => {
+                vec![
+                    Interaction {
+                        fields: to_vcols(&table_cols.page_cols.data[fkey_start..fkey_end]),
+                        count: VirtualPairCol::single_main(table_cols.page_cols.is_alloc),
+                        argument_index: t2_intersector_bus_index,
+                    },
+                    Interaction {
+                        fields: to_vcols(
+                            &[table_cols.page_cols.idx, table_cols.page_cols.data].concat(),
+                        ),
+                        count: VirtualPairCol::single_main(table_cols.out_mult),
+                        argument_index: t2_output_bus_index,
+                    },
+                ]
+            }
+        }
+    }
+
+    /// For T2:
+    /// - Receives foreign key with multiplicity out_mult on intersector_t2_bus (sent by intersector_chip)
+    fn receives(&self) -> Vec<Interaction<F>> {
+        let num_cols = self.air_width();
+        let all_cols = (0..num_cols).collect::<Vec<usize>>();
+
+        let table_cols = TableCols::<usize>::from_slice(&all_cols, self.idx_len, self.data_len);
+
+        if let TableType::T2 {
+            intersector_t2_bus_index,
+            fkey_start,
+            fkey_end,
+            ..
+        } = self.table_type
+        {
+            vec![Interaction {
+                fields: to_vcols(&table_cols.page_cols.data[fkey_start..fkey_end]),
+                count: VirtualPairCol::single_main(table_cols.out_mult),
+                argument_index: intersector_t2_bus_index,
+            }]
+        } else {
+            vec![]
+        }
+    }
+}

chips/src/inner_join/initial_table/columns.rs

@@ -0,0 +1,18 @@
+use crate::common::page_cols::PageCols;
+
+#[derive(Debug)]
+pub struct TableCols<T> {
+    pub page_cols: PageCols<T>,
+
+    /// The multiplicity with which we will send (idx, data) to output_chip
+    pub out_mult: T,
+}
+
+impl<T: Clone> TableCols<T> {
+    pub fn from_slice(cols: &[T], idx_len: usize, data_len: usize) -> TableCols<T> {
+        TableCols {
+            page_cols: PageCols::from_slice(&cols[..cols.len() - 1], idx_len, data_len),
+            out_mult: cols[idx_len + data_len + 1].clone(),
+        }
+    }
+}

chips/src/inner_join/initial_table/mod.rs

@@ -0,0 +1,42 @@
+pub mod air;
+pub mod bridge;
+pub mod columns;
+pub mod trace;
+
+// An enum for the different table types and their bus indices
+pub enum TableType {
+    T1 {
+        t1_intersector_bus_index: usize,
+        t1_output_bus_index: usize,
+    },
+    T2 {
+        fkey_start: usize,
+        fkey_end: usize,
+
+        t2_intersector_bus_index: usize,
+        intersector_t2_bus_index: usize,
+        t2_output_bus_index: usize,
+    },
+}
+
+#[derive(derive_new::new)]
+pub struct InitialTableAir {
+    idx_len: usize,
+    data_len: usize,
+
+    pub table_type: TableType,
+}
+
+impl InitialTableAir {
+    pub fn table_width(&self) -> usize {
+        1 + self.idx_len + self.data_len
+    }
+
+    pub fn aux_width(&self) -> usize {
+        1
+    }
+
+    pub fn air_width(&self) -> usize {
+        self.table_width() + self.aux_width()
+    }
+}

chips/src/inner_join/initial_table/trace.rs

@@ -0,0 +1,11 @@
+use afs_test_utils::utils::to_field_vec;
+use p3_field::Field;
+use p3_matrix::dense::RowMajorMatrix;
+
+use super::InitialTableAir;
+
+impl InitialTableAir {
+    pub fn gen_aux_trace<F: Field>(&self, out_mult: &[u32]) -> RowMajorMatrix<F> {
+        RowMajorMatrix::new_col(to_field_vec(out_mult.to_vec()))
+    }
+}