openvm_stark_backend/test_utils/dummy_airs/interaction/
dummy_interaction_air.rs

1//! Air with columns
2//! | count | fields[..] |
3//!
4//! Chip will either send or receive the fields with multiplicity count.
5
6use std::{iter, sync::Arc};
7
8use itertools::izip;
9use p3_air::{Air, BaseAir, BaseAirWithPublicValues};
10use p3_field::PrimeCharacteristicRing;
11use p3_matrix::{dense::RowMajorMatrix, Matrix};
12
13use crate::{
14    air_builders::PartitionedAirBuilder,
15    interaction::{BusIndex, InteractionBuilder},
16    prover::{
17        stacked_pcs::stacked_commit, AirProvingContext, ColMajorMatrix, CommittedTraceData,
18        CpuColMajorBackend,
19    },
20    AirRef, PartitionedBaseAir, StarkProtocolConfig,
21};
22
23pub struct DummyInteractionCols;
24impl DummyInteractionCols {
25    pub fn count_col() -> usize {
26        0
27    }
28    pub fn field_col(field_idx: usize) -> usize {
29        field_idx + 1
30    }
31}
32
33#[derive(Clone, Copy)]
34pub struct DummyInteractionAir {
35    field_width: usize,
36    /// Send if true. Receive if false.
37    pub is_send: bool,
38    bus_index: BusIndex,
39    pub count_weight: u32,
40    /// If true, then | count | and | fields[..] | are in separate main trace partitions.
41    pub partition: bool,
42}
43
44impl DummyInteractionAir {
45    pub fn new(field_width: usize, is_send: bool, bus_index: BusIndex) -> Self {
46        Self {
47            field_width,
48            is_send,
49            bus_index,
50            count_weight: 0,
51            partition: false,
52        }
53    }
54
55    pub fn partition(self) -> Self {
56        Self {
57            partition: true,
58            ..self
59        }
60    }
61
62    pub fn field_width(&self) -> usize {
63        self.field_width
64    }
65}
66
67impl<F> BaseAirWithPublicValues<F> for DummyInteractionAir {}
68impl<F> PartitionedBaseAir<F> for DummyInteractionAir {
69    fn cached_main_widths(&self) -> Vec<usize> {
70        if self.partition {
71            vec![self.field_width]
72        } else {
73            vec![]
74        }
75    }
76    fn common_main_width(&self) -> usize {
77        if self.partition {
78            1
79        } else {
80            1 + self.field_width
81        }
82    }
83}
84impl<F> BaseAir<F> for DummyInteractionAir {
85    fn width(&self) -> usize {
86        1 + self.field_width
87    }
88
89    fn preprocessed_trace(&self) -> Option<RowMajorMatrix<F>> {
90        None
91    }
92}
93
94impl<AB: InteractionBuilder + PartitionedAirBuilder> Air<AB> for DummyInteractionAir {
95    fn eval(&self, builder: &mut AB) {
96        let (fields, count) = if self.partition {
97            let local_0 = builder.common_main().row_slice(0).unwrap();
98            let local_1 = builder.cached_mains()[0].row_slice(0).unwrap();
99            let count = local_0[0];
100            let fields = local_1.to_vec();
101            (fields, count)
102        } else {
103            let main = builder.main();
104            let local = main.row_slice(0).expect("window should have two elements");
105            let count = local[DummyInteractionCols::count_col()];
106            let fields: Vec<_> = (0..self.field_width)
107                .map(|i| local[DummyInteractionCols::field_col(i)])
108                .collect();
109            (fields, count)
110        };
111        if self.is_send {
112            builder.push_interaction(self.bus_index, fields, count, self.count_weight);
113        } else {
114            builder.push_interaction(
115                self.bus_index,
116                fields,
117                AB::Expr::NEG_ONE * count,
118                self.count_weight,
119            );
120        }
121    }
122}
123
124/// Note: in principle, committing cached trace is out of scope of a chip. But this chip is for
125/// testing, so we support it for convenience.
126pub struct DummyInteractionChip<SC> {
127    config: Option<SC>,
128    data: Option<DummyInteractionData>,
129    pub air: DummyInteractionAir,
130}
131
132#[derive(Debug, Clone)]
133pub struct DummyInteractionData {
134    pub count: Vec<u32>,
135    pub fields: Vec<Vec<u32>>,
136}
137
138impl<SC> DummyInteractionChip<SC> {
139    pub fn new_without_partition(field_width: usize, is_send: bool, bus_index: BusIndex) -> Self {
140        let air = DummyInteractionAir::new(field_width, is_send, bus_index);
141        Self {
142            config: None,
143            data: None,
144            air,
145        }
146    }
147    pub fn new_with_partition(
148        config: SC,
149        field_width: usize,
150        is_send: bool,
151        bus_index: BusIndex,
152    ) -> Self {
153        let air = DummyInteractionAir::new(field_width, is_send, bus_index).partition();
154        Self {
155            config: Some(config),
156            data: None,
157            air,
158        }
159    }
160    pub fn load_data(&mut self, data: DummyInteractionData) {
161        let DummyInteractionData { count, fields } = &data;
162        let h = count.len();
163        assert_eq!(fields.len(), h);
164        let w = fields[0].len();
165        assert_eq!(self.air.field_width, w);
166        assert!(fields.iter().all(|r| r.len() == w));
167        self.data = Some(data);
168    }
169    pub fn air(&self) -> AirRef<SC>
170    where
171        SC: StarkProtocolConfig,
172    {
173        Arc::new(self.air)
174    }
175}
176
177impl<SC: StarkProtocolConfig> DummyInteractionChip<SC> {
178    pub fn generate_proving_ctx(&self) -> AirProvingContext<CpuColMajorBackend<SC>> {
179        assert!(self.data.is_some());
180        let data = self.data.clone().unwrap();
181        if self.air.partition {
182            self.generate_traces_with_partition(data)
183        } else {
184            let trace = self.generate_traces_without_partition(data);
185            AirProvingContext::simple_no_pis(ColMajorMatrix::from_row_major(&trace))
186        }
187    }
188}
189
190impl<SC: StarkProtocolConfig> DummyInteractionChip<SC> {
191    pub fn generate_traces_with_partition(
192        &self,
193        data: DummyInteractionData,
194    ) -> AirProvingContext<CpuColMajorBackend<SC>> {
195        let DummyInteractionData {
196            mut count,
197            mut fields,
198        } = data;
199        let h = count.len();
200        assert_eq!(fields.len(), h);
201        let w = fields[0].len();
202        assert_eq!(self.air.field_width, w);
203        assert!(fields.iter().all(|r| r.len() == w));
204        let h = h.next_power_of_two();
205        count.resize(h, 0);
206        fields.resize(h, vec![0; w]);
207        let common_main_val: Vec<_> = count.into_iter().map(SC::F::from_u32).collect();
208        let cached_trace_val: Vec<_> = fields.into_iter().flatten().map(SC::F::from_u32).collect();
209        let cached_trace_rm = RowMajorMatrix::new(cached_trace_val, w);
210        let cached_trace = ColMajorMatrix::from_row_major(&cached_trace_rm);
211
212        let config = self.config.as_ref().expect("params required for partition");
213        let params = config.params();
214        let (commit, data) = stacked_commit(
215            config.hasher(),
216            params.l_skip,
217            params.n_stack,
218            params.log_blowup,
219            params.k_whir(),
220            &[&cached_trace],
221        )
222        .unwrap();
223
224        let common_main_rm = RowMajorMatrix::new(common_main_val, 1);
225        AirProvingContext {
226            cached_mains: vec![CommittedTraceData {
227                commitment: commit,
228                trace: cached_trace,
229                data: Arc::new(data),
230            }],
231            common_main: ColMajorMatrix::from_row_major(&common_main_rm),
232            public_values: vec![],
233        }
234    }
235
236    fn generate_traces_without_partition(
237        &self,
238        data: DummyInteractionData,
239    ) -> RowMajorMatrix<SC::F> {
240        let DummyInteractionData { count, fields } = data;
241        let h = count.len();
242        assert_eq!(fields.len(), h);
243        let w = fields[0].len();
244        assert_eq!(self.air.field_width, w);
245        assert!(fields.iter().all(|r| r.len() == w));
246        let common_main_val: Vec<_> = izip!(count, fields)
247            .flat_map(|(count, fields)| iter::once(count).chain(fields))
248            .chain(iter::repeat(0))
249            .take((w + 1) * h.next_power_of_two())
250            .map(SC::F::from_u32)
251            .collect();
252        RowMajorMatrix::new(common_main_val, w + 1)
253    }
254}
openvm_stark_backend/test_utils/dummy_airs/interaction/dummy_interaction_air.rs

openvm_stark_backend/test_utils/dummy_airs/interaction/
dummy_interaction_air.rs
