winter_prover/trace/
poly_table.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::TraceOodFrame;
9use math::{FieldElement, StarkField};
10
11use crate::{matrix::ColumnIter, ColMatrix};
12
13// TRACE POLYNOMIAL TABLE
14// ================================================================================================
15
16/// Trace polynomials in coefficient from for all segments of the execution trace.
17///
18/// Coefficients of the polynomials for the main trace segment are always in the base field.
19/// However, coefficients of the polynomials for the auxiliary trace segment may be either in
20/// the base field, or in the extension field, depending on whether extension field is being used.
21pub struct TracePolyTable<E: FieldElement> {
22    main_trace_polys: ColMatrix<E::BaseField>,
23    aux_trace_polys: Option<ColMatrix<E>>,
24}
25
26impl<E: FieldElement> TracePolyTable<E> {
27    // CONSTRUCTOR
28    // --------------------------------------------------------------------------------------------
29    /// Creates a new table of trace polynomials from the provided main trace segment polynomials.
30    pub fn new(main_trace_polys: ColMatrix<E::BaseField>) -> Self {
31        Self { main_trace_polys, aux_trace_polys: None }
32    }
33
34    // STATE MUTATORS
35    // --------------------------------------------------------------------------------------------
36
37    /// Adds the provided auxiliary segment polynomials to this polynomial table.
38    pub fn add_aux_segment(&mut self, aux_trace_polys: ColMatrix<E>) {
39        assert!(self.aux_trace_polys.is_none());
40        assert_eq!(
41            self.main_trace_polys.num_rows(),
42            aux_trace_polys.num_rows(),
43            "polynomials in auxiliary segment must be of the same size as in the main segment"
44        );
45
46        self.aux_trace_polys = Some(aux_trace_polys);
47    }
48
49    // PUBLIC ACCESSORS
50    // --------------------------------------------------------------------------------------------
51
52    /// Returns the size of each polynomial - i.e. size of a vector needed to hold a polynomial.
53    pub fn poly_size(&self) -> usize {
54        self.main_trace_polys.num_rows()
55    }
56
57    /// Evaluates all trace polynomials (across all trace segments) at the specified point `x`.
58    pub fn evaluate_at(&self, x: E) -> Vec<E> {
59        let mut result = self.main_trace_polys.evaluate_columns_at(x);
60        for aux_polys in self.aux_trace_polys.iter() {
61            result.append(&mut aux_polys.evaluate_columns_at(x));
62        }
63        result
64    }
65
66    /// Returns an out-of-domain evaluation frame constructed by evaluating trace polynomials for
67    /// all columns at points z and z * g, where g is the generator of the trace domain.
68    pub fn get_ood_frame(&self, z: E) -> TraceOodFrame<E> {
69        let log_trace_len = self.poly_size().ilog2();
70        let g = E::from(E::BaseField::get_root_of_unity(log_trace_len));
71        let current_row = self.evaluate_at(z);
72        let next_row = self.evaluate_at(z * g);
73        let main_trace_width = self.main_trace_polys.num_cols();
74
75        TraceOodFrame::new(current_row, next_row, main_trace_width)
76    }
77
78    /// Returns an iterator over the polynomials of the main trace segment.
79    pub fn main_trace_polys(&self) -> impl Iterator<Item = &[E::BaseField]> {
80        self.main_trace_polys.columns()
81    }
82
83    /// Returns an iterator over the polynomials of the auxiliary trace segment.
84    pub fn aux_trace_polys(&self) -> impl Iterator<Item = &[E]> {
85        match self.aux_trace_polys {
86            Some(ref aux_segment_polys) => aux_segment_polys.columns(),
87            None => ColumnIter::empty(),
88        }
89    }
90
91    // TEST HELPERS
92    // --------------------------------------------------------------------------------------------
93
94    /// Returns the number of polynomials in the main segment of the trace.
95    #[cfg(test)]
96    pub fn num_main_trace_polys(&self) -> usize {
97        self.main_trace_polys.num_cols()
98    }
99
100    /// Returns a polynomial from the main segment of the trace at the specified index.
101    #[cfg(test)]
102    pub fn get_main_trace_poly(&self, idx: usize) -> &[E::BaseField] {
103        self.main_trace_polys.get_column(idx)
104    }
105}
winter_prover/trace/poly_table.rs

winter_prover/trace/
poly_table.rs
