winter_prover/constraints/
composition_poly.rs1use alloc::vec::Vec;
7
8use air::proof::QuotientOodFrame;
9use math::{fft, polynom::degree_of, FieldElement, StarkField};
10
11use super::{ColMatrix, StarkDomain};
12
13pub struct CompositionPolyTrace<E>(Vec<E>);
18
19impl<E: FieldElement> CompositionPolyTrace<E> {
20 pub fn new(evaluations: Vec<E>) -> Self {
25 assert!(
26 evaluations.len().is_power_of_two(),
27 "length of composition polynomial trace must be a power of 2, but was {}",
28 evaluations.len(),
29 );
30
31 Self(evaluations)
32 }
33
34 pub fn num_rows(&self) -> usize {
36 self.0.len()
37 }
38
39 pub fn into_inner(self) -> Vec<E> {
41 self.0
42 }
43}
44
45pub struct CompositionPoly<E: FieldElement> {
53 data: ColMatrix<E>,
54}
55
56impl<E: FieldElement> CompositionPoly<E> {
57 pub fn new(
59 composition_trace: CompositionPolyTrace<E>,
60 domain: &StarkDomain<E::BaseField>,
61 num_cols: usize,
62 ) -> Self {
63 assert!(
64 domain.trace_length() < composition_trace.num_rows(),
65 "trace length must be smaller than length of composition polynomial trace"
66 );
67
68 let mut trace = composition_trace.into_inner();
69
70 let inv_twiddles = fft::get_inv_twiddles::<E::BaseField>(trace.len());
73 fft::interpolate_poly_with_offset(&mut trace, &inv_twiddles, domain.offset());
74
75 let polys = segment(trace, domain.trace_length(), num_cols);
76
77 CompositionPoly { data: ColMatrix::new(polys) }
78 }
79
80 pub fn num_columns(&self) -> usize {
86 self.data.num_cols()
87 }
88
89 pub fn column_len(&self) -> usize {
91 self.data.num_rows()
92 }
93
94 #[allow(unused)]
96 pub fn column_degree(&self) -> usize {
97 self.column_len() - 1
98 }
99
100 pub fn get_ood_frame(&self, z: E) -> QuotientOodFrame<E> {
102 let log_trace_len = self.column_len().ilog2();
103 let g = E::from(E::BaseField::get_root_of_unity(log_trace_len));
104 let current_row = self.data.evaluate_columns_at(z);
105 let next_row = self.data.evaluate_columns_at(z * g);
106
107 QuotientOodFrame::new(current_row, next_row)
108 }
109
110 pub fn data(&self) -> &ColMatrix<E> {
112 &self.data
113 }
114
115 pub fn into_columns(self) -> Vec<Vec<E>> {
117 self.data.into_columns()
118 }
119}
120
121fn segment<E: FieldElement>(
129 coefficients: Vec<E>,
130 trace_len: usize,
131 num_cols: usize,
132) -> Vec<Vec<E>> {
133 debug_assert!(degree_of(&coefficients) < trace_len * num_cols);
134
135 coefficients
136 .chunks(trace_len)
137 .take(num_cols)
138 .map(|slice| slice.to_vec())
139 .collect()
140}
141
142#[cfg(test)]
146mod tests {
147
148 use alloc::vec::Vec;
149
150 use math::fields::f128::BaseElement;
151
152 #[test]
153 fn segment() {
154 let values = (0u128..16).map(BaseElement::new).collect::<Vec<_>>();
155 let actual = super::segment(values, 4, 4);
156
157 #[rustfmt::skip]
158 let expected = vec![
159 vec![BaseElement::new(0), BaseElement::new(1), BaseElement::new(2), BaseElement::new(3)],
160 vec![BaseElement::new(4), BaseElement::new(5), BaseElement::new(6), BaseElement::new(7)],
161 vec![BaseElement::new(8), BaseElement::new(9), BaseElement::new(10), BaseElement::new(11)],
162 vec![BaseElement::new(12), BaseElement::new(13), BaseElement::new(14), BaseElement::new(15)],
163 ];
164
165 assert_eq!(expected, actual)
166 }
167}