openvm_circuit_primitives/var_range/
mod.rs

1//! A chip which provides a lookup table for range checking a variable `x` has `b` bits
2//! where `b` can be any integer in `[0, range_max_bits]` without using preprocessed trace.
3//! In other words, the same chip can be used to range check for different bit sizes.
4//! We define `0` to have `0` bits.
5
6use core::mem::size_of;
7use std::{
8    borrow::{Borrow, BorrowMut},
9    sync::{atomic::AtomicU32, Arc},
10};
11
12use openvm_circuit_primitives_derive::AlignedBorrow;
13use openvm_cpu_backend::CpuBackend;
14use openvm_stark_backend::{
15    interaction::InteractionBuilder,
16    p3_air::{Air, AirBuilder, BaseAir},
17    p3_field::{Field, PrimeCharacteristicRing, PrimeField32},
18    p3_matrix::{dense::RowMajorMatrix, Matrix},
19    prover::AirProvingContext,
20    BaseAirWithPublicValues, PartitionedBaseAir, StarkProtocolConfig, Val,
21};
22use tracing::instrument;
23
24use crate::{Chip, ColumnsAir, StructReflection, StructReflectionHelper};
25
26mod bus;
27pub use bus::*;
28
29#[cfg(feature = "cuda")]
30mod cuda;
31#[cfg(feature = "cuda")]
32pub use cuda::*;
33
34#[cfg(test)]
35pub mod tests;
36
37#[derive(Default, AlignedBorrow, StructReflection, Copy, Clone)]
38#[repr(C)]
39pub struct VariableRangeCols<T> {
40    /// The value being range checked
41    pub value: T,
42    /// The maximum number of bits for this value
43    pub max_bits: T,
44    /// Helper column storing 2^max_bits, used in constraints
45    pub two_to_max_bits: T,
46    /// Number of range checks requested for each (value, max_bits) pair
47    pub mult: T,
48}
49
50pub const NUM_VARIABLE_RANGE_COLS: usize = size_of::<VariableRangeCols<u8>>();
51
52#[derive(Clone, Copy, Debug, derive_new::new, ColumnsAir)]
53#[columns_via(VariableRangeCols<u8>)]
54pub struct VariableRangeCheckerAir {
55    pub bus: VariableRangeCheckerBus,
56}
57
58impl VariableRangeCheckerAir {
59    pub fn range_max_bits(&self) -> usize {
60        self.bus.range_max_bits
61    }
62}
63
64impl<F: Field> BaseAirWithPublicValues<F> for VariableRangeCheckerAir {}
65impl<F: Field> PartitionedBaseAir<F> for VariableRangeCheckerAir {}
66impl<F: Field> BaseAir<F> for VariableRangeCheckerAir {
67    fn width(&self) -> usize {
68        VariableRangeCols::<F>::width()
69    }
70}
71
72impl<AB: InteractionBuilder> Air<AB> for VariableRangeCheckerAir {
73    fn eval(&self, builder: &mut AB) {
74        let main = builder.main();
75
76        let (local, next) = (
77            main.row_slice(0).expect("window should have two elements"),
78            main.row_slice(1).expect("window should have two elements"),
79        );
80        let local: &VariableRangeCols<AB::Var> = (*local).borrow();
81        let next: &VariableRangeCols<AB::Var> = (*next).borrow();
82
83        // First row: start at [value=0, max_bits=0, two_to_max_bits=1]
84        builder.when_first_row().assert_zero(local.value);
85        builder.when_first_row().assert_zero(local.max_bits);
86        builder.when_first_row().assert_one(local.two_to_max_bits);
87
88        // Transition constraints use a "monotonic sum" approach instead of selector-based
89        // branching. The key insight is that (value + two_to_max_bits) equals
90        // (row_index + 1), forming a strictly increasing sequence. Combined with the last-row
91        // constraint, this forces the unique valid trace enumeration.
92        let max_bits_delta = next.max_bits - local.max_bits;
93
94        // max_bits can only stay the same or increment by 1
95        builder
96            .when_transition()
97            .assert_bool(max_bits_delta.clone());
98
99        // value can only increment by 1 or wrap back to 0
100        builder
101            .when_transition()
102            .when(next.value)
103            .assert_eq(next.value, local.value + AB::Expr::ONE);
104
105        // two_to_max_bits doubles whenever max_bits increments
106        builder.when_transition().assert_eq(
107            next.two_to_max_bits,
108            local.two_to_max_bits * (AB::Expr::ONE + max_bits_delta),
109        );
110
111        // (value + two_to_max_bits) increases by exactly 1 each row
112        builder.when_transition().assert_eq(
113            local.value + local.two_to_max_bits + AB::Expr::ONE,
114            next.value + next.two_to_max_bits,
115        );
116
117        // Last row: end at [value=0, max_bits=range_max_bits+1, mult=0]
118        // If value ever skips wrapping to 0, the monotonic sum constraint forces it to keep
119        // incrementing, making this final state unreachable.
120        builder.when_last_row().assert_zero(local.value);
121        builder.when_last_row().assert_eq(
122            local.max_bits,
123            AB::F::from_usize(self.bus.range_max_bits + 1),
124        );
125        builder.when_last_row().assert_zero(local.mult);
126
127        self.bus
128            .receive(local.value, local.max_bits)
129            .eval(builder, local.mult);
130    }
131}
132
133pub struct VariableRangeCheckerChip {
134    pub air: VariableRangeCheckerAir,
135    pub count: Vec<AtomicU32>,
136}
137
138pub type SharedVariableRangeCheckerChip = Arc<VariableRangeCheckerChip>;
139
140impl VariableRangeCheckerChip {
141    pub fn new(bus: VariableRangeCheckerBus) -> Self {
142        let num_rows = (1 << (bus.range_max_bits + 1)) as usize;
143        let count = (0..num_rows).map(|_| AtomicU32::new(0)).collect();
144        Self {
145            air: VariableRangeCheckerAir::new(bus),
146            count,
147        }
148    }
149
150    pub fn bus(&self) -> VariableRangeCheckerBus {
151        self.air.bus
152    }
153
154    pub fn range_max_bits(&self) -> usize {
155        self.air.range_max_bits()
156    }
157
158    pub fn air_width(&self) -> usize {
159        NUM_VARIABLE_RANGE_COLS
160    }
161
162    #[instrument(
163        name = "VariableRangeCheckerChip::add_count",
164        skip(self),
165        level = "trace"
166    )]
167    pub fn add_count(&self, value: u32, max_bits: usize) {
168        // index is 2^max_bits + value - 1
169        // if each [value, max_bits] is valid, the sends multiset will be exactly the receives
170        // multiset
171        let idx = (1 << max_bits) + (value as usize) - 1;
172        assert!(
173            idx < self.count.len(),
174            "range exceeded: {} >= {}",
175            idx,
176            self.count.len()
177        );
178        let val_atomic = &self.count[idx];
179        val_atomic.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
180    }
181
182    pub fn clear(&self) {
183        for i in 0..self.count.len() {
184            self.count[i].store(0, std::sync::atomic::Ordering::Relaxed);
185        }
186    }
187
188    /// Generates trace and resets the internal counters all to 0.
189    pub fn generate_trace<F: Field>(&self) -> RowMajorMatrix<F> {
190        let mut rows = F::zero_vec(self.count.len() * NUM_VARIABLE_RANGE_COLS);
191        for (i, row) in rows.chunks_exact_mut(NUM_VARIABLE_RANGE_COLS).enumerate() {
192            let cols: &mut VariableRangeCols<F> = (*row).borrow_mut();
193            let max_bits = (i + 1).ilog2();
194            let two_to_max_bits = 1 << max_bits;
195            let value = i + 1 - two_to_max_bits;
196
197            cols.value = F::from_usize(value);
198            cols.max_bits = F::from_u32(max_bits);
199            cols.two_to_max_bits = F::from_usize(two_to_max_bits);
200            cols.mult = F::from_u32(self.count[i].swap(0, std::sync::atomic::Ordering::Relaxed));
201        }
202        RowMajorMatrix::new(rows, NUM_VARIABLE_RANGE_COLS)
203    }
204
205    /// Range checks that `value` is `bits` bits by decomposing into `limbs` where all but
206    /// last limb is `range_max_bits` bits. Assumes there are enough limbs.
207    pub fn decompose<F: Field>(&self, mut value: u32, bits: usize, limbs: &mut [F]) {
208        debug_assert!(
209            limbs.len() >= bits.div_ceil(self.range_max_bits()),
210            "Not enough limbs: len {}",
211            limbs.len()
212        );
213        let mask = (1 << self.range_max_bits()) - 1;
214        let mut bits_remaining = bits;
215        for limb in limbs.iter_mut() {
216            let limb_u32 = value & mask;
217            *limb = F::from_u32(limb_u32);
218            self.add_count(limb_u32, bits_remaining.min(self.range_max_bits()));
219
220            value >>= self.range_max_bits();
221            bits_remaining = bits_remaining.saturating_sub(self.range_max_bits());
222        }
223        debug_assert_eq!(value, 0);
224        debug_assert_eq!(bits_remaining, 0);
225    }
226}
227
228// We allow any `R` type so this can work with arbitrary record arenas.
229impl<R, SC: StarkProtocolConfig> Chip<R, CpuBackend<SC>> for VariableRangeCheckerChip
230where
231    Val<SC>: PrimeField32,
232{
233    /// Generates trace and resets the internal counters all to 0.
234    fn generate_proving_ctx(&self, _: R) -> AirProvingContext<CpuBackend<SC>> {
235        let trace_row_maj = self.generate_trace::<Val<SC>>();
236        AirProvingContext::simple_no_pis(trace_row_maj)
237    }
238
239    fn constant_trace_height(&self) -> Option<usize> {
240        Some(1 << (self.air.range_max_bits() + 1))
241    }
242}
openvm_circuit_primitives/var_range/mod.rs

openvm_circuit_primitives/var_range/
mod.rs
