trace to evals

starkware-libs · Dec 23, 2024 · a42bd83 · a42bd83
1 parent 8ade304
commit a42bd83
Showing 1 changed file with 16 additions and 12 deletions.
diff --git a/crates/air_utils/src/trace/iterable_trace.rs b/crates/air_utils/src/trace/iterable_trace.rs
@@ -4,10 +4,11 @@ use bytemuck::{cast_slice, Zeroable};
 use itertools::Itertools;
 use rayon::iter::plumbing::{bridge, Consumer, Producer, ProducerCallback, UnindexedConsumer};
 use rayon::prelude::*;
+use stwo_prover::core::backend::simd::column::BaseColumn;
 use stwo_prover::core::backend::simd::m31::{PackedM31, N_LANES};
 use stwo_prover::core::backend::simd::SimdBackend;
 use stwo_prover::core::fields::m31::M31;
-use stwo_prover::core::poly::circle::CircleEvaluation;
+use stwo_prover::core::poly::circle::{CanonicCoset, CircleEvaluation};
 use stwo_prover::core::poly::BitReversedOrder;
 
 /// A 2D Matrix of [`PackedM31`] values.
@@ -51,16 +52,16 @@ use stwo_prover::core::poly::BitReversedOrder;
 pub struct IterableTrace<const N: usize> {
     data: [Vec<PackedM31>; N],
 
-    /// Number of M31 rows in each column.
-    length: usize,
+    /// Log number of M31 rows in each column.
+    log_size: u32,
 }
 
 impl<const N: usize> IterableTrace<N> {
     pub fn zeroed(log_size: u32) -> Self {
         let length = 1 << log_size;
         let n_simd_elems = length / N_LANES;
         let data = [(); N].map(|_| vec![PackedM31::zeroed(); n_simd_elems]);
-        Self { data, length }
+        Self { data, log_size }
     }
 
     /// # Safety
@@ -70,12 +71,8 @@ impl<const N: usize> IterableTrace<N> {
         todo!()
     }
 
-    pub fn len(&self) -> usize {
-        self.length
-    }
-
-    pub fn is_empty(&self) -> bool {
-        self.length == 0
+    pub fn log_size(&self) -> u32 {
+        self.log_size
     }
 
     pub fn iter_mut(&mut self) -> RowIterMut<'_, N> {
@@ -101,11 +98,18 @@ impl<const N: usize> IterableTrace<N> {
     }
 
     pub fn to_evals(self) -> [CircleEvaluation<SimdBackend, M31, BitReversedOrder>; N] {
-        todo!()
+        let domain = CanonicCoset::new(self.log_size).circle_domain();
+        self.data.map(|column| {
+            let eval = BaseColumn {
+                data: column,
+                length: 1 << self.log_size,
+            };
+            CircleEvaluation::<SimdBackend, M31, BitReversedOrder>::new(domain, eval)
+        })
     }
 
     pub fn pretty_print(&self, row_limit: usize) -> String {
-        assert!(row_limit <= self.len());
+        assert!(row_limit <= 1 << self.log_size);
         let cpu_trace: Vec<&[u32]> = self
             .data
             .iter()