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
24    /// the field modulus**.
25    pub fn range_check<T: FieldAlgebra>(
26        &self,
27        x: impl Into<T>,
28        max_bits: usize,
29    ) -> BitsCheckBusInteraction<T>
30    where
31        T::F: PrimeField32,
32    {
33        debug_assert!((1 << max_bits) <= self.range_max);
34        debug_assert!(self.range_max < T::F::ORDER_U32 / 2);
35        let shift = self.range_max - (1 << max_bits);
36        BitsCheckBusInteraction {
37            x: x.into(),
38            shift,
39            bus: self.inner,
40        }
41    }
42
43    pub fn send<T>(&self, x: impl Into<T>) -> RangeCheckBusInteraction<T> {
44        self.push(x, true)
45    }
46
47    pub fn receive<T>(&self, x: impl Into<T>) -> RangeCheckBusInteraction<T> {
48        self.push(x, false)
49    }
50
51    pub fn push<T>(&self, x: impl Into<T>, is_lookup: bool) -> RangeCheckBusInteraction<T> {
52        RangeCheckBusInteraction {
53            x: x.into(),
54            bus: self.inner,
55            is_lookup,
56        }
57    }
58
59    pub fn index(&self) -> BusIndex {
60        self.inner.index
61    }
62}
63
64#[derive(Clone, Copy, Debug)]
65pub struct BitsCheckBusInteraction<T> {
66    pub x: T,
67    pub shift: u32,
68    pub bus: LookupBus,
69}
70
71#[derive(Clone, Copy, Debug)]
72pub struct RangeCheckBusInteraction<T> {
73    pub x: T,
74
75    pub bus: LookupBus,
76    pub is_lookup: bool,
77}
78
79impl<T: FieldAlgebra> RangeCheckBusInteraction<T> {
80    /// Finalizes and sends/receives over the RangeCheck bus.
81    pub fn eval<AB>(self, builder: &mut AB, count: impl Into<AB::Expr>)
82    where
83        AB: InteractionBuilder<Expr = T>,
84    {
85        if self.is_lookup {
86            self.bus.lookup_key(builder, [self.x], count);
87        } else {
88            self.bus.add_key_with_lookups(builder, [self.x], count);
89        }
90    }
91}
92
93impl<T: FieldAlgebra> BitsCheckBusInteraction<T> {
94    /// Send interaction(s) to range check for max bits over the RangeCheck bus.
95    pub fn eval<AB>(self, builder: &mut AB, count: impl Into<AB::Expr>)
96    where
97        AB: InteractionBuilder<Expr = T>,
98    {
99        let count = count.into();
100        if self.shift > 0 {
101            // if 2^max_bits < range_max, then we also range check that `x + (range_max -
102            // 2^max_bits) < range_max`
103            // - this will hold if `x < 2^max_bits` (necessary)
104            // - if `x < range_max` then we know the integer value `x.as_canonical_u32() +
105            //   (range_max - 2^max_bits) < 2*range_max`. **Assuming that `2*range_max <
106            //   F::MODULUS`, then additionally knowing `x + (range_max - 2^max_bits) < range_max`
107            //   implies `x < 2^max_bits`.
108            self.bus.lookup_key(
109                builder,
110                [self.x.clone() + AB::Expr::from_canonical_u32(self.shift)],
111                count.clone(),
112            );
113        }
114        self.bus.lookup_key(builder, [self.x], count);
115    }
116}