openvm_circuit_primitives/range_tuple/
mod.rs

1//! Range check a tuple simultaneously.
2//! When you know you want to range check `(x, y)` to `x_bits, y_bits` respectively
3//! and `2^{x_bits + y_bits} < ~2^20`, then you can use this chip to do the range check in one
4//! interaction versus the two interactions necessary if you were to use
5//! [VariableRangeCheckerChip](super::var_range::VariableRangeCheckerChip) instead.
6
7use std::{
8    borrow::Borrow,
9    sync::{atomic::AtomicU32, Arc},
10};
11
12use openvm_cpu_backend::CpuBackend;
13use openvm_stark_backend::{
14    interaction::InteractionBuilder,
15    p3_air::{Air, AirBuilder, BaseAir, PairBuilder},
16    p3_field::{Field, PrimeCharacteristicRing, PrimeField32},
17    p3_matrix::{dense::RowMajorMatrix, Matrix},
18    prover::AirProvingContext,
19    BaseAirWithPublicValues, PartitionedBaseAir, StarkProtocolConfig, Val,
20};
21
22use crate::{Chip, ColumnsAir};
23
24mod bus;
25pub use bus::*;
26
27#[cfg(feature = "cuda")]
28mod cuda;
29#[cfg(feature = "cuda")]
30pub use cuda::*;
31
32#[cfg(test)]
33pub mod tests;
34
35#[derive(Copy, Clone)]
36pub struct RangeTupleColsRef<'a, T> {
37    /// Contains all possible tuple combinations within specified ranges. Has size N.
38    pub tuple: &'a [T],
39    /// Number of range checks requested for each tuple combination
40    pub mult: &'a T,
41}
42
43impl<'a, T> RangeTupleColsRef<'a, T> {
44    fn from_slice<const N: usize>(slice: &'a [T]) -> Self {
45        let (tuple, rest) = slice.split_at(N);
46        Self {
47            tuple,
48            mult: &rest[0],
49        }
50    }
51}
52
53pub struct RangeTupleColsRefMut<'a, T> {
54    /// Contains all possible tuple combinations within specified ranges. Has size N.
55    pub tuple: &'a mut [T],
56    /// Number of range checks requested for each tuple combination
57    pub mult: &'a mut T,
58}
59
60impl<'a, T> RangeTupleColsRefMut<'a, T> {
61    fn from_slice_mut<const N: usize>(slice: &'a mut [T]) -> Self {
62        let (tuple, rest) = slice.split_at_mut(N);
63        Self {
64            tuple,
65            mult: &mut rest[0],
66        }
67    }
68}
69
70#[derive(Clone, Copy, Debug)]
71pub struct RangeTupleCheckerAir<const N: usize> {
72    pub bus: RangeTupleCheckerBus<N>,
73}
74
75impl<const N: usize> RangeTupleCheckerAir<N> {
76    pub fn height(&self) -> u32 {
77        self.bus.sizes.iter().product()
78    }
79}
80impl<F: Field, const N: usize> BaseAirWithPublicValues<F> for RangeTupleCheckerAir<N> {}
81impl<F: Field, const N: usize> PartitionedBaseAir<F> for RangeTupleCheckerAir<N> {}
82impl<const N: usize> ColumnsAir for RangeTupleCheckerAir<N> {
83    // Implemented manually because struct reflection does not work on `RangeTupleColsRef`.
84    fn columns(&self) -> Option<Vec<String>> {
85        Some(
86            (0..N)
87                .map(|i| format!("tuple[{}]", i))
88                .chain(std::iter::once("mult".to_string()))
89                .collect(),
90        )
91    }
92}
93
94impl<F: Field, const N: usize> BaseAir<F> for RangeTupleCheckerAir<N> {
95    fn width(&self) -> usize {
96        N + 1
97    }
98}
99
100// An explanation to the constraints is available in the README.
101impl<AB: InteractionBuilder + PairBuilder, const N: usize> Air<AB> for RangeTupleCheckerAir<N> {
102    fn eval(&self, builder: &mut AB) {
103        let main = builder.main();
104        let (local, next) = (
105            main.row_slice(0).expect("window should have two elements"),
106            main.row_slice(1).expect("window should have two elements"),
107        );
108        let local = RangeTupleColsRef::from_slice::<N>((*local).borrow());
109        let next = RangeTupleColsRef::from_slice::<N>((*next).borrow());
110
111        // (T1): The trace starts with `(0, ..., 0)`.
112        // (T2): The trace ends with `(size[0]-1, ..., size[N-1]-1)`.
113        for i in 0..N {
114            builder.when_first_row().assert_zero(local.tuple[i]);
115            builder
116                .when_last_row()
117                .assert_eq(local.tuple[i], AB::F::from_u32(self.bus.sizes[i] - 1));
118        }
119
120        // (T4): Between consecutive tuples, column `0` can stay the same or increment.
121        builder
122            .when_transition()
123            .assert_bool(next.tuple[0] - local.tuple[0]);
124        // (T5): Between consecutive tuples, all other columns can stay the same, increment, or
125        // wrap.
126        for i in 1..N - 1 {
127            builder
128                .when_ne(next.tuple[i] - local.tuple[i], AB::Expr::ZERO)
129                .when_ne(next.tuple[i] - local.tuple[i], AB::Expr::ONE)
130                .assert_eq(local.tuple[i], AB::F::from_u32(self.bus.sizes[i] - 1));
131            builder
132                .when_ne(next.tuple[i] - local.tuple[i], AB::Expr::ZERO)
133                .when_ne(next.tuple[i] - local.tuple[i], AB::Expr::ONE)
134                .assert_eq(next.tuple[i], AB::Expr::ZERO);
135        }
136        // (T3): Between consecutive tuples, column `N-1` can increment or wrap.
137        builder
138            .when_ne(next.tuple[N - 1] - local.tuple[N - 1], AB::Expr::ONE)
139            .assert_eq(
140                local.tuple[N - 1],
141                AB::F::from_u32(self.bus.sizes[N - 1] - 1),
142            );
143        builder
144            .when_ne(next.tuple[N - 1] - local.tuple[N - 1], AB::Expr::ONE)
145            .assert_eq(next.tuple[N - 1], AB::Expr::ZERO);
146
147        // (T6): Between consecutive tuples, column `i` increments or wraps if and only if column
148        // `i+1` wraps.
149        for i in 0..N - 1 {
150            let x = next.tuple[i] - local.tuple[i];
151            let y = next.tuple[i + 1] - local.tuple[i + 1];
152            let a = -AB::F::from_u32(self.bus.sizes[i] - 1);
153            let b = -AB::F::from_u32(self.bus.sizes[i + 1] - 1);
154            // See range_tuple/README.md
155            builder.assert_zero(
156                y.clone() * (y.clone() - AB::Expr::ONE) * (-x.clone() * (a + AB::F::ONE) + a)
157                    + x.clone() * x.clone() * (y.clone() * (b + b - AB::F::ONE) - b * b),
158            );
159        }
160
161        self.bus
162            .receive(local.tuple.to_vec())
163            .eval(builder, *local.mult);
164    }
165}
166
167#[derive(Debug)]
168pub struct RangeTupleCheckerChip<const N: usize> {
169    pub air: RangeTupleCheckerAir<N>,
170    pub count: Vec<Arc<AtomicU32>>,
171}
172
173pub type SharedRangeTupleCheckerChip<const N: usize> = Arc<RangeTupleCheckerChip<N>>;
174
175impl<const N: usize> RangeTupleCheckerChip<N> {
176    pub fn new(bus: RangeTupleCheckerBus<N>) -> Self {
177        assert!(N > 1, "RangeTupleChecker requires at least 2 dimensions");
178        // size = 1 is not useful, and breaks the completeness guarantee of this chip
179        assert!(
180            bus.sizes.iter().all(|&s| s > 1),
181            "RangeTupleChecker requires all sizes to be > 1 (size=1 dimensions break \
182             the carry coupling constraint)"
183        );
184        let range_max = bus.sizes.iter().product();
185        assert!(
186            range_max > 0 && (range_max & (range_max - 1)) == 0,
187            "RangeTupleChecker requires range_max ({}) to be a power of 2",
188            range_max
189        );
190        let count = (0..range_max)
191            .map(|_| Arc::new(AtomicU32::new(0)))
192            .collect();
193
194        Self {
195            air: RangeTupleCheckerAir { bus },
196            count,
197        }
198    }
199
200    pub fn bus(&self) -> &RangeTupleCheckerBus<N> {
201        &self.air.bus
202    }
203
204    pub fn sizes(&self) -> &[u32; N] {
205        &self.air.bus.sizes
206    }
207
208    pub fn add_count(&self, ids: &[u32]) {
209        let index = ids
210            .iter()
211            .zip(self.air.bus.sizes.iter())
212            .fold(0, |acc, (id, sz)| acc * sz + id) as usize;
213        assert!(
214            index < self.count.len(),
215            "range exceeded: {} >= {}",
216            index,
217            self.count.len()
218        );
219        let val_atomic = &self.count[index];
220        val_atomic.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
221    }
222
223    pub fn clear(&self) {
224        for val in &self.count {
225            val.store(0, std::sync::atomic::Ordering::Relaxed);
226        }
227    }
228
229    pub fn generate_trace<F: Field + PrimeField32>(&self) -> RowMajorMatrix<F> {
230        let mut rows = F::zero_vec(self.count.len() * (N + 1));
231
232        for (i, row) in rows.chunks_exact_mut(N + 1).enumerate() {
233            let cols = RangeTupleColsRefMut::from_slice_mut::<N>(row);
234            let mut tmp_idx = i as u32;
235            for j in (0..N).rev() {
236                cols.tuple[j] = F::from_u32(tmp_idx % self.air.bus.sizes[j]);
237                tmp_idx /= self.air.bus.sizes[j];
238            }
239            *cols.mult = F::from_u32(self.count[i].swap(0, std::sync::atomic::Ordering::Relaxed));
240        }
241
242        RowMajorMatrix::new(rows, N + 1)
243    }
244}
245
246impl<R, SC: StarkProtocolConfig, const N: usize> Chip<R, CpuBackend<SC>>
247    for RangeTupleCheckerChip<N>
248where
249    Val<SC>: PrimeField32,
250{
251    fn generate_proving_ctx(&self, _: R) -> AirProvingContext<CpuBackend<SC>> {
252        let trace_row_maj = self.generate_trace::<Val<SC>>();
253        AirProvingContext::simple_no_pis(trace_row_maj)
254    }
255
256    fn constant_trace_height(&self) -> Option<usize> {
257        Some(self.air.height() as usize)
258    }
259}
openvm_circuit_primitives/range_tuple/mod.rs

openvm_circuit_primitives/range_tuple/
mod.rs
