abstract Claim for multiple traces (#93)

* add Trace enum

* fix comments

* fix comments

* fix marker visibility

* fix

crates/brainfuck_prover/src/brainfuck_air/mod.rs

@@ -1,4 +1,7 @@
-use crate::components::memory::{self, table::MemoryTable};
+use crate::components::{
+    memory::table::{MemoryColumn, MemoryTable},
+    Claim,
+};
 use brainfuck_vm::machine::Machine;
 use stwo_prover::core::{
     air::{Component, ComponentProver},
@@ -27,7 +30,7 @@ pub struct BrainfuckProof<H: MerkleHasher> {
 /// It includes the common claim values such as the initial and final states
 /// and the claim of each component.
 pub struct BrainfuckClaim {
-    pub memory: memory::component::Claim,
+    pub memory: Claim<MemoryColumn>,
 }
 
 impl BrainfuckClaim {

crates/brainfuck_prover/src/components/instruction/table.rs

@@ -1,4 +1,4 @@
-use crate::components::{memory::component::Claim, TraceError, TraceEval};
+use crate::components::{Claim, TraceColumn, TraceError, TraceEval};
 use brainfuck_vm::{
     instruction::VALID_INSTRUCTIONS_BF, machine::ProgramMemory, registers::Registers,
 };
@@ -128,7 +128,7 @@ impl InstructionTable {
     ///
     /// # Errors
     /// Returns [`TraceError::EmptyTrace`] if the table is empty.
-    pub fn trace_evaluation(&self) -> Result<(TraceEval, Claim), TraceError> {
+    pub fn trace_evaluation(&self) -> Result<(TraceEval, Claim<InstructionColumn>), TraceError> {
         let n_rows = self.table.len() as u32;
         // If the table is empty, there is no data to evaluate, so return an error.
         if n_rows == 0 {
@@ -171,7 +171,7 @@ impl InstructionTable {
         let trace = trace.into_iter().map(|col| CircleEvaluation::new(domain, col)).collect();
 
         // Return the evaluated trace and a claim containing the log size of the domain.
-        Ok((trace, Claim { log_size }))
+        Ok((trace, Claim::<InstructionColumn>::new(log_size)))
     }
 }
 
@@ -233,16 +233,18 @@ impl InstructionColumn {
             Self::Ni => 2,
         }
     }
+}
 
-    /// Returns the total number of columns in the Instruction trace
-    pub const fn count() -> usize {
+impl TraceColumn for InstructionColumn {
+    fn count() -> usize {
         3
     }
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::components::Claim;
     use brainfuck_vm::{
         compiler::Compiler, instruction::InstructionType, test_helper::create_test_machine,
     };
@@ -583,7 +585,7 @@ mod tests {
             .collect();
 
         // Create the expected claim.
-        let expected_claim = Claim { log_size: expected_log_size };
+        let expected_claim = Claim::<InstructionColumn>::new(expected_log_size);
 
         // Assert equality of the claim.
         assert_eq!(claim, expected_claim);

crates/brainfuck_prover/src/components/io/table.rs

@@ -1,4 +1,4 @@
-use crate::components::{memory::component::Claim, TraceEval};
+use crate::components::{Claim, TraceColumn, TraceEval};
 use brainfuck_vm::{instruction::InstructionType, registers::Registers};
 use stwo_prover::core::{
     backend::{
@@ -96,12 +96,12 @@ impl<const N: u32> IOTable<N> {
     /// # Returns
     /// A tuple containing the evaluated trace and claim for STARK proof.
     /// If the table is empty, returns an empty trace and a claim with a log size of 0.
-    pub fn trace_evaluation(&self) -> (TraceEval, Claim) {
+    pub fn trace_evaluation(&self) -> (TraceEval, Claim<IoColumn>) {
         let n_rows = self.table.len() as u32;
 
         // It is possible that the table is empty because the program has no input or output.
         if n_rows == 0 {
-            return (TraceEval::new(), Claim { log_size: 0 });
+            return (TraceEval::new(), Claim::<IoColumn>::new(0));
         }
 
         // Compute `log_n_rows`, the base-2 logarithm of the number of rows.
@@ -116,7 +116,7 @@ impl<const N: u32> IOTable<N> {
 
         // Populate the column with data from the table rows.
         for (index, row) in self.table.iter().enumerate().take(1 << log_n_rows) {
-            trace[0].data[index] = row.mv.into();
+            trace[IoColumn::Io.index()].data[index] = row.mv.into();
         }
 
         // Create a circle domain using a canonical coset.
@@ -126,7 +126,7 @@ impl<const N: u32> IOTable<N> {
         let trace = trace.into_iter().map(|col| CircleEvaluation::new(domain, col)).collect();
 
         // Return the evaluated trace and a claim containing the log size of the domain.
-        (trace, Claim { log_size })
+        (trace, Claim::<IoColumn>::new(log_size))
     }
 }
 
@@ -158,6 +158,28 @@ pub type InputTable = IOTable<{ InstructionType::ReadChar.to_u32() }>;
 /// outputs (when the current instruction `ci` equals '.').
 pub type OutputTable = IOTable<{ InstructionType::PutChar.to_u32() }>;
 
+/// Enum representing the column indices in the IO trace.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum IoColumn {
+    /// Column representing the input/output operations.
+    Io,
+}
+
+impl IoColumn {
+    /// Returns the index of the column in the IO table.
+    pub const fn index(self) -> usize {
+        match self {
+            Self::Io => 0,
+        }
+    }
+}
+
+impl TraceColumn for IoColumn {
+    fn count() -> usize {
+        1
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -329,7 +351,7 @@ mod tests {
             expected_columns.into_iter().map(|col| CircleEvaluation::new(domain, col)).collect();
 
         // Create the expected claim.
-        let expected_claim = Claim { log_size: expected_log_size };
+        let expected_claim = Claim::<IoColumn>::new(expected_log_size);
 
         // Assert equality of the claim.
         assert_eq!(claim, expected_claim, "The claim should match the expected claim.");

crates/brainfuck_prover/src/components/memory/component.rs

@@ -1,51 +1,4 @@
-use super::table::MemoryColumn;
-use stwo_prover::core::{
-    channel::Channel,
-    fields::{qm31::SecureField, secure_column::SECURE_EXTENSION_DEGREE},
-    pcs::TreeVec,
-};
-
-/// The claim for the Memory component
-#[derive(Debug, Eq, PartialEq)]
-pub struct Claim {
-    pub log_size: u32,
-}
-
-impl Claim {
-    /// Returns the `log_size` of the each type of trace commited for the Memory component:
-    /// - Preprocessed trace,
-    /// - Main trace,
-    /// - Interaction trace.
-    ///
-    /// The number of columns of each trace is known before actually evaluating them.
-    /// The `log_size` is known once the main trace has been evaluated
-    /// (the log2 of the size of the [`super::table::MemoryTable`], to which we add
-    /// [`stwo_prover::core::backend::simd::m31::LOG_N_LANES`]
-    /// for the [`stwo_prover::core::backend::simd::SimdBackend`])
-    ///
-    /// Each element of the [`TreeVec`] is dedicated to the commitment of one type of trace.
-    /// First element is for the preprocessed trace, second for the main trace and third for the
-    /// interaction one.
-    ///
-    /// NOTE: Currently only the main trace is provided.
-    pub fn log_sizes(&self) -> TreeVec<Vec<u32>> {
-        // TODO: Add the preprocessed and interaction trace correct sizes
-        let preprocessed_trace_log_sizes: Vec<u32> = vec![];
-        let trace_log_sizes = vec![self.log_size; MemoryColumn::count()];
-        let interaction_trace_log_sizes: Vec<u32> = vec![self.log_size; SECURE_EXTENSION_DEGREE];
-        TreeVec::new(vec![
-            preprocessed_trace_log_sizes,
-            trace_log_sizes,
-            interaction_trace_log_sizes,
-        ])
-    }
-
-    /// Mix the log size of the Memory table to the Fiat-Shamir [`Channel`],
-    /// to bound the channel randomness and the trace.
-    pub fn mix_into(&self, channel: &mut impl Channel) {
-        channel.mix_u64(self.log_size.into());
-    }
-}
+use stwo_prover::core::{channel::Channel, fields::qm31::SecureField};
 
 /// The claim of the interaction phase 2 (with the logUp protocol).
 ///

crates/brainfuck_prover/src/components/memory/table.rs

@@ -1,5 +1,5 @@
-use super::component::{Claim, InteractionClaim};
-use crate::components::{TraceError, TraceEval};
+use super::component::InteractionClaim;
+use crate::components::{Claim, TraceColumn, TraceError, TraceEval};
 use brainfuck_vm::registers::Registers;
 use num_traits::One;
 use stwo_prover::{
@@ -203,7 +203,7 @@ impl MemoryTable {
     ///
     /// # Errors
     /// Returns [`TraceError::EmptyTrace`] if the table is empty.
-    pub fn trace_evaluation(&self) -> Result<(TraceEval, Claim), TraceError> {
+    pub fn trace_evaluation(&self) -> Result<(TraceEval, Claim<MemoryColumn>), TraceError> {
         let n_rows = self.table.len() as u32;
         if n_rows == 0 {
             return Err(TraceError::EmptyTrace);
@@ -225,7 +225,7 @@ impl MemoryTable {
         let trace = trace.into_iter().map(|col| CircleEvaluation::new(domain, col)).collect();
 
         // TODO: Confirm that the log_size in `Claim` is `log_size`, including the SIMD lanes
-        Ok((trace, Claim { log_size }))
+        Ok((trace, Claim::<MemoryColumn>::new(log_size)))
     }
 }
 
@@ -269,11 +269,17 @@ impl MemoryColumn {
     }
 
     /// Returns the total number of columns in the Memory table
-    pub const fn count() -> usize {
+    pub const fn column_count() -> usize {
         4
     }
 }
 
+impl TraceColumn for MemoryColumn {
+    fn count() -> usize {
+        Self::column_count()
+    }
+}
+
 /// The interaction elements drawn for the extension column of the Memory component.
 ///
 /// The logUp protocol uses these elements to combine the values of the different
@@ -283,7 +289,7 @@ impl MemoryColumn {
 /// There are 3 lookup elements in the Memory component, as only the 'real' registers
 /// are used: `clk`, `mp` and `mv`. `d` is used to eventually nullify the numerator.
 #[derive(Clone, Debug, Eq, PartialEq)]
-pub struct MemoryElements(LookupElements<{ MemoryColumn::count() - 1 }>);
+pub struct MemoryElements(LookupElements<{ MemoryColumn::column_count() - 1 }>);
 
 impl MemoryElements {
     /// Provides dummy lookup elements.
@@ -590,7 +596,7 @@ mod tests {
             .into_iter()
             .map(|col| CircleEvaluation::new(domain, col))
             .collect();
-        let expected_claim = Claim { log_size: expected_log_size };
+        let expected_claim = Claim::<MemoryColumn>::new(expected_log_size);
 
         assert_eq!(claim, expected_claim);
         for col_index in 0..expected_trace.len() {

crates/brainfuck_prover/src/components/mod.rs

@@ -1,6 +1,8 @@
 use stwo_prover::core::{
     backend::simd::SimdBackend,
+    channel::Channel,
     fields::m31::BaseField,
+    pcs::TreeVec,
     poly::{circle::CircleEvaluation, BitReversedOrder},
     ColumnVec,
 };
@@ -21,3 +23,59 @@ pub enum TraceError {
     #[error("The trace is empty.")]
     EmptyTrace,
 }
+
+/// Represents a claim associated with a specific trace in the Brainfuck STARK proving system.
+#[derive(Debug, Eq, PartialEq)]
+pub struct Claim<T: TraceColumn> {
+    /// Logarithmic size (`log2`) of the evaluated trace.
+    pub log_size: u32,
+    /// Marker for the trace type.
+    _marker: std::marker::PhantomData<T>,
+}
+
+impl<T: TraceColumn> Claim<T> {
+    /// Creates a new claim for the given trace type.
+    pub const fn new(log_size: u32) -> Self {
+        Self { log_size, _marker: std::marker::PhantomData }
+    }
+
+    /// Returns the `log_size` for each type of trace committed for the given trace type:
+    /// - Preprocessed trace,
+    /// - Main trace,
+    /// - Interaction trace.
+    ///
+    /// The number of columns of each trace is known before actually evaluating them.
+    /// The `log_size` is known once the main trace has been evaluated
+    /// (the log2 of the size of the [`super::table::MemoryTable`], to which we add
+    /// [`stwo_prover::core::backend::simd::m31::LOG_N_LANES`]
+    /// for the [`stwo_prover::core::backend::simd::SimdBackend`])
+    ///
+    /// Each element of the [`TreeVec`] is dedicated to the commitment of one type of trace.
+    /// First element is for the preprocessed trace, second for the main trace and third for the
+    /// interaction one.
+    ///
+    /// NOTE: Currently only the main trace is provided.
+    pub fn log_sizes(&self) -> TreeVec<Vec<u32>> {
+        // TODO: Add the preprocessed and interaction trace correct sizes
+        let preprocessed_trace_log_sizes: Vec<u32> = vec![];
+        let trace_log_sizes = vec![self.log_size; T::count()];
+        let interaction_trace_log_sizes: Vec<u32> = vec![];
+        TreeVec::new(vec![
+            preprocessed_trace_log_sizes,
+            trace_log_sizes,
+            interaction_trace_log_sizes,
+        ])
+    }
+
+    /// Mix the log size of the table to the Fiat-Shamir [`Channel`],
+    /// to bound the channel randomness and the trace.
+    pub fn mix_into(&self, channel: &mut impl Channel) {
+        channel.mix_u64(self.log_size.into());
+    }
+}
+
+/// Represents columns of a trace.
+pub trait TraceColumn {
+    /// Returns the number of columns associated with the specific trace type.
+    fn count() -> usize;
+}