openvm_deferral_circuit/poseidon2/
bus.rs

1use std::iter::once;
2
3use itertools::Itertools;
4use openvm_circuit::system::poseidon2::PERIPHERY_POSEIDON2_CHUNK_SIZE;
5use openvm_stark_backend::{
6    interaction::{BusIndex, InteractionBuilder, LookupBus},
7    p3_field::PrimeCharacteristicRing,
8};
9
10#[derive(Clone, Copy, Debug, PartialEq, Eq)]
11pub struct DeferralPoseidon2Bus(pub LookupBus);
12
13impl DeferralPoseidon2Bus {
14    pub const fn new(index: BusIndex) -> Self {
15        Self(LookupBus::new(index))
16    }
17
18    #[inline(always)]
19    pub fn index(&self) -> BusIndex {
20        self.0.index
21    }
22
23    #[must_use]
24    pub fn lookup<T>(
25        &self,
26        left: [impl Into<T>; PERIPHERY_POSEIDON2_CHUNK_SIZE],
27        right: [impl Into<T>; PERIPHERY_POSEIDON2_CHUNK_SIZE],
28        output: [impl Into<T>; PERIPHERY_POSEIDON2_CHUNK_SIZE],
29        is_compress: impl Into<T>,
30    ) -> DeferralPoseidon2Interaction<T> {
31        DeferralPoseidon2Interaction {
32            left: left.map(Into::into),
33            right: right.map(Into::into),
34            output: output.map(Into::into),
35            is_compress: is_compress.into(),
36            bus: self.0,
37        }
38    }
39}
40
41#[derive(Clone, Copy, Debug)]
42pub struct DeferralPoseidon2Interaction<T> {
43    pub left: [T; PERIPHERY_POSEIDON2_CHUNK_SIZE],
44    pub right: [T; PERIPHERY_POSEIDON2_CHUNK_SIZE],
45    pub output: [T; PERIPHERY_POSEIDON2_CHUNK_SIZE],
46    pub is_compress: T,
47    pub bus: LookupBus,
48}
49
50impl<T: PrimeCharacteristicRing> DeferralPoseidon2Interaction<T> {
51    pub fn eval<AB>(self, builder: &mut AB, count: impl Into<AB::Expr>)
52    where
53        AB: InteractionBuilder<Expr = T>,
54    {
55        let key: [_; 3 * PERIPHERY_POSEIDON2_CHUNK_SIZE + 1] = self
56            .left
57            .into_iter()
58            .chain(self.right)
59            .chain(self.output)
60            .chain(once(self.is_compress))
61            .collect_array()
62            .unwrap();
63        self.bus.lookup_key(builder, key, count);
64    }
65}