openvm_circuit_primitives/var_range/
bus.rs

1use openvm_stark_backend::{
2    interaction::{BusIndex, InteractionBuilder, LookupBus},
3    p3_field::FieldAlgebra,
4};
5
6// Represents a bus for (x, bits) where either (x, bits) = (0, 0) or
7// x is in [0, 2^bits) and bits is in [1, range_max_bits]
8#[derive(Clone, Copy, Debug, PartialEq, Eq)]
9pub struct VariableRangeCheckerBus {
10    pub inner: LookupBus,
11    pub range_max_bits: usize,
12}
13
14impl VariableRangeCheckerBus {
15    pub const fn new(index: BusIndex, range_max_bits: usize) -> Self {
16        Self {
17            inner: LookupBus::new(index),
18            range_max_bits,
19        }
20    }
21
22    #[inline(always)]
23    pub fn index(&self) -> BusIndex {
24        self.inner.index
25    }
26
27    #[must_use]
28    pub fn send<T>(
29        &self,
30        value: impl Into<T>,
31        max_bits: impl Into<T>,
32    ) -> VariableRangeCheckerBusInteraction<T> {
33        self.push(value, max_bits, true)
34    }
35
36    #[must_use]
37    pub fn receive<T>(
38        &self,
39        value: impl Into<T>,
40        max_bits: impl Into<T>,
41    ) -> VariableRangeCheckerBusInteraction<T> {
42        self.push(value, max_bits, false)
43    }
44
45    // Equivalent to `self.send(value, max_bits)` where max_bits is a usize constant
46    #[must_use]
47    pub fn range_check<T>(
48        &self,
49        value: impl Into<T>,
50        max_bits: usize,
51    ) -> VariableRangeCheckerBusInteraction<T>
52    where
53        T: FieldAlgebra,
54    {
55        debug_assert!(max_bits <= self.range_max_bits);
56        self.push(value, T::from_canonical_usize(max_bits), true)
57    }
58
59    pub fn push<T>(
60        &self,
61        value: impl Into<T>,
62        max_bits: impl Into<T>,
63        is_lookup: bool,
64    ) -> VariableRangeCheckerBusInteraction<T> {
65        VariableRangeCheckerBusInteraction {
66            value: value.into(),
67            max_bits: max_bits.into(),
68            bus: self.inner,
69            is_lookup,
70        }
71    }
72}
73
74#[derive(Clone, Copy, Debug)]
75pub struct VariableRangeCheckerBusInteraction<T> {
76    pub value: T,
77    pub max_bits: T,
78    pub bus: LookupBus,
79    pub is_lookup: bool,
80}
81
82impl<T: FieldAlgebra> VariableRangeCheckerBusInteraction<T> {
83    pub fn eval<AB>(self, builder: &mut AB, count: impl Into<AB::Expr>)
84    where
85        AB: InteractionBuilder<Expr = T>,
86    {
87        let key = [self.value, self.max_bits];
88        if self.is_lookup {
89            self.bus.lookup_key(builder, key, count);
90        } else {
91            self.bus.add_key_with_lookups(builder, key, count);
92        }
93    }
94}