openvm_deferral_circuit/count/
bus.rs

1use openvm_stark_backend::{
2    interaction::{BusIndex, InteractionBuilder, LookupBus},
3    p3_field::PrimeCharacteristicRing,
4};
5
6#[derive(Clone, Copy, Debug, PartialEq, Eq)]
7pub struct DeferralCircuitCountBus(LookupBus);
8
9impl DeferralCircuitCountBus {
10    pub const fn new(index: BusIndex) -> Self {
11        Self(LookupBus::new(index))
12    }
13
14    #[inline(always)]
15    pub fn index(&self) -> BusIndex {
16        self.0.index
17    }
18
19    #[must_use]
20    pub fn send<T>(&self, deferral_idx: impl Into<T>) -> DeferralCircuitCountInteraction<T> {
21        self.push(deferral_idx, true)
22    }
23
24    #[must_use]
25    pub fn receive<T>(&self, deferral_idx: impl Into<T>) -> DeferralCircuitCountInteraction<T> {
26        self.push(deferral_idx, false)
27    }
28
29    pub fn push<T>(
30        &self,
31        deferral_idx: impl Into<T>,
32        is_lookup: bool,
33    ) -> DeferralCircuitCountInteraction<T> {
34        DeferralCircuitCountInteraction {
35            deferral_idx: deferral_idx.into(),
36            bus: self.0,
37            is_lookup,
38        }
39    }
40}
41
42#[derive(Clone, Copy, Debug)]
43pub struct DeferralCircuitCountInteraction<T> {
44    pub deferral_idx: T,
45    pub bus: LookupBus,
46    pub is_lookup: bool,
47}
48
49impl<T: PrimeCharacteristicRing> DeferralCircuitCountInteraction<T> {
50    pub fn eval<AB>(self, builder: &mut AB, count: impl Into<AB::Expr>)
51    where
52        AB: InteractionBuilder<Expr = T>,
53    {
54        let key = [self.deferral_idx];
55        if self.is_lookup {
56            self.bus.lookup_key(builder, key, count);
57        } else {
58            self.bus.add_key_with_lookups(builder, key, count);
59        }
60    }
61}