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 math::{fft, polynom::degree_of, FieldElement};
9
10use super::{ColMatrix, StarkDomain};
11
12// CONSTRAINT COMPOSITION POLYNOMIAL TRACE
13// ================================================================================================
14
15/// Represents merged evaluations of all constraint evaluations.
16pub struct CompositionPolyTrace<E>(Vec<E>);
17
18impl<E: FieldElement> CompositionPolyTrace<E> {
19    /// Returns a new instance of [CompositionPolyTrace] instantiated from the provided evaluations.
20    ///
21    /// # Panics
22    /// Panics if the number of evaluations is not a power of 2.
23    pub fn new(evaluations: Vec<E>) -> Self {
24        assert!(
25            evaluations.len().is_power_of_two(),
26            "length of composition polynomial trace must be a power of 2, but was {}",
27            evaluations.len(),
28        );
29
30        Self(evaluations)
31    }
32
33    /// Returns the number of evaluations in this trace.
34    pub fn num_rows(&self) -> usize {
35        self.0.len()
36    }
37
38    /// Returns the internal vector representing this trace.
39    pub fn into_inner(self) -> Vec<E> {
40        self.0
41    }
42}
43
44// CONSTRAINT COMPOSITION POLYNOMIAL
45// ================================================================================================
46/// A composition polynomial split into columns with each column being of length equal to
47/// trace_length.
48///
49/// For example, if the composition polynomial has degree 2N - 1, where N is the trace length,
50/// it will be stored as two columns of size N (each of degree N - 1).
51pub struct CompositionPoly<E: FieldElement> {
52    data: ColMatrix<E>,
53}
54
55impl<E: FieldElement> CompositionPoly<E> {
56    /// Returns a new composition polynomial.
57    pub fn new(
58        composition_trace: CompositionPolyTrace<E>,
59        domain: &StarkDomain<E::BaseField>,
60        num_cols: usize,
61    ) -> Self {
62        assert!(
63            domain.trace_length() < composition_trace.num_rows(),
64            "trace length must be smaller than length of composition polynomial trace"
65        );
66
67        let mut trace = composition_trace.into_inner();
68
69        // at this point, combined_poly contains evaluations of the combined constraint polynomial;
70        // we interpolate this polynomial to transform it into coefficient form.
71        let inv_twiddles = fft::get_inv_twiddles::<E::BaseField>(trace.len());
72        fft::interpolate_poly_with_offset(&mut trace, &inv_twiddles, domain.offset());
73
74        let polys = segment(trace, domain.trace_length(), num_cols);
75
76        CompositionPoly { data: ColMatrix::new(polys) }
77    }
78
79    // PUBLIC ACCESSORS
80    // --------------------------------------------------------------------------------------------
81
82    /// Returns the number of individual column polynomials used to describe this composition
83    /// polynomial.
84    pub fn num_columns(&self) -> usize {
85        self.data.num_cols()
86    }
87
88    /// Returns the length of individual column polynomials; this is guaranteed to be a power of 2.
89    pub fn column_len(&self) -> usize {
90        self.data.num_rows()
91    }
92
93    /// Returns the degree of individual column polynomial.
94    #[allow(unused)]
95    pub fn column_degree(&self) -> usize {
96        self.column_len() - 1
97    }
98
99    /// Returns evaluations of all composition polynomial columns at point z.
100    pub fn evaluate_at(&self, z: E) -> Vec<E> {
101        self.data.evaluate_columns_at(z)
102    }
103
104    /// Returns a reference to the matrix of individual column polynomials.
105    pub fn data(&self) -> &ColMatrix<E> {
106        &self.data
107    }
108
109    /// Transforms this composition polynomial into a vector of individual column polynomials.
110    pub fn into_columns(self) -> Vec<Vec<E>> {
111        self.data.into_columns()
112    }
113}
114
115// HELPER FUNCTIONS
116// ================================================================================================
117
118/// Splits polynomial coefficients into the specified number of columns. The coefficients are split
119/// in such a way that each resulting column has the same degree. For example, a polynomial
120/// a * x^3 + b * x^2 + c * x + d, can be rewritten as: (c * x + d) + x^2 * (a * x + b), and then
121/// the two columns will be: (c * x + d) and (a * x + b).
122fn segment<E: FieldElement>(
123    coefficients: Vec<E>,
124    trace_len: usize,
125    num_cols: usize,
126) -> Vec<Vec<E>> {
127    debug_assert!(degree_of(&coefficients) < trace_len * num_cols);
128
129    coefficients
130        .chunks(trace_len)
131        .take(num_cols)
132        .map(|slice| slice.to_vec())
133        .collect()
134}
135
136// TESTS
137// ================================================================================================
138
139#[cfg(test)]
140mod tests {
141
142    use alloc::vec::Vec;
143
144    use math::fields::f128::BaseElement;
145
146    #[test]
147    fn segment() {
148        let values = (0u128..16).map(BaseElement::new).collect::<Vec<_>>();
149        let actual = super::segment(values, 4, 4);
150
151        #[rustfmt::skip]
152        let expected = vec![
153            vec![BaseElement::new(0), BaseElement::new(1), BaseElement::new(2), BaseElement::new(3)],
154            vec![BaseElement::new(4), BaseElement::new(5), BaseElement::new(6), BaseElement::new(7)],
155            vec![BaseElement::new(8), BaseElement::new(9), BaseElement::new(10), BaseElement::new(11)],
156            vec![BaseElement::new(12), BaseElement::new(13), BaseElement::new(14), BaseElement::new(15)],
157        ];
158
159        assert_eq!(expected, actual)
160    }
161}