refactor: new-type for preprocessed trace (#281)

* move traits to utils

* type safety for const columns

* move constant to preprocessed struct

prover/src/machine2/chips/add.rs

@@ -171,15 +171,15 @@ impl MachineChip for AddChip {
 
 #[cfg(test)]
 mod test {
-    use crate::machine2::chips::CpuChip;
+    use crate::machine2::{chips::CpuChip, trace::PreprocessedTraces};
 
     use super::*;
     use nexus_vm::{
         riscv::{BasicBlock, BuiltinOpcode, Instruction, InstructionType, Opcode},
         trace::k_trace_direct,
     };
 
-    const LOG_SIZE: u32 = Traces::MIN_LOG_SIZE;
+    const LOG_SIZE: u32 = PreprocessedTraces::MIN_LOG_SIZE;
 
     #[rustfmt::skip]
     fn setup_basic_block_ir() -> Vec<BasicBlock>

prover/src/machine2/chips/bit_op.rs

@@ -37,7 +37,7 @@ use crate::machine2::{
     trace::{
         eval::{preprocessed_trace_eval, trace_eval, TraceEval},
         sidenote::SideNote,
-        ProgramStep, Traces, Word,
+        PreprocessedTraces, ProgramStep, Traces, Word,
     },
     traits::{ExecuteChip, MachineChip},
 };
@@ -186,7 +186,7 @@ impl MachineChip for BitOpChip {
     /// data[vec_row] contains sixteen rows. A single write_frac() adds sixteen rows.
     fn fill_interaction_trace(
         original_traces: &Traces,
-        preprocessed_trace: &Traces,
+        preprocessed_trace: &PreprocessedTraces,
         lookup_element: &LookupElements<MAX_LOOKUP_TUPLE_SIZE>,
     ) -> ColumnVec<CircleEvaluation<SimdBackend, BaseField, BitReversedOrder>> {
         let mut logup_trace_gen = LogupTraceGenerator::new(original_traces.log_size());
@@ -331,7 +331,7 @@ mod test {
         trace::k_trace_direct,
     };
 
-    const LOG_SIZE: u32 = Traces::MIN_LOG_SIZE;
+    const LOG_SIZE: u32 = PreprocessedTraces::MIN_LOG_SIZE;
 
     fn setup_basic_block_ir() -> Vec<BasicBlock> {
         #[rustfmt::skip]

prover/src/machine2/chips/range256.rs

@@ -23,7 +23,7 @@ use crate::machine2::{
     trace::{
         eval::{preprocessed_trace_eval, trace_eval},
         sidenote::SideNote,
-        ProgramStep, Traces,
+        PreprocessedTraces, ProgramStep, Traces,
     },
     traits::MachineChip,
 };
@@ -67,7 +67,7 @@ impl MachineChip for Range256Chip {
     /// data[vec_row] contains sixteen rows. A single write_frac() adds sixteen numbers.
     fn fill_interaction_trace(
         original_traces: &Traces,
-        preprocessed_traces: &Traces,
+        preprocessed_traces: &PreprocessedTraces,
         lookup_element: &LookupElements<12>,
     ) -> ColumnVec<CircleEvaluation<SimdBackend, BaseField, BitReversedOrder>> {
         let mut logup_trace_gen = LogupTraceGenerator::new(original_traces.log_size());
@@ -207,7 +207,7 @@ mod test {
                 &mut side_note,
             );
         }
-        let mut preprocessed_256_rows = Traces::empty_preprocessed_trace(LOG_SIZE);
+        let mut preprocessed_256_rows = PreprocessedTraces::empty(LOG_SIZE);
         preprocessed_256_rows.fill_is_first();
         preprocessed_256_rows.fill_range256();
         assert_chip::<Range256Chip>(traces, Some(preprocessed_256_rows));
@@ -231,7 +231,7 @@ mod test {
     }
 
     fn range256_chip_fail_out_of_range() {
-        const LOG_SIZE: u32 = Traces::MIN_LOG_SIZE;
+        const LOG_SIZE: u32 = PreprocessedTraces::MIN_LOG_SIZE;
         let (config, twiddles) = test_params(LOG_SIZE);
         let mut traces = Traces::new(LOG_SIZE);
         let mut side_note = SideNote::default();

prover/src/machine2/chips/range32.rs

@@ -22,7 +22,7 @@ use crate::machine2::{
     trace::{
         eval::{preprocessed_trace_eval, trace_eval},
         sidenote::SideNote,
-        ProgramStep, Traces,
+        PreprocessedTraces, ProgramStep, Traces,
     },
     traits::MachineChip,
 };
@@ -58,7 +58,7 @@ impl MachineChip for Range32Chip {
     /// data[vec_row] contains sixteen rows. A single write_frac() adds sixteen numbers.
     fn fill_interaction_trace(
         original_traces: &Traces,
-        preprocessed_traces: &Traces,
+        preprocessed_traces: &PreprocessedTraces,
         lookup_element: &LookupElements<12>,
     ) -> ColumnVec<CircleEvaluation<SimdBackend, BaseField, BitReversedOrder>> {
         let mut logup_trace_gen = LogupTraceGenerator::new(original_traces.log_size());
@@ -170,7 +170,7 @@ mod test {
                 &mut side_note,
             );
         }
-        let mut preprocessed_32_rows = Traces::empty_preprocessed_trace(LOG_SIZE);
+        let mut preprocessed_32_rows = PreprocessedTraces::empty(LOG_SIZE);
         preprocessed_32_rows.fill_is_first();
         preprocessed_32_rows.fill_range32();
         assert_chip::<Range32Chip>(traces, Some(preprocessed_32_rows));
@@ -194,7 +194,7 @@ mod test {
     }
 
     fn range32_chip_fail_out_of_range() {
-        const LOG_SIZE: u32 = Traces::MIN_LOG_SIZE;
+        const LOG_SIZE: u32 = PreprocessedTraces::MIN_LOG_SIZE;
         let (config, twiddles) = test_params(LOG_SIZE);
         let mut traces = Traces::new(LOG_SIZE);
         let mut side_note = SideNote::default();

prover/src/machine2/chips/range_bool.rs

@@ -73,6 +73,7 @@ mod test {
 
     use crate::machine2::components::{MachineComponent, MachineEval};
 
+    use crate::machine2::trace::PreprocessedTraces;
     use crate::machine2::traits::MachineChip;
     use crate::test_utils::{assert_chip, commit_traces, test_params, CommittedTraces};
 
@@ -106,7 +107,7 @@ mod test {
                 &mut side_note,
             );
         }
-        let preprocessed_bool_rows = Traces::empty_preprocessed_trace(LOG_SIZE);
+        let preprocessed_bool_rows = PreprocessedTraces::empty(LOG_SIZE);
         assert_chip::<RangeBoolChip>(traces, Some(preprocessed_bool_rows));
     }
 

prover/src/machine2/chips/register_mem_check.rs

@@ -32,7 +32,8 @@ use crate::machine2::{
         eval::{preprocessed_trace_eval, trace_eval, TraceEval},
         regs::AccessResult,
         sidenote::SideNote,
-        FromBaseFields, ProgramStep, Traces,
+        utils::FromBaseFields,
+        PreprocessedTraces, ProgramStep, Traces,
     },
     traits::MachineChip,
 };
@@ -144,7 +145,7 @@ impl MachineChip for RegisterMemCheckChip {
     }
     fn fill_interaction_trace(
         original_traces: &Traces,
-        preprocessed_trace: &Traces,
+        preprocessed_trace: &PreprocessedTraces,
         lookup_element: &LookupElements<MAX_LOOKUP_TUPLE_SIZE>,
     ) -> ColumnVec<CircleEvaluation<SimdBackend, BaseField, BitReversedOrder>> {
         let mut logup_trace_gen = LogupTraceGenerator::new(original_traces.log_size());
@@ -238,7 +239,7 @@ impl RegisterMemCheckChip {
     fn add_initial_reg(
         logup_trace_gen: &mut LogupTraceGenerator,
         original_traces: &Traces,
-        preprocessed_trace: &Traces,
+        preprocessed_trace: &PreprocessedTraces,
         lookup_element: &LookupElements<MAX_LOOKUP_TUPLE_SIZE>,
     ) {
         let [is_first_32] =
@@ -259,7 +260,7 @@ impl RegisterMemCheckChip {
     fn subtract_final_reg(
         logup_trace_gen: &mut LogupTraceGenerator,
         original_traces: &Traces,
-        preprocessed_trace: &Traces,
+        preprocessed_trace: &PreprocessedTraces,
         lookup_element: &LookupElements<MAX_LOOKUP_TUPLE_SIZE>,
     ) {
         let [is_first_32] =
@@ -311,7 +312,7 @@ impl RegisterMemCheckChip {
     fn add_cur_reg(
         logup_trace_gen: &mut LogupTraceGenerator,
         original_traces: &Traces,
-        preprocessed_trace: &Traces,
+        preprocessed_trace: &PreprocessedTraces,
         lookup_element: &LookupElements<MAX_LOOKUP_TUPLE_SIZE>,
         accessed: Column,
         reg_address: Column,
@@ -370,7 +371,7 @@ mod test {
     use crate::{
         machine2::{
             chips::{AddChip, CpuChip, RegisterMemCheckChip},
-            trace::{ProgramStep, Traces},
+            trace::{PreprocessedTraces, ProgramStep, Traces},
             traits::MachineChip,
         },
         test_utils::assert_chip,
@@ -470,7 +471,7 @@ mod test {
                 &mut side_note,
             );
         }
-        let mut preprocessed_column = Traces::empty_preprocessed_trace(LOG_SIZE);
+        let mut preprocessed_column = PreprocessedTraces::empty(LOG_SIZE);
         preprocessed_column.fill_is_first();
         preprocessed_column.fill_is_first32();
         preprocessed_column.fill_row_idx();

prover/src/machine2/chips/slt.rs

@@ -184,15 +184,18 @@ impl MachineChip for SltChip {
 
 #[cfg(test)]
 mod test {
-    use crate::machine2::chips::{AddChip, CpuChip, SubChip};
+    use crate::machine2::{
+        chips::{AddChip, CpuChip, SubChip},
+        trace::PreprocessedTraces,
+    };
 
     use super::*;
     use nexus_vm::{
         riscv::{BasicBlock, BuiltinOpcode, Instruction, InstructionType, Opcode},
         trace::k_trace_direct,
     };
 
-    const LOG_SIZE: u32 = Traces::MIN_LOG_SIZE;
+    const LOG_SIZE: u32 = PreprocessedTraces::MIN_LOG_SIZE;
 
     #[rustfmt::skip]
     fn setup_basic_block_ir() -> Vec<BasicBlock>

prover/src/machine2/chips/sltu.rs

@@ -130,15 +130,18 @@ impl MachineChip for SltuChip {
 
 #[cfg(test)]
 mod test {
-    use crate::machine2::chips::{AddChip, CpuChip}; // needed for ADDI to put a non-zero value in a register
+    use crate::machine2::{
+        chips::{AddChip, CpuChip}, // needed for ADDI to put a non-zero value in a register
+        trace::PreprocessedTraces,
+    };
 
     use super::*;
     use nexus_vm::{
         riscv::{BasicBlock, BuiltinOpcode, Instruction, InstructionType, Opcode},
         trace::k_trace_direct,
     };
 
-    const LOG_SIZE: u32 = Traces::MIN_LOG_SIZE;
+    const LOG_SIZE: u32 = PreprocessedTraces::MIN_LOG_SIZE;
 
     #[rustfmt::skip]
     fn setup_basic_block_ir() -> Vec<BasicBlock>

prover/src/machine2/chips/sub.rs

@@ -140,13 +140,16 @@ impl MachineChip for SubChip {
 #[cfg(test)]
 mod test {
     use super::*;
-    use crate::machine2::chips::{cpu::CpuChip, AddChip};
+    use crate::machine2::{
+        chips::{cpu::CpuChip, AddChip},
+        trace::PreprocessedTraces,
+    };
     use nexus_vm::{
         riscv::{BasicBlock, BuiltinOpcode, Instruction, InstructionType, Opcode},
         trace::k_trace_direct,
     };
 
-    const LOG_SIZE: u32 = Traces::MIN_LOG_SIZE;
+    const LOG_SIZE: u32 = PreprocessedTraces::MIN_LOG_SIZE;
 
     #[rustfmt::skip]
     fn setup_basic_block_ir() -> Vec<BasicBlock>

prover/src/machine2/mod.rs

@@ -13,7 +13,7 @@ use stwo_prover::{
 };
 
 use nexus_vm::trace::Trace;
-use trace::{sidenote::SideNote, ProgramStep};
+use trace::{sidenote::SideNote, PreprocessedTraces, ProgramStep};
 
 pub mod chips;
 pub mod components;
@@ -88,7 +88,7 @@ impl<C: MachineChip + Sync> Machine<C> {
         tree_builder.commit(prover_channel);
 
         let lookup_elements = LookupElements::draw(prover_channel);
-        let preprocessed_trace = trace::Traces::new_preprocessed_trace(LOG_SIZE);
+        let preprocessed_trace = PreprocessedTraces::new(LOG_SIZE);
         let mut tree_builder = commitment_scheme.tree_builder();
         let interaction_trace =
             C::fill_interaction_trace(&prover_traces, &preprocessed_trace, &lookup_elements);