-
Notifications
You must be signed in to change notification settings - Fork 38
/
air_ext.rs
126 lines (116 loc) · 4.22 KB
/
air_ext.rs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
use std::iter::zip;
use itertools::Itertools;
use super::accumulation::{DomainEvaluationAccumulator, PointEvaluationAccumulator};
use super::{Air, AirProver, ComponentTrace};
use crate::core::backend::Backend;
use crate::core::channel::{Blake2sChannel, Channel};
use crate::core::circle::CirclePoint;
use crate::core::fields::m31::BaseField;
use crate::core::fields::qm31::SecureField;
use crate::core::poly::circle::{CanonicCoset, CircleEvaluation, CirclePoly, SecureCirclePoly};
use crate::core::poly::BitReversedOrder;
use crate::core::prover::LOG_BLOWUP_FACTOR;
use crate::core::{ComponentVec, InteractionElements};
pub trait AirExt: Air {
fn composition_log_degree_bound(&self) -> u32 {
self.components()
.iter()
.map(|component| component.max_constraint_log_degree_bound())
.max()
.unwrap()
}
fn trace_commitment_domains(&self) -> Vec<CanonicCoset> {
self.column_log_sizes()
.iter()
.map(|&log_size| CanonicCoset::new(log_size + LOG_BLOWUP_FACTOR))
.collect_vec()
}
fn mask_points(
&self,
point: CirclePoint<SecureField>,
) -> ComponentVec<Vec<CirclePoint<SecureField>>> {
let mut component_points = ComponentVec(Vec::new());
for component in self.components() {
let points = component.mask_points(point);
component_points.push(points);
}
component_points
}
fn interaction_elements(&self, channel: &mut Blake2sChannel) -> InteractionElements {
InteractionElements(
self.components()
.iter()
.flat_map(|component| component.interaction_element_ids())
.sorted()
.dedup()
.map(|id| {
let element = channel.draw_felt().0 .0;
(id, element)
})
.collect(),
)
}
fn eval_composition_polynomial_at_point(
&self,
point: CirclePoint<SecureField>,
mask_values: &ComponentVec<Vec<SecureField>>,
random_coeff: SecureField,
) -> SecureField {
let mut evaluation_accumulator = PointEvaluationAccumulator::new(random_coeff);
zip(self.components(), &mask_values.0).for_each(|(component, mask)| {
component.evaluate_constraint_quotients_at_point(
point,
mask,
&mut evaluation_accumulator,
)
});
evaluation_accumulator.finalize()
}
fn column_log_sizes(&self) -> Vec<u32> {
self.components()
.iter()
.flat_map(|component| component.trace_log_degree_bounds())
.collect()
}
fn component_traces<'a, B: Backend>(
&'a self,
polynomials: &'a [CirclePoly<B>],
evals: &'a [CircleEvaluation<B, BaseField, BitReversedOrder>],
) -> Vec<ComponentTrace<'_, B>> {
let poly_iter = &mut polynomials.iter();
let eval_iter = &mut evals.iter();
self.components()
.iter()
.map(|component| {
let n_columns = component.trace_log_degree_bounds().len();
let polys = poly_iter.take(n_columns).collect();
let evals = eval_iter.take(n_columns).collect();
ComponentTrace::new(polys, evals)
})
.collect()
}
}
impl<A: Air + ?Sized> AirExt for A {}
pub trait AirProverExt<B: Backend>: AirProver<B> {
fn compute_composition_polynomial(
&self,
random_coeff: SecureField,
component_traces: &[ComponentTrace<'_, B>],
) -> SecureCirclePoly<B> {
let total_constraints: usize = self
.prover_components()
.iter()
.map(|c| c.n_constraints())
.sum();
let mut accumulator = DomainEvaluationAccumulator::new(
random_coeff,
self.composition_log_degree_bound(),
total_constraints,
);
zip(self.prover_components(), component_traces).for_each(|(component, trace)| {
component.evaluate_constraint_quotients_on_domain(trace, &mut accumulator)
});
accumulator.finalize()
}
}
impl<B: Backend, A: AirProver<B>> AirProverExt<B> for A {}