openvm_stark_backend/keygen/
view.rs

1use itertools::Itertools;
2use p3_field::{ExtensionField, Field};
3
4use crate::{
5    config::{Com, StarkGenericConfig, Val},
6    keygen::types::{LinearConstraint, MultiStarkVerifyingKey, StarkVerifyingKey},
7};
8
9#[derive(Clone, derive_new::new)]
10pub struct MultiStarkVerifyingKeyView<'a, Val, Com> {
11    pub per_air: Vec<&'a StarkVerifyingKey<Val, Com>>,
12    /// Trace height constraints are *not* filtered by AIR. When computing the dot product, this
13    /// will be indexed into by air_id.
14    pub trace_height_constraints: &'a [LinearConstraint],
15    pub pre_hash: Com,
16}
17
18impl<SC: StarkGenericConfig> MultiStarkVerifyingKey<SC> {
19    /// Returns a view with all airs.
20    pub(crate) fn full_view(&self) -> MultiStarkVerifyingKeyView<'_, Val<SC>, Com<SC>> {
21        self.view(&(0..self.inner.per_air.len()).collect_vec())
22    }
23    pub(crate) fn view(
24        &self,
25        air_ids: &[usize],
26    ) -> MultiStarkVerifyingKeyView<'_, Val<SC>, Com<SC>> {
27        MultiStarkVerifyingKeyView {
28            per_air: air_ids.iter().map(|&id| &self.inner.per_air[id]).collect(),
29            trace_height_constraints: &self.inner.trace_height_constraints,
30            pre_hash: self.pre_hash.clone(),
31        }
32    }
33}
34
35impl<Val, Com: Clone> MultiStarkVerifyingKeyView<'_, Val, Com> {
36    /// Returns the preprocessed commit of each AIR. If the AIR does not have a preprocessed trace,
37    /// returns None.
38    pub fn preprocessed_commits(&self) -> Vec<Option<Com>> {
39        self.per_air
40            .iter()
41            .map(|vk| {
42                vk.preprocessed_data
43                    .as_ref()
44                    .map(|data| data.commit.clone())
45            })
46            .collect()
47    }
48
49    /// Returns all non-empty preprocessed commits.
50    pub fn flattened_preprocessed_commits(&self) -> Vec<Com> {
51        self.preprocessed_commits().into_iter().flatten().collect()
52    }
53
54    pub fn num_phases(&self) -> usize {
55        self.per_air
56            .iter()
57            .map(|vk| {
58                // Consistency check
59                let num = vk.params.width.after_challenge.len();
60                assert_eq!(num, vk.params.num_challenges_to_sample.len());
61                assert_eq!(num, vk.params.num_exposed_values_after_challenge.len());
62                num
63            })
64            .max()
65            .unwrap_or(0)
66    }
67
68    pub fn num_challenges_per_phase(&self) -> Vec<usize> {
69        let num_phases = self.num_phases();
70        (0..num_phases)
71            .map(|phase_idx| self.num_challenges_in_phase(phase_idx))
72            .collect()
73    }
74
75    pub fn num_challenges_in_phase(&self, phase_idx: usize) -> usize {
76        self.per_air
77            .iter()
78            .flat_map(|vk| vk.params.num_challenges_to_sample.get(phase_idx))
79            .copied()
80            .max()
81            .unwrap_or_else(|| panic!("No challenges used in challenge phase {phase_idx}"))
82    }
83
84    /// Returns the main width for each AIR.
85    pub fn main_widths(&self) -> Vec<usize> {
86        self.per_air
87            .iter()
88            .map(|vk| vk.params.width.main_width())
89            .collect()
90    }
91
92    /// Returns the total width for each AIR.
93    pub fn total_widths<E>(&self) -> Vec<usize>
94    where
95        Val: Field,
96        E: ExtensionField<Val>,
97    {
98        self.per_air
99            .iter()
100            .map(|vk| vk.params.width.total_width(E::D))
101            .collect()
102    }
103
104    /// Returns the number of interactions for each AIR.
105    pub fn num_interactions(&self) -> Vec<usize> {
106        self.per_air
107            .iter()
108            .map(|vk| vk.symbolic_constraints.interactions.len())
109            .collect()
110    }
111}