openvm_circuit_primitives/range/
mod.rs

Help
1//! Range check for a fixed bit size via preprocessed trace.
2//!
3//! Caution: We almost always prefer to use the [VariableRangeCheckerChip](super::var_range::VariableRangeCheckerChip) instead of this chip.
4// Adapted from Valida
5
6use core::mem::size_of;
7use std::{
8    borrow::{Borrow, BorrowMut},
9    sync::atomic::AtomicU32,
10};
11
12use openvm_circuit_primitives_derive::AlignedBorrow;
13use openvm_stark_backend::{
14    interaction::InteractionBuilder,
15    p3_air::{Air, BaseAir, PairBuilder},
16    p3_field::Field,
17    p3_matrix::{dense::RowMajorMatrix, Matrix},
18    rap::{BaseAirWithPublicValues, PartitionedBaseAir},
19};
20
21mod bus;
22
23#[cfg(test)]
24pub mod tests;
25
26pub use bus::*;
27
28#[derive(Default, AlignedBorrow, Copy, Clone)]
29#[repr(C)]
30pub struct RangeCols<T> {
31    /// Number of range checks for each value
32    pub mult: T,
33}
34
35#[derive(Default, AlignedBorrow, Copy, Clone)]
36#[repr(C)]
37pub struct RangePreprocessedCols<T> {
38    /// Contains all possible values within range [0, max)
39    pub counter: T,
40}
41
42pub const NUM_RANGE_COLS: usize = size_of::<RangeCols<u8>>();
43pub const NUM_RANGE_PREPROCESSED_COLS: usize = size_of::<RangePreprocessedCols<u8>>();
44
45#[derive(Clone, Copy, Debug, derive_new::new)]
46pub struct RangeCheckerAir {
47    pub bus: RangeCheckBus,
48}
49
50impl RangeCheckerAir {
51    pub fn range_max(&self) -> u32 {
52        self.bus.range_max
53    }
54}
55
56impl<F: Field> BaseAirWithPublicValues<F> for RangeCheckerAir {}
57impl<F: Field> PartitionedBaseAir<F> for RangeCheckerAir {}
58impl<F: Field> BaseAir<F> for RangeCheckerAir {
59    fn width(&self) -> usize {
60        NUM_RANGE_COLS
61    }
62
63    fn preprocessed_trace(&self) -> Option<RowMajorMatrix<F>> {
64        // Create lookup table with all values 0..range_max
65        let column = (0..self.range_max()).map(F::from_canonical_u32).collect();
66        Some(RowMajorMatrix::new_col(column))
67    }
68}
69
70impl<AB: InteractionBuilder + PairBuilder> Air<AB> for RangeCheckerAir {
71    fn eval(&self, builder: &mut AB) {
72        let preprocessed = builder.preprocessed();
73        let prep_local = preprocessed.row_slice(0);
74        let prep_local: &RangePreprocessedCols<AB::Var> = (*prep_local).borrow();
75        let main = builder.main();
76        let local = main.row_slice(0);
77        let local: &RangeCols<AB::Var> = (*local).borrow();
78        // Omit creating separate bridge.rs file for brevity
79        self.bus
80            .receive(prep_local.counter)
81            .eval(builder, local.mult);
82    }
83}
84
85pub struct RangeCheckerChip {
86    pub air: RangeCheckerAir,
87    /// Tracks multiplicity of each value in range
88    count: Vec<AtomicU32>,
89}
90
91impl RangeCheckerChip {
92    pub fn new(bus: RangeCheckBus) -> Self {
93        let mut count = vec![];
94        for _ in 0..bus.range_max {
95            count.push(AtomicU32::new(0));
96        }
97
98        Self {
99            air: RangeCheckerAir::new(bus),
100            count,
101        }
102    }
103
104    pub fn bus(&self) -> RangeCheckBus {
105        self.air.bus
106    }
107
108    pub fn range_max(&self) -> u32 {
109        self.air.range_max()
110    }
111
112    pub fn add_count(&self, val: u32) {
113        // Increment the count for this value when range checked
114        let val_atomic = &self.count[val as usize];
115        val_atomic.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
116    }
117
118    pub fn generate_trace<F: Field>(&self) -> RowMajorMatrix<F> {
119        let mut rows = F::zero_vec(self.air.range_max() as usize * NUM_RANGE_COLS);
120        for (n, row) in rows.chunks_exact_mut(NUM_RANGE_COLS).enumerate() {
121            let cols: &mut RangeCols<F> = (*row).borrow_mut();
122            // Set multiplicity for each value in range
123            cols.mult =
124                F::from_canonical_u32(self.count[n].load(std::sync::atomic::Ordering::SeqCst));
125        }
126        RowMajorMatrix::new(rows, NUM_RANGE_COLS)
127    }
128}
openvm_circuit_primitives/range/mod.rs

openvm_circuit_primitives/range/
mod.rs
