openvm_circuit_primitives/range_tuple/
bus.rs

1use openvm_stark_backend::{
2    interaction::{BusIndex, InteractionBuilder, LookupBus},
3    p3_field::FieldAlgebra,
4};
5
6#[derive(Clone, Copy, Debug, PartialEq, Eq)]
7pub struct RangeTupleCheckerBus<const N: usize> {
8    pub inner: LookupBus,
9    pub sizes: [u32; N],
10}
11
12impl<const N: usize> RangeTupleCheckerBus<N> {
13    pub fn new(index: BusIndex, sizes: [u32; N]) -> Self {
14        let mut product = 1u32;
15        for &size in sizes.iter() {
16            product = product
17                .checked_mul(size)
18                .expect("The number of the range tuple checker rows is too large");
19        }
20        Self {
21            inner: LookupBus::new(index),
22            sizes,
23        }
24    }
25
26    #[must_use]
27    pub fn send<T>(&self, tuple: Vec<impl Into<T>>) -> RangeTupleCheckerBusInteraction<T> {
28        self.push(tuple, true)
29    }
30
31    #[must_use]
32    pub fn receive<T>(&self, tuple: Vec<impl Into<T>>) -> RangeTupleCheckerBusInteraction<T> {
33        self.push(tuple, false)
34    }
35
36    pub fn push<T>(
37        &self,
38        tuple: Vec<impl Into<T>>,
39        is_lookup: bool,
40    ) -> RangeTupleCheckerBusInteraction<T> {
41        RangeTupleCheckerBusInteraction {
42            tuple: tuple.into_iter().map(|t| t.into()).collect(),
43            bus: self.inner,
44            is_lookup,
45        }
46    }
47}
48
49#[derive(Clone, Debug)]
50pub struct RangeTupleCheckerBusInteraction<T> {
51    pub tuple: Vec<T>,
52    pub bus: LookupBus,
53    pub is_lookup: bool,
54}
55
56impl<T: FieldAlgebra> RangeTupleCheckerBusInteraction<T> {
57    pub fn eval<AB>(self, builder: &mut AB, count: impl Into<AB::Expr>)
58    where
59        AB: InteractionBuilder<Expr = T>,
60    {
61        if self.is_lookup {
62            self.bus.lookup_key(builder, self.tuple, count);
63        } else {
64            self.bus.add_key_with_lookups(builder, self.tuple, count);
65        }
66    }
67}