openvm_stark_backend/keygen/
view.rs

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