Skip to content

Commit

Permalink
chore: IndexScanPageAir use from_partitioned_slice (#97)
Browse files Browse the repository at this point in the history
* chore: IndexScanPageAir use from_partitioned_slice

* chore: directly use local_page and local_aux

* chore: address comments

* chore: use tip
  • Loading branch information
bfan05 authored and Alex Zhao committed Jun 28, 2024
1 parent 6d58ac3 commit 41bb88b
Show file tree
Hide file tree
Showing 3 changed files with 76 additions and 47 deletions.
22 changes: 10 additions & 12 deletions chips/src/single_page_index_scan/page_index_scan_input/air.rs
Original file line number Diff line number Diff line change
Expand Up @@ -81,12 +81,6 @@ where

let local_page = page_main.row_slice(0);
let local_aux = aux_main.row_slice(0);
let local_vec = local_page
.iter()
.chain(local_aux.iter())
.cloned()
.collect::<Vec<AB::Var>>();
let local = local_vec.as_slice();

// get the idx_limb_bits and decomp, which will be used to generate local_cols
let (idx_limb_bits, decomp) = match &self.variant_air {
Expand Down Expand Up @@ -121,8 +115,12 @@ where
PageIndexScanInputAirVariants::Eq(..) => Comp::Eq,
};

let local_cols = PageIndexScanInputCols::<AB::Var>::from_slice(
local,
let PageIndexScanInputCols {
page_cols,
local_cols,
} = PageIndexScanInputCols::<AB::Var>::from_partitioned_slice(
&local_page,
&local_aux,
self.idx_len,
self.data_len,
idx_limb_bits.clone(),
Expand All @@ -136,7 +134,7 @@ where
}
// constrain that we send the row iff the row is allocated and satisfies the predicate
builder.assert_eq(
local_cols.page_cols.is_alloc * local_cols.satisfies_pred,
page_cols.is_alloc * local_cols.satisfies_pred,
local_cols.send_row,
);
// constrain that satisfies_pred and send_row are boolean indicators
Expand Down Expand Up @@ -175,7 +173,7 @@ where
}) => Some(IsLessThanTupleCols {
io: IsLessThanTupleIOCols {
// idx < x
x: local_cols.page_cols.idx.clone(),
x: page_cols.idx.clone(),
y: local_cols.x.clone(),
// use the strict_comp_ind
tuple_less_than: strict_comp_ind.unwrap(),
Expand All @@ -193,7 +191,7 @@ where
io: IsLessThanTupleIOCols {
// idx > x
x: local_cols.x.clone(),
y: local_cols.page_cols.idx.clone(),
y: page_cols.idx.clone(),
// use the strict_comp_ind
tuple_less_than: strict_comp_ind.unwrap(),
},
Expand All @@ -215,7 +213,7 @@ where
}) => {
let is_equal_vec_cols = IsEqualVecCols {
io: IsEqualVecIOCols {
x: local_cols.page_cols.idx.clone(),
x: page_cols.idx.clone(),
y: local_cols.x.clone(),
// use the equal_comp_ind
prod: equal_comp_ind.unwrap(),
Expand Down
10 changes: 5 additions & 5 deletions chips/src/single_page_index_scan/page_index_scan_input/bridge.rs
Original file line number Diff line number Diff line change
Expand Up @@ -85,7 +85,7 @@ impl<F: PrimeField64> AirBridge<F> for PageIndexScanInputAir {

interactions.push(Interaction {
fields: virtual_cols,
count: VirtualPairCol::single_main(cols_numbered.send_row),
count: VirtualPairCol::single_main(cols_numbered.local_cols.send_row),
argument_index: self.page_bus_index,
});

Expand All @@ -94,7 +94,7 @@ impl<F: PrimeField64> AirBridge<F> for PageIndexScanInputAir {
let (is_less_than_tuple_aux_flattened, strict_comp_ind): (
Option<Vec<usize>>,
Option<usize>,
) = match cols_numbered.aux_cols {
) = match cols_numbered.local_cols.aux_cols {
PageIndexScanInputAuxCols::Lt(StrictCompAuxCols {
is_less_than_tuple_aux,
..
Expand All @@ -104,7 +104,7 @@ impl<F: PrimeField64> AirBridge<F> for PageIndexScanInputAir {
..
}) => (
Some(is_less_than_tuple_aux.flatten()),
Some(cols_numbered.satisfies_pred),
Some(cols_numbered.local_cols.satisfies_pred),
),
PageIndexScanInputAuxCols::Lte(NonStrictCompAuxCols {
satisfies_strict_comp,
Expand Down Expand Up @@ -135,7 +135,7 @@ impl<F: PrimeField64> AirBridge<F> for PageIndexScanInputAir {
let is_less_than_tuple_cols = IsLessThanTupleCols {
io: IsLessThanTupleIOCols {
x: cols_numbered.page_cols.idx.clone(),
y: cols_numbered.x.clone(),
y: cols_numbered.local_cols.x.clone(),
tuple_less_than: strict_comp_ind.unwrap(),
},
aux: IsLessThanTupleAuxCols::from_slice(
Expand All @@ -158,7 +158,7 @@ impl<F: PrimeField64> AirBridge<F> for PageIndexScanInputAir {
}) => {
let is_less_than_tuple_cols = IsLessThanTupleCols {
io: IsLessThanTupleIOCols {
x: cols_numbered.x.clone(),
x: cols_numbered.local_cols.x.clone(),
y: cols_numbered.page_cols.idx.clone(),
tuple_less_than: strict_comp_ind.unwrap(),
},
Expand Down
91 changes: 61 additions & 30 deletions chips/src/single_page_index_scan/page_index_scan_input/columns.rs
Original file line number Diff line number Diff line change
Expand Up @@ -28,37 +28,29 @@ pub enum PageIndexScanInputAuxCols<T> {
Gt(StrictCompAuxCols<T>),
}

pub struct PageIndexScanInputCols<T> {
pub page_cols: PageCols<T>,
pub struct PageIndexScanInputLocalCols<T> {
pub x: Vec<T>,
pub satisfies_pred: T,
pub send_row: T,
pub aux_cols: PageIndexScanInputAuxCols<T>,
}

impl<T: Clone> PageIndexScanInputCols<T> {
impl<T: Clone> PageIndexScanInputLocalCols<T> {
pub fn from_slice(
slc: &[T],
idx_len: usize,
data_len: usize,
idx_limb_bits: Vec<usize>,
decomp: usize,
cmp: Comp,
) -> Self {
let page_cols = PageCols {
is_alloc: slc[0].clone(),
idx: slc[1..idx_len + 1].to_vec(),
data: slc[idx_len + 1..idx_len + data_len + 1].to_vec(),
};

let x = slc[idx_len + data_len + 1..2 * idx_len + data_len + 1].to_vec();
let satisfies_pred = slc[2 * idx_len + data_len + 1].clone();
let send_row = slc[2 * idx_len + data_len + 2].clone();
let x = slc[0..idx_len].to_vec();
let satisfies_pred = slc[idx_len].clone();
let send_row = slc[idx_len + 1].clone();

let aux_cols = match cmp {
Comp::Lt => PageIndexScanInputAuxCols::Lt(StrictCompAuxCols {
is_less_than_tuple_aux: IsLessThanTupleAuxCols::from_slice(
&slc[2 * idx_len + data_len + 3..],
&slc[idx_len + 2..],
idx_limb_bits,
decomp,
idx_len,
Expand All @@ -68,49 +60,44 @@ impl<T: Clone> PageIndexScanInputCols<T> {
let less_than_tuple_aux_width =
IsLessThanTupleAuxCols::<T>::get_width(idx_limb_bits.clone(), decomp, idx_len);
PageIndexScanInputAuxCols::Lte(NonStrictCompAuxCols {
satisfies_strict_comp: slc[2 * idx_len + data_len + 3].clone(),
satisfies_eq_comp: slc[2 * idx_len + data_len + 4].clone(),
satisfies_strict_comp: slc[idx_len + 2].clone(),
satisfies_eq_comp: slc[idx_len + 3].clone(),
is_less_than_tuple_aux: IsLessThanTupleAuxCols::from_slice(
&slc[2 * idx_len + data_len + 5
..2 * idx_len + data_len + 5 + less_than_tuple_aux_width],
&slc[idx_len + 4..idx_len + 4 + less_than_tuple_aux_width],
idx_limb_bits,
decomp,
idx_len,
),
is_equal_vec_aux: IsEqualVecAuxCols::from_slice(
&slc[2 * idx_len + data_len + 5 + less_than_tuple_aux_width..],
&slc[idx_len + 4 + less_than_tuple_aux_width..],
idx_len,
),
})
}
Comp::Eq => PageIndexScanInputAuxCols::Eq(EqCompAuxCols {
is_equal_vec_aux: IsEqualVecAuxCols::from_slice(
&slc[2 * idx_len + data_len + 3..],
idx_len,
),
is_equal_vec_aux: IsEqualVecAuxCols::from_slice(&slc[idx_len + 2..], idx_len),
}),
Comp::Gte => {
let less_than_tuple_aux_width =
IsLessThanTupleAuxCols::<T>::get_width(idx_limb_bits.clone(), decomp, idx_len);
PageIndexScanInputAuxCols::Gte(NonStrictCompAuxCols {
satisfies_strict_comp: slc[2 * idx_len + data_len + 3].clone(),
satisfies_eq_comp: slc[2 * idx_len + data_len + 4].clone(),
satisfies_strict_comp: slc[idx_len + 2].clone(),
satisfies_eq_comp: slc[idx_len + 3].clone(),
is_less_than_tuple_aux: IsLessThanTupleAuxCols::from_slice(
&slc[2 * idx_len + data_len + 5
..2 * idx_len + data_len + 5 + less_than_tuple_aux_width],
&slc[idx_len + 4..idx_len + 4 + less_than_tuple_aux_width],
idx_limb_bits,
decomp,
idx_len,
),
is_equal_vec_aux: IsEqualVecAuxCols::from_slice(
&slc[2 * idx_len + data_len + 5 + less_than_tuple_aux_width..],
&slc[idx_len + 4 + less_than_tuple_aux_width..],
idx_len,
),
})
}
Comp::Gt => PageIndexScanInputAuxCols::Gt(StrictCompAuxCols {
is_less_than_tuple_aux: IsLessThanTupleAuxCols::from_slice(
&slc[2 * idx_len + data_len + 3..],
&slc[idx_len + 2..],
idx_limb_bits,
decomp,
idx_len,
Expand All @@ -119,13 +106,57 @@ impl<T: Clone> PageIndexScanInputCols<T> {
};

Self {
page_cols,
x,
satisfies_pred,
send_row,
aux_cols,
}
}
}

pub struct PageIndexScanInputCols<T> {
pub page_cols: PageCols<T>,
pub local_cols: PageIndexScanInputLocalCols<T>,
}

impl<T: Clone> PageIndexScanInputCols<T> {
pub fn from_partitioned_slice(
page_slc: &[T],
aux_slc: &[T],
idx_len: usize,
data_len: usize,
idx_limb_bits: Vec<usize>,
decomp: usize,
cmp: Comp,
) -> Self {
let page_cols = PageCols::from_slice(page_slc, idx_len, data_len);
let local_cols =
PageIndexScanInputLocalCols::from_slice(aux_slc, idx_len, idx_limb_bits, decomp, cmp);

Self {
page_cols,
local_cols,
}
}

pub fn from_slice(
slc: &[T],
idx_len: usize,
data_len: usize,
idx_limb_bits: Vec<usize>,
decomp: usize,
cmp: Comp,
) -> Self {
Self::from_partitioned_slice(
&slc[..idx_len + data_len + 1],
&slc[idx_len + data_len + 1..],
idx_len,
data_len,
idx_limb_bits,
decomp,
cmp,
)
}

pub fn get_width(
idx_len: usize,
Expand Down

0 comments on commit 41bb88b

Please sign in to comment.