openvm_circuit_primitives/range/
bus.rs

1use openvm_stark_backend::{
2    interaction::{BusIndex, InteractionBuilder, LookupBus},
3    p3_field::{FieldAlgebra, PrimeField32},
4};
5
6/// Represents a bus for `x` where `x` must lie in the range `[0, range_max)`.
7#[derive(Clone, Copy, Debug, PartialEq, Eq)]
8pub struct RangeCheckBus {
9    pub inner: LookupBus,
10    pub range_max: u32,
11}
12
13impl RangeCheckBus {
14    pub const fn new(index: BusIndex, range_max: u32) -> Self {
15        Self {
16            inner: LookupBus::new(index),
17            range_max,
18        }
19    }
20
21    /// Range check that `x` is in the range `[0, 2^max_bits)`.
22    ///
23    /// This can be used when `2^max_bits < self.range_max` **if `2 * self.range_max` is less than the field modulus**.
24    pub fn range_check<T: FieldAlgebra>(
25        &self,
26        x: impl Into<T>,
27        max_bits: usize,
28    ) -> BitsCheckBusInteraction<T>
29    where
30        T::F: PrimeField32,
31    {
32        debug_assert!((1 << max_bits) <= self.range_max);
33        debug_assert!(self.range_max < T::F::ORDER_U32 / 2);
34        let shift = self.range_max - (1 << max_bits);
35        BitsCheckBusInteraction {
36            x: x.into(),
37            shift,
38            bus: self.inner,
39        }
40    }
41
42    pub fn send<T>(&self, x: impl Into<T>) -> RangeCheckBusInteraction<T> {
43        self.push(x, true)
44    }
45
46    pub fn receive<T>(&self, x: impl Into<T>) -> RangeCheckBusInteraction<T> {
47        self.push(x, false)
48    }
49
50    pub fn push<T>(&self, x: impl Into<T>, is_lookup: bool) -> RangeCheckBusInteraction<T> {
51        RangeCheckBusInteraction {
52            x: x.into(),
53            bus: self.inner,
54            is_lookup,
55        }
56    }
57
58    pub fn index(&self) -> BusIndex {
59        self.inner.index
60    }
61}
62
63#[derive(Clone, Copy, Debug)]
64pub struct BitsCheckBusInteraction<T> {
65    pub x: T,
66    pub shift: u32,
67    pub bus: LookupBus,
68}
69
70#[derive(Clone, Copy, Debug)]
71pub struct RangeCheckBusInteraction<T> {
72    pub x: T,
73
74    pub bus: LookupBus,
75    pub is_lookup: bool,
76}
77
78impl<T: FieldAlgebra> RangeCheckBusInteraction<T> {
79    /// Finalizes and sends/receives over the RangeCheck bus.
80    pub fn eval<AB>(self, builder: &mut AB, count: impl Into<AB::Expr>)
81    where
82        AB: InteractionBuilder<Expr = T>,
83    {
84        if self.is_lookup {
85            self.bus.lookup_key(builder, [self.x], count);
86        } else {
87            self.bus.add_key_with_lookups(builder, [self.x], count);
88        }
89    }
90}
91
92impl<T: FieldAlgebra> BitsCheckBusInteraction<T> {
93    /// Send interaction(s) to range check for max bits over the RangeCheck bus.
94    pub fn eval<AB>(self, builder: &mut AB, count: impl Into<AB::Expr>)
95    where
96        AB: InteractionBuilder<Expr = T>,
97    {
98        let count = count.into();
99        if self.shift > 0 {
100            // if 2^max_bits < range_max, then we also range check that `x + (range_max - 2^max_bits) < range_max`
101            // - this will hold if `x < 2^max_bits` (necessary)
102            // - if `x < range_max` then we know the integer value `x.as_canonical_u32() + (range_max - 2^max_bits) < 2*range_max`.
103            //   **Assuming that `2*range_max < F::MODULUS`, then additionally knowing `x + (range_max - 2^max_bits) < range_max` implies `x < 2^max_bits`.
104            self.bus.lookup_key(
105                builder,
106                [self.x.clone() + AB::Expr::from_canonical_u32(self.shift)],
107                count.clone(),
108            );
109        }
110        self.bus.lookup_key(builder, [self.x], count);
111    }
112}