openvm_circuit_primitives/range/
mod.rs
1use core::mem::size_of;
8use std::{
9 borrow::{Borrow, BorrowMut},
10 sync::atomic::AtomicU32,
11};
12
13use openvm_circuit_primitives_derive::AlignedBorrow;
14use openvm_stark_backend::{
15 interaction::InteractionBuilder,
16 p3_air::{Air, BaseAir, PairBuilder},
17 p3_field::Field,
18 p3_matrix::{dense::RowMajorMatrix, Matrix},
19 rap::{BaseAirWithPublicValues, PartitionedBaseAir},
20};
21
22mod bus;
23
24#[cfg(test)]
25pub mod tests;
26
27pub use bus::*;
28
29#[derive(Default, AlignedBorrow, Copy, Clone)]
30#[repr(C)]
31pub struct RangeCols<T> {
32 pub mult: T,
34}
35
36#[derive(Default, AlignedBorrow, Copy, Clone)]
37#[repr(C)]
38pub struct RangePreprocessedCols<T> {
39 pub counter: T,
41}
42
43pub const NUM_RANGE_COLS: usize = size_of::<RangeCols<u8>>();
44pub const NUM_RANGE_PREPROCESSED_COLS: usize = size_of::<RangePreprocessedCols<u8>>();
45
46#[derive(Clone, Copy, Debug, derive_new::new)]
47pub struct RangeCheckerAir {
48 pub bus: RangeCheckBus,
49}
50
51impl RangeCheckerAir {
52 pub fn range_max(&self) -> u32 {
53 self.bus.range_max
54 }
55}
56
57impl<F: Field> BaseAirWithPublicValues<F> for RangeCheckerAir {}
58impl<F: Field> PartitionedBaseAir<F> for RangeCheckerAir {}
59impl<F: Field> BaseAir<F> for RangeCheckerAir {
60 fn width(&self) -> usize {
61 NUM_RANGE_COLS
62 }
63
64 fn preprocessed_trace(&self) -> Option<RowMajorMatrix<F>> {
65 let column = (0..self.range_max()).map(F::from_canonical_u32).collect();
67 Some(RowMajorMatrix::new_col(column))
68 }
69}
70
71impl<AB: InteractionBuilder + PairBuilder> Air<AB> for RangeCheckerAir {
72 fn eval(&self, builder: &mut AB) {
73 let preprocessed = builder.preprocessed();
74 let prep_local = preprocessed.row_slice(0);
75 let prep_local: &RangePreprocessedCols<AB::Var> = (*prep_local).borrow();
76 let main = builder.main();
77 let local = main.row_slice(0);
78 let local: &RangeCols<AB::Var> = (*local).borrow();
79 self.bus
81 .receive(prep_local.counter)
82 .eval(builder, local.mult);
83 }
84}
85
86pub struct RangeCheckerChip {
87 pub air: RangeCheckerAir,
88 count: Vec<AtomicU32>,
90}
91
92impl RangeCheckerChip {
93 pub fn new(bus: RangeCheckBus) -> Self {
94 let mut count = vec![];
95 for _ in 0..bus.range_max {
96 count.push(AtomicU32::new(0));
97 }
98
99 Self {
100 air: RangeCheckerAir::new(bus),
101 count,
102 }
103 }
104
105 pub fn bus(&self) -> RangeCheckBus {
106 self.air.bus
107 }
108
109 pub fn range_max(&self) -> u32 {
110 self.air.range_max()
111 }
112
113 pub fn add_count(&self, val: u32) {
114 let val_atomic = &self.count[val as usize];
116 val_atomic.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
117 }
118
119 pub fn generate_trace<F: Field>(&self) -> RowMajorMatrix<F> {
120 let mut rows = F::zero_vec(self.air.range_max() as usize * NUM_RANGE_COLS);
121 for (n, row) in rows.chunks_exact_mut(NUM_RANGE_COLS).enumerate() {
122 let cols: &mut RangeCols<F> = (*row).borrow_mut();
123 cols.mult =
125 F::from_canonical_u32(self.count[n].load(std::sync::atomic::Ordering::SeqCst));
126 }
127 RowMajorMatrix::new(rows, NUM_RANGE_COLS)
128 }
129}