Skip to main content

openvm_stark_backend/test_utils/dummy_airs/interaction/
self_interaction_air.rs

1use itertools::{fold, Itertools};
2use p3_air::{Air, AirBuilder, BaseAir, BaseAirWithPublicValues};
3use p3_field::PrimeCharacteristicRing;
4use p3_matrix::{dense::RowMajorMatrix, Matrix};
5
6use crate::{
7    interaction::{BusIndex, InteractionBuilder},
8    prover::{AirProvingContext, ColMajorMatrix, CpuColMajorBackend},
9    PartitionedBaseAir, StarkProtocolConfig,
10};
11
12#[derive(Debug, Clone, Copy)]
13pub struct SelfInteractionAir {
14    pub width: usize,
15    pub bus_index: BusIndex,
16}
17
18impl<F> BaseAir<F> for SelfInteractionAir {
19    fn width(&self) -> usize {
20        self.width
21    }
22}
23impl<F> BaseAirWithPublicValues<F> for SelfInteractionAir {}
24impl<F> PartitionedBaseAir<F> for SelfInteractionAir {}
25
26impl<AB: AirBuilder + InteractionBuilder> Air<AB> for SelfInteractionAir {
27    fn eval(&self, builder: &mut AB) {
28        let main = builder.main();
29
30        let (local, next) = (
31            main.row_slice(0).expect("window should have two elements"),
32            main.row_slice(1).expect("window should have two elements"),
33        );
34        let mut local: Vec<<AB as AirBuilder>::Expr> =
35            (*local).iter().map(|v| (*v).into()).collect_vec();
36        let mut next: Vec<<AB as AirBuilder>::Expr> =
37            (*next).iter().map(|v| (*v).into()).collect_vec();
38
39        let local_sum = fold(&local, AB::Expr::ZERO, |acc, val| acc + val.clone());
40        let next_sum = fold(&local, AB::Expr::ZERO, |acc, val| acc + val.clone());
41
42        // Interaction where count is constant
43        builder.push_interaction(self.bus_index, local.clone(), AB::Expr::ONE, 1);
44        builder.push_interaction(self.bus_index, next.clone(), AB::Expr::NEG_ONE, 1);
45
46        // Interaction where count is an expression + common with interaction below
47        builder.push_interaction(self.bus_index, local.clone(), local_sum.clone(), 1);
48        builder.push_interaction(self.bus_index, next.clone(), -next_sum.clone(), 1);
49
50        // Interaction where count == fields[0]
51        builder.push_interaction(self.bus_index, local.clone(), local[0].clone(), 1);
52        builder.push_interaction(self.bus_index, next.clone(), -next[0].clone(), 1);
53
54        local.reverse();
55        next.reverse();
56
57        // Interaction where count is an expression + common with interaction above
58        builder.push_interaction(self.bus_index, local, local_sum, 1);
59        builder.push_interaction(self.bus_index, next, -next_sum, 1);
60    }
61}
62
63#[derive(Debug, Clone, Copy)]
64pub struct SelfInteractionChip {
65    pub width: usize,
66    pub log_height: usize,
67}
68
69impl SelfInteractionChip {
70    pub fn generate_proving_ctx<SC: StarkProtocolConfig>(
71        &self,
72    ) -> AirProvingContext<CpuColMajorBackend<SC>> {
73        assert!(self.width > 0);
74        let mut trace = vec![SC::F::ZERO; (1 << self.log_height) * self.width];
75        for (row_idx, chunk) in trace.chunks_mut(self.width).enumerate() {
76            for (i, val) in chunk.iter_mut().enumerate() {
77                *val = SC::F::from_usize((row_idx + i) % self.width);
78            }
79        }
80        let rm = RowMajorMatrix::new(trace, self.width);
81        AirProvingContext::simple_no_pis(ColMajorMatrix::from_row_major(&rm))
82    }
83}