openvm_stark_backend/test_utils/dummy_airs/interaction/
self_interaction_air.rs1use 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 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 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 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 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}