Skip to main content

openvm_stark_backend/interaction/
debug.rs

1use std::collections::{BTreeMap, HashMap};
2
3use itertools::Itertools;
4use p3_field::Field;
5use p3_matrix::{dense::RowMajorMatrixView, Matrix};
6
7use super::{BusIndex, SymbolicInteraction};
8use crate::air_builders::symbolic::{
9    symbolic_expression::SymbolicEvaluator,
10    symbolic_variable::{Entry, SymbolicVariable},
11};
12
13/// The actual interactions that are sent/received during a single run
14/// of trace generation. For debugging purposes only.
15#[derive(Default, Clone, Debug)]
16pub struct LogicalInteractions<F: Field> {
17    /// Bus index => (fields => (air_idx, count))
18    #[allow(clippy::type_complexity)]
19    pub at_bus: BTreeMap<BusIndex, HashMap<Vec<F>, Vec<(usize, F)>>>,
20}
21
22pub fn generate_logical_interactions<F: Field>(
23    air_idx: usize,
24    all_interactions: &[SymbolicInteraction<F>],
25    preprocessed: &Option<RowMajorMatrixView<F>>,
26    partitioned_main: &[RowMajorMatrixView<F>],
27    public_values: &[F],
28    logical_interactions: &mut LogicalInteractions<F>,
29) {
30    if all_interactions.is_empty() {
31        return;
32    }
33
34    let height = partitioned_main[0].height();
35
36    for n in 0..height {
37        let evaluator = Evaluator {
38            preprocessed,
39            partitioned_main,
40            public_values,
41            height,
42            local_index: n,
43        };
44        for interaction in all_interactions {
45            let fields = interaction
46                .message
47                .iter()
48                .map(|expr| evaluator.eval_expr(expr))
49                .collect_vec();
50            let count = evaluator.eval_expr(&interaction.count);
51            if count.is_zero() {
52                continue;
53            }
54            logical_interactions
55                .at_bus
56                .entry(interaction.bus_index)
57                .or_default()
58                .entry(fields)
59                .or_default()
60                .push((air_idx, count));
61        }
62    }
63}
64
65struct Evaluator<'a, F: Field> {
66    pub preprocessed: &'a Option<RowMajorMatrixView<'a, F>>,
67    pub partitioned_main: &'a [RowMajorMatrixView<'a, F>],
68    pub public_values: &'a [F],
69    pub height: usize,
70    pub local_index: usize,
71}
72
73impl<F: Field> SymbolicEvaluator<F, F> for Evaluator<'_, F> {
74    fn eval_const(&self, c: F) -> F {
75        c
76    }
77    fn eval_var(&self, symbolic_var: SymbolicVariable<F>) -> F {
78        let n = self.local_index;
79        let height = self.height;
80        let index = symbolic_var.index;
81        match symbolic_var.entry {
82            Entry::Preprocessed { offset } => self
83                .preprocessed
84                .unwrap()
85                .get((n + offset) % height, index)
86                .expect("matrix index out of bounds"),
87            Entry::Main { part_index, offset } => self.partitioned_main[part_index]
88                .get((n + offset) % height, index)
89                .expect("matrix index out of bounds"),
90            Entry::Public => self.public_values[index],
91            Entry::Challenge => unreachable!("There should be no challenge variables"),
92        }
93    }
94    fn eval_is_first_row(&self) -> F {
95        unreachable!()
96    }
97    fn eval_is_last_row(&self) -> F {
98        unreachable!()
99    }
100    fn eval_is_transition(&self) -> F {
101        unreachable!()
102    }
103}