winter_prover/constraints/
composition_poly.rs

1// Copyright (c) Facebook, Inc. and its affiliates.
2//
3// This source code is licensed under the MIT license found in the
4// LICENSE file in the root directory of this source tree.
5
6use alloc::vec::Vec;
7
8use air::proof::QuotientOodFrame;
9use math::{fft, polynom::degree_of, FieldElement, StarkField};
10
11use super::{ColMatrix, StarkDomain};
12
13// CONSTRAINT COMPOSITION POLYNOMIAL TRACE
14// ================================================================================================
15
16/// Represents merged evaluations of all constraint evaluations.
17pub struct CompositionPolyTrace<E>(Vec<E>);
18
19impl<E: FieldElement> CompositionPolyTrace<E> {
20    /// Returns a new instance of [CompositionPolyTrace] instantiated from the provided evaluations.
21    ///
22    /// # Panics
23    /// Panics if the number of evaluations is not a power of 2.
24    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    /// Returns the number of evaluations in this trace.
35    pub fn num_rows(&self) -> usize {
36        self.0.len()
37    }
38
39    /// Returns the internal vector representing this trace.
40    pub fn into_inner(self) -> Vec<E> {
41        self.0
42    }
43}
44
45// CONSTRAINT COMPOSITION POLYNOMIAL
46// ================================================================================================
47/// A composition polynomial split into columns with each column being of length equal to
48/// trace_length.
49///
50/// For example, if the composition polynomial has degree 2N - 1, where N is the trace length,
51/// it will be stored as two columns of size N (each of degree N - 1).
52pub struct CompositionPoly<E: FieldElement> {
53    data: ColMatrix<E>,
54}
55
56impl<E: FieldElement> CompositionPoly<E> {
57    /// Returns a new composition polynomial.
58    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        // at this point, combined_poly contains evaluations of the combined constraint polynomial;
71        // we interpolate this polynomial to transform it into coefficient form.
72        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    // PUBLIC ACCESSORS
81    // --------------------------------------------------------------------------------------------
82
83    /// Returns the number of individual column polynomials used to describe this composition
84    /// polynomial.
85    pub fn num_columns(&self) -> usize {
86        self.data.num_cols()
87    }
88
89    /// Returns the length of individual column polynomials; this is guaranteed to be a power of 2.
90    pub fn column_len(&self) -> usize {
91        self.data.num_rows()
92    }
93
94    /// Returns the degree of individual column polynomial.
95    #[allow(unused)]
96    pub fn column_degree(&self) -> usize {
97        self.column_len() - 1
98    }
99
100    /// Returns evaluations of all composition polynomial columns at points `z` and `g * z`.
101    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    /// Returns a reference to the matrix of individual column polynomials.
111    pub fn data(&self) -> &ColMatrix<E> {
112        &self.data
113    }
114
115    /// Transforms this composition polynomial into a vector of individual column polynomials.
116    pub fn into_columns(self) -> Vec<Vec<E>> {
117        self.data.into_columns()
118    }
119}
120
121// HELPER FUNCTIONS
122// ================================================================================================
123
124/// Splits polynomial coefficients into the specified number of columns. The coefficients are split
125/// in such a way that each resulting column has the same degree. For example, a polynomial
126/// a * x^3 + b * x^2 + c * x + d, can be rewritten as: (c * x + d) + x^2 * (a * x + b), and then
127/// the two columns will be: (c * x + d) and (a * x + b).
128fn 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// TESTS
143// ================================================================================================
144
145#[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}