add Trace enum

crates/brainfuck_prover/src/brainfuck_air/mod.rs

@@ -1,4 +1,4 @@
-use crate::components::memory::{self, table::MemoryTable};
+use crate::components::{memory::table::MemoryTable, Claim};
 use brainfuck_vm::machine::Machine;
 use stwo_prover::core::{
     air::{Component, ComponentProver},
@@ -27,7 +27,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,
 }
 
 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, Trace, TraceError, TraceEval};
 use brainfuck_vm::{
     instruction::VALID_INSTRUCTIONS_BF, machine::ProgramMemory, registers::Registers,
 };
@@ -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 { log_size, trace: Trace::Instruction }))
     }
 }
 
@@ -242,6 +242,8 @@ impl InstructionColumn {
 
 #[cfg(test)]
 mod tests {
+    use crate::components::{Claim, Trace};
+
     use super::*;
     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 { log_size: expected_log_size, trace: Trace::Instruction };
 
         // 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, Trace, TraceEval};
 use brainfuck_vm::{instruction::InstructionType, registers::Registers};
 use stwo_prover::core::{
     backend::{
@@ -101,7 +101,7 @@ impl<const N: u32> IOTable<N> {
 
         // 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 { log_size: 0, trace: Trace::Io });
         }
 
         // 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 { log_size, trace: Trace::Io })
     }
 }
 
@@ -158,8 +158,31 @@ 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,
+        }
+    }
+
+    /// Returns the total number of columns in the IO table.
+    pub const fn count() -> usize {
+        1
+    }
+}
+
 #[cfg(test)]
 mod tests {
+    use crate::components::Trace;
+
     use super::*;
     use num_traits::One;
 
@@ -329,7 +352,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 { log_size: expected_log_size, trace: Trace::Io };
 
         // Assert equality of the claim.
         assert_eq!(claim, expected_claim, "The claim should match the expected claim.");

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

@@ -1,2 +1 @@
-pub mod component;
 pub mod table;

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

@@ -1,5 +1,4 @@
-use super::component::Claim;
-use crate::components::{TraceError, TraceEval};
+use crate::components::{Claim, Trace, TraceError, TraceEval};
 use brainfuck_vm::registers::Registers;
 use num_traits::One;
 use stwo_prover::core::{
@@ -218,7 +217,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 { log_size, trace: Trace::Memory }))
     }
 }
 
@@ -269,6 +268,8 @@ impl MemoryColumn {
 
 #[cfg(test)]
 mod tests {
+    use crate::components::Trace;
+
     use super::*;
     use num_traits::Zero;
 
@@ -485,7 +486,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 { log_size: expected_log_size, trace: Trace::Memory };
 
         assert_eq!(claim, expected_claim);
         for col_index in 0..expected_trace.len() {

crates/brainfuck_prover/src/components/mod.rs

@@ -1,6 +1,11 @@
+use instruction::table::InstructionColumn;
+use io::table::IoColumn;
+use memory::table::MemoryColumn;
 use stwo_prover::core::{
     backend::simd::SimdBackend,
+    channel::Channel,
     fields::m31::BaseField,
+    pcs::TreeVec,
     poly::{circle::CircleEvaluation, BitReversedOrder},
     ColumnVec,
 };
@@ -21,3 +26,73 @@ pub enum TraceError {
     #[error("The trace is empty.")]
     EmptyTrace,
 }
+
+/// Represents the different trace types used in the Brainfuck STARK proving system.
+#[derive(Debug, Eq, PartialEq)]
+pub enum Trace {
+    /// Memory access trace.
+    Memory,
+    /// Instruction execution trace.
+    Instruction,
+    /// Input/output trace.
+    Io,
+    /// Processor trace (register).
+    Processor,
+}
+
+impl Trace {
+    /// Returns the number of columns associated with the specific trace type.
+    pub fn column_count(&self) -> usize {
+        match self {
+            Self::Memory => MemoryColumn::count(),
+            Self::Instruction => InstructionColumn::count(),
+            Self::Io => IoColumn::count(),
+            Self::Processor => unimplemented!(),
+        }
+    }
+}
+
+/// Represents a claim associated with a specific trace in the Brainfuck STARK proving system.
+#[derive(Debug, Eq, PartialEq)]
+pub struct Claim {
+    /// Logarithmic size (`log2`) of the evaluated trace.
+    pub log_size: u32,
+    /// Type of the associated trace.
+    pub trace: Trace,
+}
+
+impl Claim {
+    /// 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; self.trace.column_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());
+    }
+}