openvm-org
diff --git a/‎chips/src/common/page.rs
Lines changed: 4 additions & 0 deletions b/‎chips/src/common/page.rs
Lines changed: 4 additions & 0 deletions
diff --git a/‎chips/src/group_by/group_by_input/air.rs
Lines changed: 139 additions & 0 deletions b/‎chips/src/group_by/group_by_input/air.rs
Lines changed: 139 additions & 0 deletions
diff --git a/‎chips/src/group_by/group_by_input/bridge.rs
Lines changed: 93 additions & 0 deletions b/‎chips/src/group_by/group_by_input/bridge.rs
Lines changed: 93 additions & 0 deletions
diff --git a/‎chips/src/group_by/group_by_input/columns.rs
Lines changed: 145 additions & 0 deletions b/‎chips/src/group_by/group_by_input/columns.rs
Lines changed: 145 additions & 0 deletions
@@ -37,6 +37,10 @@ impl Page {
         }
     }
 
+    pub fn to_2d_vec(&self) -> Vec<Vec<u32>> {
+        self.rows.iter().map(|row| row.to_vec()).collect()
+    }
+
     /// Returns a random page with the given parameters in the proper format
     /// Note that max_idx and max_data are not inclusive
     pub fn random(
 
@@ -0,0 +1,139 @@
+use std::iter;
+
+use afs_stark_backend::air_builders::PartitionedAirBuilder;
+use p3_air::{Air, AirBuilder, BaseAir};
+use p3_field::Field;
+use p3_matrix::Matrix;
+
+use crate::common::page_cols::PageCols;
+use crate::is_equal_vec::columns::IsEqualVecIOCols;
+use crate::sub_chip::{AirConfig, SubAir};
+
+use super::columns::{GroupByAuxCols, GroupByCols, GroupedPageCols};
+use super::GroupByAir;
+
+impl<F: Field> BaseAir<F> for GroupByAir {
+    fn width(&self) -> usize {
+        self.get_width()
+    }
+}
+
+impl<AB: PartitionedAirBuilder> Air<AB> for GroupByAir
+where
+    AB::M: Clone,
+{
+    /// Re-references builder into page_trace and aux_trace, then slices into local and next rows
+    /// to evaluate using SubAir::eval(GroupByAir)
+    fn eval(&self, builder: &mut AB) {
+        let page_trace: &<AB as AirBuilder>::M = &builder.partitioned_main()[0];
+        let aux_trace: &<AB as AirBuilder>::M = &builder.partitioned_main()[1];
+
+        // get the current row and the next row
+        let (local_page, next_page) = (page_trace.row_slice(0), page_trace.row_slice(1));
+        let (local_aux, next_aux) = (aux_trace.row_slice(0), aux_trace.row_slice(1));
+
+        let local_cols = GroupByCols::from_partitioned_slice(&local_page, &local_aux, self);
+        let next_cols = GroupByCols::from_partitioned_slice(&next_page, &next_aux, self);
+        drop((local_page, next_page, local_aux, next_aux));
+
+        SubAir::eval(
+            self,
+            builder,
+            (local_cols.page, next_cols.page),
+            (local_cols.aux, next_cols.aux),
+        );
+    }
+}
+
+impl AirConfig for GroupByAir {
+    type Cols<T> = GroupByCols<T>;
+}
+
+impl<AB: PartitionedAirBuilder> SubAir<AB> for GroupByAir {
+    /// `io` consists of `(local_page, next_page)` only the page, including `is_alloc`
+    type IoView = (PageCols<AB::Var>, PageCols<AB::Var>);
+    /// `aux.0` is `local.aux`, `aux.1` is `next.aux`.
+    ///
+    /// `aux` consists of everything that isn't `io`, including
+    /// `sorted_group_by`, `sorted_group_by_alloc`, `aggregated`, and `partial_aggregated`
+    type AuxView = (GroupByAuxCols<AB::Var>, GroupByAuxCols<AB::Var>);
+
+    /// Constrains `sorted_group_by` along with `partial_aggregated` to hold correct values
+    /// with minimal constraints.
+    ///
+    /// In fact `sorted_group_by` is not necessarily sorted. The only constraints are that
+    /// allocated rows are placed at the beginning, and like rows are placed together.
+    ///
+    /// Like rows being placed together is enforced by the constraints on `MyFinalPage`, since
+    /// all rows marked `final` are sent to MyFinalPage and hence must be pairwise distinct.
+    fn eval(
+        &self,
+        builder: &mut AB,
+        (local_page, next_page): Self::IoView,
+        (local_aux, next_aux): Self::AuxView,
+    ) {
+        // Return either grouped columns (if not None) or the respective columns from the page
+        // (the latter occurs if the page is already assumed to be grouped appropriated)
+        let get_grouped = |page: PageCols<AB::Var>, grouped: Option<GroupedPageCols<AB::Var>>| {
+            grouped.unwrap_or_else(|| GroupedPageCols {
+                is_alloc: page.is_alloc,
+                group_by: self.group_by_cols.iter().map(|&i| page.data[i]).collect(),
+                to_aggregate: page.data[self.aggregated_col],
+            })
+        };
+        let local_grouped = get_grouped(local_page, local_aux.grouped);
+        let next_grouped = get_grouped(next_page, next_aux.grouped);
+
+        let is_equal_io = IsEqualVecIOCols {
+            x: iter::once(local_grouped.is_alloc)
+                .chain(local_grouped.group_by)
+                .collect(),
+            y: iter::once(next_grouped.is_alloc)
+                .chain(next_grouped.group_by)
+                .collect(),
+            prod: local_aux.eq_next,
+        };
+
+        // constrain eq_next to hold the correct value
+        SubAir::eval(
+            &self.is_equal_vec_air,
+            &mut builder.when_transition(),
+            is_equal_io,
+            local_aux.is_equal_vec_aux,
+        );
+
+        // if grouped is_alloc changes, then is_final must be 1, even if eq_next is also 1
+        let grouped_is_alloc_diff = local_grouped.is_alloc - next_grouped.is_alloc;
+        builder.when_transition().assert_eq(
+            grouped_is_alloc_diff.clone(),
+            grouped_is_alloc_diff * local_aux.is_final,
+        );
+
+        // constrain is_final to be 1 iff eq_next == false and local_grouped.is_alloc is 1
+        builder.assert_eq(
+            local_aux.is_final,
+            local_grouped.is_alloc - local_grouped.is_alloc * local_aux.eq_next,
+        );
+
+        // constrain last vector equality to 0
+        // because previously only constrained on transition
+        builder.when_last_row().assert_zero(local_aux.eq_next);
+
+        // initialize partial sum at first row
+        builder
+            .when_first_row()
+            .assert_eq(local_aux.partial_aggregated, local_grouped.to_aggregate);
+
+        // constrain partials to sum correctly
+        builder.when_transition().assert_eq(
+            next_aux.partial_aggregated,
+            local_aux.eq_next * local_aux.partial_aggregated + next_grouped.to_aggregate,
+        );
+
+        // constrain allocated rows come first
+        builder.when_transition().assert_eq(
+            next_grouped.is_alloc * local_grouped.is_alloc,
+            next_grouped.is_alloc,
+        );
+    }
+}
@@ -0,0 +1,93 @@
+use std::iter;
+
+use crate::sub_chip::SubAirBridge;
+use afs_stark_backend::interaction::{AirBridge, Interaction};
+use p3_air::VirtualPairCol;
+use p3_field::PrimeField64;
+
+use super::columns::GroupByCols;
+use super::GroupByAir;
+
+impl<F: PrimeField64> AirBridge<F> for GroupByAir {
+    fn sends(&self) -> Vec<Interaction<F>> {
+        let col_indices_vec: Vec<usize> = (0..self.get_width()).collect();
+        let col_indices = GroupByCols::from_slice(&col_indices_vec, self);
+        SubAirBridge::sends(self, col_indices)
+    }
+
+    fn receives(&self) -> Vec<Interaction<F>> {
+        let col_indices_vec: Vec<usize> = (0..self.get_width()).collect();
+        let col_indices = GroupByCols::from_slice(&col_indices_vec, self);
+        SubAirBridge::receives(self, col_indices)
+    }
+}
+
+impl<F: PrimeField64> SubAirBridge<F> for GroupByAir {
+    /// Sends desired columns (group_by and to_aggregate) from input page internally with count
+    /// `is_alloc`, and sends answer columns with count `is_final`.
+    fn sends(&self, col_indices: GroupByCols<usize>) -> Vec<Interaction<F>> {
+        assert_eq!(col_indices.aux.grouped.is_some(), !self.sorted);
+        let group_by_col_indices = if let Some(grouped) = col_indices.aux.grouped {
+            grouped.group_by
+        } else {
+            self.group_by_cols
+                .iter()
+                .map(|&i| col_indices.page.data[i])
+                .collect()
+        };
+        // fields = group_by cols, partial_aggregated
+        // count = is_final
+        let output_sent_fields = group_by_col_indices
+            .into_iter()
+            .chain(iter::once(col_indices.aux.partial_aggregated))
+            .map(VirtualPairCol::single_main)
+            .collect();
+        let output_count = VirtualPairCol::single_main(col_indices.aux.is_final);
+        let mut interactions = vec![Interaction {
+            fields: output_sent_fields,
+            count: output_count,
+            argument_index: self.output_bus,
+        }];
+        if !self.sorted {
+            // Must do internal grouping of page based on group_by columns
+            // Sends from columns in input page to internal bus
+            let internal_sent_fields = self
+                .group_by_cols
+                .iter()
+                .chain(iter::once(&self.aggregated_col))
+                .map(|&i| VirtualPairCol::single_main(col_indices.page.data[i]))
+                .collect();
+            let internal_count = VirtualPairCol::single_main(col_indices.page.is_alloc);
+
+            interactions.push(Interaction {
+                fields: internal_sent_fields,
+                count: internal_count,
+                argument_index: self.internal_bus,
+            });
+        }
+        interactions
+    }
+
+    /// Receives desired columns (`sorted_group_by` and `aggregated`) internally with count
+    /// `is_alloc`.
+    fn receives(&self, col_indices: GroupByCols<usize>) -> Vec<Interaction<F>> {
+        if self.sorted {
+            return vec![];
+        }
+
+        let grouped = col_indices.aux.grouped.unwrap();
+        let internal_received_fields: Vec<VirtualPairCol<F>> = grouped
+            .group_by
+            .into_iter()
+            .chain(iter::once(grouped.to_aggregate))
+            .map(VirtualPairCol::single_main)
+            .collect();
+        let internal_count = VirtualPairCol::single_main(grouped.is_alloc);
+
+        vec![Interaction {
+            fields: internal_received_fields,
+            count: internal_count,
+            argument_index: self.internal_bus,
+        }]
+    }
+}
@@ -0,0 +1,145 @@
+use super::GroupByAir;
+use crate::{common::page_cols::PageCols, is_equal_vec::columns::IsEqualVecAuxCols};
+use std::ops::Range;
+
+/// Since `GroupByChip` contains a `LessThanChip` subchip and an `IsEqualVecChip` subchip, a subset of
+/// the columns are those of the `LessThanChip` and `IsEqualVecChip`.
+///
+/// The `io` columns consist only of the cached page, because output is sent to `MyFinalPage`. The
+/// `aux` columns are all other columns.
+///
+/// Implements two methods:
+///
+/// * `from_slice`: Takes a slice and returns a `GroupByCols` struct.
+/// * `index_map`: Returns a `GroupByColsIndexMap` struct, used to index all other structs and
+///   defines the order of segments in a slice.
+pub struct GroupByCols<T> {
+    /// Page with `idx_len = 0` because the page does not need an index.
+    pub page: PageCols<T>,
+    pub aux: GroupByAuxCols<T>,
+}
+
+/// The `aux` columns are all non-cached columns.
+pub struct GroupByAuxCols<T> {
+    /// If page is not already grouped, extra columns are needed to do the grouping.
+    pub grouped: Option<GroupedPageCols<T>>,
+    pub partial_aggregated: T,
+    pub is_final: T,
+    pub eq_next: T,
+    pub is_equal_vec_aux: IsEqualVecAuxCols<T>,
+}
+
+/// The columns that are relevant to the GroupBy operation.
+pub struct GroupedPageCols<T> {
+    pub is_alloc: T,
+    pub group_by: Vec<T>,
+    pub to_aggregate: T,
+}
+
+/// Maps parts of the `GroupByCols` to their indices. Note that `sorted_group_by_combined_range` is
+/// a range containing `sorted_group_by_alloc` and `sorted_group_by_range`. Indexes by the
+/// respective partition, not by the complete row.
+pub struct GroupByColsIndexMap {
+    pub allocated_idx: usize,
+    pub page_range: Range<usize>,
+    pub sorted_group_by_alloc: usize,
+    pub sorted_group_by_range: Range<usize>,
+    pub sorted_group_by_combined_range: Range<usize>,
+    pub to_aggregate: usize,
+    pub partial_aggregated: usize,
+    pub is_final: usize,
+    pub eq_next: usize,
+    pub is_equal_vec_aux_range: Range<usize>,
+}
+
+impl<T: Clone> GroupByCols<T> {
+    /// Takes a slice and returns a `GroupByCols` struct.
+    pub fn from_slice(slc: &[T], group_by_air: &GroupByAir) -> Self {
+        assert_eq!(slc.len(), group_by_air.get_width());
+        Self::from_partitioned_slice(
+            &slc[..group_by_air.page_width],
+            &slc[group_by_air.page_width..],
+            group_by_air,
+        )
+    }
+
+    pub fn from_partitioned_slice(page: &[T], aux: &[T], group_by_air: &GroupByAir) -> Self {
+        assert_eq!(page.len(), group_by_air.page_width);
+        assert_eq!(
+            aux.len(),
+            group_by_air.get_width() - group_by_air.page_width
+        );
+
+        let index_map = GroupByCols::<T>::index_map(group_by_air);
+
+        let grouped = if !group_by_air.sorted {
+            Some(GroupedPageCols {
+                is_alloc: aux[index_map.sorted_group_by_alloc].clone(),
+                group_by: aux[index_map.sorted_group_by_range].to_vec(),
+                to_aggregate: aux[index_map.to_aggregate].clone(),
+            })
+        } else {
+            None
+        };
+
+        let partial_aggregated = aux[index_map.partial_aggregated].clone();
+        let is_final = aux[index_map.is_final].clone();
+        let eq_next = aux[index_map.eq_next].clone();
+        let is_equal_vec_aux = IsEqualVecAuxCols::from_slice(
+            &aux[index_map.is_equal_vec_aux_range],
+            group_by_air.group_by_cols.len() + 1,
+        );
+
+        let data_len = group_by_air.page_width - 1;
+        Self {
+            page: PageCols::from_slice(page, 0, data_len),
+            aux: GroupByAuxCols {
+                grouped,
+                partial_aggregated,
+                is_final,
+                eq_next,
+                is_equal_vec_aux,
+            },
+        }
+    }
+
+    /// Returns a `GroupByColsIndexMap` struct, used to index all other structs and defines the
+    /// order of segments in a slice. Indexes by the respective partition, not by the complete row.
+    pub fn index_map(group_by_air: &GroupByAir) -> GroupByColsIndexMap {
+        let num_group_by = group_by_air.group_by_cols.len();
+        let eq_vec_width = IsEqualVecAuxCols::<T>::get_width(num_group_by + 1);
+
+        let allocated_idx = 0;
+        let page_range = if !group_by_air.sorted {
+            allocated_idx + 1..group_by_air.page_width
+        } else {
+            0..0
+        };
+        let sorted_group_by_alloc = 0;
+        let sorted_group_by_range =
+            sorted_group_by_alloc + 1..sorted_group_by_alloc + 1 + num_group_by;
+        let sorted_group_by_combined_range = sorted_group_by_alloc..sorted_group_by_range.end;
+        let aggregated_idx = sorted_group_by_range.end;
+        let partial_aggregated_idx = if !group_by_air.sorted {
+            aggregated_idx + 1
+        } else {
+            0
+        };
+        let is_final_idx = partial_aggregated_idx + 1;
+        let eq_next_idx = is_final_idx + 1;
+        let is_equal_vec_aux_range = eq_next_idx + 1..eq_next_idx + 1 + eq_vec_width;
+
+        GroupByColsIndexMap {
+            allocated_idx,
+            page_range,
+            sorted_group_by_alloc,
+            sorted_group_by_range,
+            sorted_group_by_combined_range,
+            to_aggregate: aggregated_idx,
+            partial_aggregated: partial_aggregated_idx,
+            is_final: is_final_idx,
+            eq_next: eq_next_idx,
+            is_equal_vec_aux_range,
+        }
+    }
+}
Original file line number	Diff line number	Diff line change
`@@ -37,6 +37,10 @@ impl Page {`
`37`	`37`	`}`
`38`	`38`	`}`
`39`	`39`
	`40`	`+ pub fn to_2d_vec(&self) -> Vec<Vec<u32>> {`
	`41`	`+ self.rows.iter().map(\|row\| row.to_vec()).collect()`
	`42`	`+ }`
	`43`	`+`
`40`	`44`	`/// Returns a random page with the given parameters in the proper format`
`41`	`45`	`/// Note that max_idx and max_data are not inclusive`
`42`	`46`	`pub fn random(`