stwo_gpu/core/air/
components.rs

1use core::iter::zip;
2
3use std_shims::{vec, Vec};
4
5use super::accumulation::PointEvaluationAccumulator;
6use super::Component;
7use crate::core::circle::CirclePoint;
8use crate::core::fields::qm31::SecureField;
9use crate::core::pcs::TreeVec;
10use crate::core::verifier::PREPROCESSED_TRACE_IDX;
11use crate::core::ColumnVec;
12
13pub struct Components<'a> {
14    pub components: Vec<&'a dyn Component>,
15    pub n_preprocessed_columns: usize,
16}
17
18impl Components<'_> {
19    pub fn composition_log_degree_bound(&self) -> u32 {
20        self.components
21            .iter()
22            .map(|component| component.max_constraint_log_degree_bound())
23            .max()
24            .unwrap()
25    }
26
27    pub fn mask_points(
28        &self,
29        point: CirclePoint<SecureField>,
30        max_log_degree_bound: u32,
31        include_all_preprocessed_columns: bool,
32    ) -> TreeVec<ColumnVec<Vec<CirclePoint<SecureField>>>> {
33        let mut mask_points = TreeVec::concat_cols(
34            self.components
35                .iter()
36                .map(|component| component.mask_points(point, max_log_degree_bound)),
37        );
38
39        let preprocessed_mask_points = &mut mask_points[PREPROCESSED_TRACE_IDX];
40        if include_all_preprocessed_columns {
41            *preprocessed_mask_points = vec![vec![point]; self.n_preprocessed_columns];
42        } else {
43            *preprocessed_mask_points = vec![vec![]; self.n_preprocessed_columns];
44            for component in &self.components {
45                for idx in component.preprocessed_column_indices() {
46                    preprocessed_mask_points[idx] = vec![point];
47                }
48            }
49        }
50
51        mask_points
52    }
53
54    pub fn eval_composition_polynomial_at_point(
55        &self,
56        point: CirclePoint<SecureField>,
57        mask_values: &TreeVec<Vec<Vec<SecureField>>>,
58        random_coeff: SecureField,
59        max_log_degree_bound: u32,
60    ) -> SecureField {
61        let mut evaluation_accumulator = PointEvaluationAccumulator::new(random_coeff);
62        for component in &self.components {
63            component.evaluate_constraint_quotients_at_point(
64                point,
65                mask_values,
66                &mut evaluation_accumulator,
67                max_log_degree_bound,
68            )
69        }
70        evaluation_accumulator.finalize()
71    }
72
73    pub fn column_log_sizes(&self) -> TreeVec<ColumnVec<u32>> {
74        let mut preprocessed_columns_trace_log_sizes = vec![0; self.n_preprocessed_columns];
75        let mut visited_columns = vec![false; self.n_preprocessed_columns];
76
77        let mut column_log_sizes = TreeVec::concat_cols(self.components.iter().map(|component| {
78            let component_trace_log_sizes = component.trace_log_degree_bounds();
79
80            for (column_index, &log_size) in zip(
81                component.preprocessed_column_indices(),
82                &component_trace_log_sizes[PREPROCESSED_TRACE_IDX],
83            ) {
84                let column_log_size = &mut preprocessed_columns_trace_log_sizes[column_index];
85                if visited_columns[column_index] {
86                    assert!(
87                        *column_log_size == log_size,
88                        "Preprocessed column size mismatch for column {column_index}"
89                    );
90                } else {
91                    *column_log_size = log_size;
92                    visited_columns[column_index] = true;
93                }
94            }
95
96            component_trace_log_sizes
97        }));
98
99        assert!(
100            visited_columns.iter().all(|&updated| updated),
101            "Column size not set for all reprocessed columns"
102        );
103
104        column_log_sizes[PREPROCESSED_TRACE_IDX] = preprocessed_columns_trace_log_sizes;
105
106        column_log_sizes
107    }
108}
stwo_gpu/core/air/components.rs

stwo_gpu/core/air/
components.rs
