From 48c9aac9556f7693d898c428240baf79856e8f25 Mon Sep 17 00:00:00 2001
From: Thomas Coratger <thomas.coratger@gmail.com>
Date: Tue, 26 Nov 2024 13:53:12 +0100
Subject: [PATCH 1/5] add Trace enum

---
 .../brainfuck_prover/src/brainfuck_air/mod.rs |  4 +-
 .../src/components/instruction/table.rs       |  8 +-
 .../src/components/io/table.rs                | 33 ++++++--
 .../src/components/memory/component.rs        | 44 -----------
 .../src/components/memory/mod.rs              |  1 -
 .../src/components/memory/table.rs            |  9 ++-
 crates/brainfuck_prover/src/components/mod.rs | 75 +++++++++++++++++++
 7 files changed, 115 insertions(+), 59 deletions(-)
 delete mode 100644 crates/brainfuck_prover/src/components/memory/component.rs

diff --git a/crates/brainfuck_prover/src/brainfuck_air/mod.rs b/crates/brainfuck_prover/src/brainfuck_air/mod.rs
index 9cabbe6..6594ef4 100644
--- a/crates/brainfuck_prover/src/brainfuck_air/mod.rs
+++ b/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 {
diff --git a/crates/brainfuck_prover/src/components/instruction/table.rs b/crates/brainfuck_prover/src/components/instruction/table.rs
index bef1950..e0138eb 100644
--- a/crates/brainfuck_prover/src/components/instruction/table.rs
+++ b/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);
diff --git a/crates/brainfuck_prover/src/components/io/table.rs b/crates/brainfuck_prover/src/components/io/table.rs
index 12d30b5..25b2de7 100644
--- a/crates/brainfuck_prover/src/components/io/table.rs
+++ b/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.");
diff --git a/crates/brainfuck_prover/src/components/memory/component.rs b/crates/brainfuck_prover/src/components/memory/component.rs
deleted file mode 100644
index 6c24207..0000000
--- a/crates/brainfuck_prover/src/components/memory/component.rs
+++ /dev/null
@@ -1,44 +0,0 @@
-use super::table::MemoryColumn;
-use stwo_prover::core::{channel::Channel, 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![];
-        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());
-    }
-}
diff --git a/crates/brainfuck_prover/src/components/memory/mod.rs b/crates/brainfuck_prover/src/components/memory/mod.rs
index 570dae2..13971b0 100644
--- a/crates/brainfuck_prover/src/components/memory/mod.rs
+++ b/crates/brainfuck_prover/src/components/memory/mod.rs
@@ -1,2 +1 @@
-pub mod component;
 pub mod table;
diff --git a/crates/brainfuck_prover/src/components/memory/table.rs b/crates/brainfuck_prover/src/components/memory/table.rs
index d5ae4c5..072c943 100644
--- a/crates/brainfuck_prover/src/components/memory/table.rs
+++ b/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() {
diff --git a/crates/brainfuck_prover/src/components/mod.rs b/crates/brainfuck_prover/src/components/mod.rs
index 4c66996..f853238 100644
--- a/crates/brainfuck_prover/src/components/mod.rs
+++ b/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());
+    }
+}

From 64ff66e97f79afca717b0e91d04e2f6c6b4ec244 Mon Sep 17 00:00:00 2001
From: Thomas Coratger <thomas.coratger@gmail.com>
Date: Tue, 26 Nov 2024 16:11:31 +0100
Subject: [PATCH 2/5] fix comments

---
 .../src/components/instruction/table.rs               |  9 ++++-----
 crates/brainfuck_prover/src/components/io/table.rs    | 11 +++++------
 .../brainfuck_prover/src/components/memory/table.rs   |  9 ++++-----
 crates/brainfuck_prover/src/components/mod.rs         |  6 +++---
 4 files changed, 16 insertions(+), 19 deletions(-)

diff --git a/crates/brainfuck_prover/src/components/instruction/table.rs b/crates/brainfuck_prover/src/components/instruction/table.rs
index e0138eb..91e03a7 100644
--- a/crates/brainfuck_prover/src/components/instruction/table.rs
+++ b/crates/brainfuck_prover/src/components/instruction/table.rs
@@ -1,4 +1,4 @@
-use crate::components::{Claim, Trace, TraceError, TraceEval};
+use crate::components::{Claim, TraceError, TraceEval, TraceType};
 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, trace: Trace::Instruction }))
+        Ok((trace, Claim { log_size, trace: TraceType::Instruction }))
     }
 }
 
@@ -242,9 +242,8 @@ impl InstructionColumn {
 
 #[cfg(test)]
 mod tests {
-    use crate::components::{Claim, Trace};
-
     use super::*;
+    use crate::components::{Claim, TraceType};
     use brainfuck_vm::{
         compiler::Compiler, instruction::InstructionType, test_helper::create_test_machine,
     };
@@ -585,7 +584,7 @@ mod tests {
             .collect();
 
         // Create the expected claim.
-        let expected_claim = Claim { log_size: expected_log_size, trace: Trace::Instruction };
+        let expected_claim = Claim { log_size: expected_log_size, trace: TraceType::Instruction };
 
         // Assert equality of the claim.
         assert_eq!(claim, expected_claim);
diff --git a/crates/brainfuck_prover/src/components/io/table.rs b/crates/brainfuck_prover/src/components/io/table.rs
index 25b2de7..8dddf89 100644
--- a/crates/brainfuck_prover/src/components/io/table.rs
+++ b/crates/brainfuck_prover/src/components/io/table.rs
@@ -1,4 +1,4 @@
-use crate::components::{Claim, Trace, TraceEval};
+use crate::components::{Claim, TraceEval, TraceType};
 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, trace: Trace::Io });
+            return (TraceEval::new(), Claim { log_size: 0, trace: TraceType::Io });
         }
 
         // Compute `log_n_rows`, the base-2 logarithm of the number of rows.
@@ -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: Trace::Io })
+        (trace, Claim { log_size, trace: TraceType::Io })
     }
 }
 
@@ -181,9 +181,8 @@ impl IoColumn {
 
 #[cfg(test)]
 mod tests {
-    use crate::components::Trace;
-
     use super::*;
+    use crate::components::TraceType;
     use num_traits::One;
 
     type TestIOTable = IOTable<10>;
@@ -352,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, trace: Trace::Io };
+        let expected_claim = Claim { log_size: expected_log_size, trace: TraceType::Io };
 
         // Assert equality of the claim.
         assert_eq!(claim, expected_claim, "The claim should match the expected claim.");
diff --git a/crates/brainfuck_prover/src/components/memory/table.rs b/crates/brainfuck_prover/src/components/memory/table.rs
index 072c943..f8ebf11 100644
--- a/crates/brainfuck_prover/src/components/memory/table.rs
+++ b/crates/brainfuck_prover/src/components/memory/table.rs
@@ -1,4 +1,4 @@
-use crate::components::{Claim, Trace, TraceError, TraceEval};
+use crate::components::{Claim, TraceError, TraceEval, TraceType};
 use brainfuck_vm::registers::Registers;
 use num_traits::One;
 use stwo_prover::core::{
@@ -217,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, trace: Trace::Memory }))
+        Ok((trace, Claim { log_size, trace: TraceType::Memory }))
     }
 }
 
@@ -268,9 +268,8 @@ impl MemoryColumn {
 
 #[cfg(test)]
 mod tests {
-    use crate::components::Trace;
-
     use super::*;
+    use crate::components::TraceType;
     use num_traits::Zero;
 
     #[test]
@@ -486,7 +485,7 @@ mod tests {
             .into_iter()
             .map(|col| CircleEvaluation::new(domain, col))
             .collect();
-        let expected_claim = Claim { log_size: expected_log_size, trace: Trace::Memory };
+        let expected_claim = Claim { log_size: expected_log_size, trace: TraceType::Memory };
 
         assert_eq!(claim, expected_claim);
         for col_index in 0..expected_trace.len() {
diff --git a/crates/brainfuck_prover/src/components/mod.rs b/crates/brainfuck_prover/src/components/mod.rs
index f853238..b7919c0 100644
--- a/crates/brainfuck_prover/src/components/mod.rs
+++ b/crates/brainfuck_prover/src/components/mod.rs
@@ -29,7 +29,7 @@ pub enum TraceError {
 
 /// Represents the different trace types used in the Brainfuck STARK proving system.
 #[derive(Debug, Eq, PartialEq)]
-pub enum Trace {
+pub enum TraceType {
     /// Memory access trace.
     Memory,
     /// Instruction execution trace.
@@ -40,7 +40,7 @@ pub enum Trace {
     Processor,
 }
 
-impl Trace {
+impl TraceType {
     /// Returns the number of columns associated with the specific trace type.
     pub fn column_count(&self) -> usize {
         match self {
@@ -58,7 +58,7 @@ pub struct Claim {
     /// Logarithmic size (`log2`) of the evaluated trace.
     pub log_size: u32,
     /// Type of the associated trace.
-    pub trace: Trace,
+    pub trace: TraceType,
 }
 
 impl Claim {

From 5e7f2a02655ea6b7abe697adbd3a20d9e2ccdd1d Mon Sep 17 00:00:00 2001
From: Thomas Coratger <thomas.coratger@gmail.com>
Date: Tue, 26 Nov 2024 16:38:15 +0100
Subject: [PATCH 3/5] fix comments

---
 .../brainfuck_prover/src/brainfuck_air/mod.rs |  7 ++-
 .../src/components/instruction/table.rs       | 15 +++---
 .../src/components/io/table.rs                | 16 +++---
 .../src/components/memory/table.rs            | 14 +++---
 crates/brainfuck_prover/src/components/mod.rs | 49 ++++++-------------
 5 files changed, 44 insertions(+), 57 deletions(-)

diff --git a/crates/brainfuck_prover/src/brainfuck_air/mod.rs b/crates/brainfuck_prover/src/brainfuck_air/mod.rs
index 6594ef4..0f17abb 100644
--- a/crates/brainfuck_prover/src/brainfuck_air/mod.rs
+++ b/crates/brainfuck_prover/src/brainfuck_air/mod.rs
@@ -1,4 +1,7 @@
-use crate::components::{memory::table::MemoryTable, Claim};
+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: Claim,
+    pub memory: Claim<MemoryColumn>,
 }
 
 impl BrainfuckClaim {
diff --git a/crates/brainfuck_prover/src/components/instruction/table.rs b/crates/brainfuck_prover/src/components/instruction/table.rs
index 91e03a7..22e4a44 100644
--- a/crates/brainfuck_prover/src/components/instruction/table.rs
+++ b/crates/brainfuck_prover/src/components/instruction/table.rs
@@ -1,4 +1,4 @@
-use crate::components::{Claim, TraceError, TraceEval, TraceType};
+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, trace: TraceType::Instruction }))
+        Ok((trace, Claim::<InstructionColumn>::new(log_size)))
     }
 }
 
@@ -233,9 +233,10 @@ 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
     }
 }
@@ -243,7 +244,7 @@ impl InstructionColumn {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::components::{Claim, TraceType};
+    use crate::components::Claim;
     use brainfuck_vm::{
         compiler::Compiler, instruction::InstructionType, test_helper::create_test_machine,
     };
@@ -584,7 +585,7 @@ mod tests {
             .collect();
 
         // Create the expected claim.
-        let expected_claim = Claim { log_size: expected_log_size, trace: TraceType::Instruction };
+        let expected_claim = Claim::<InstructionColumn>::new(expected_log_size);
 
         // Assert equality of the claim.
         assert_eq!(claim, expected_claim);
diff --git a/crates/brainfuck_prover/src/components/io/table.rs b/crates/brainfuck_prover/src/components/io/table.rs
index 8dddf89..0c2149f 100644
--- a/crates/brainfuck_prover/src/components/io/table.rs
+++ b/crates/brainfuck_prover/src/components/io/table.rs
@@ -1,4 +1,4 @@
-use crate::components::{Claim, TraceEval, TraceType};
+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, trace: TraceType::Io });
+            return (TraceEval::new(), Claim::<IoColumn>::new(0));
         }
 
         // Compute `log_n_rows`, the base-2 logarithm of the number of rows.
@@ -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: TraceType::Io })
+        (trace, Claim::<IoColumn>::new(log_size))
     }
 }
 
@@ -172,9 +172,10 @@ impl IoColumn {
             Self::Io => 0,
         }
     }
+}
 
-    /// Returns the total number of columns in the IO table.
-    pub const fn count() -> usize {
+impl TraceColumn for IoColumn {
+    fn count() -> usize {
         1
     }
 }
@@ -182,7 +183,6 @@ impl IoColumn {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::components::TraceType;
     use num_traits::One;
 
     type TestIOTable = IOTable<10>;
@@ -351,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, trace: TraceType::Io };
+        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.");
diff --git a/crates/brainfuck_prover/src/components/memory/table.rs b/crates/brainfuck_prover/src/components/memory/table.rs
index f8ebf11..60c5857 100644
--- a/crates/brainfuck_prover/src/components/memory/table.rs
+++ b/crates/brainfuck_prover/src/components/memory/table.rs
@@ -1,4 +1,4 @@
-use crate::components::{Claim, TraceError, TraceEval, TraceType};
+use crate::components::{Claim, TraceColumn, TraceError, TraceEval};
 use brainfuck_vm::registers::Registers;
 use num_traits::One;
 use stwo_prover::core::{
@@ -195,7 +195,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);
@@ -217,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, trace: TraceType::Memory }))
+        Ok((trace, Claim::<MemoryColumn>::new(log_size)))
     }
 }
 
@@ -259,9 +259,10 @@ impl MemoryColumn {
             Self::D => 3,
         }
     }
+}
 
-    /// Returns the total number of columns in the Memory table
-    pub const fn count() -> usize {
+impl TraceColumn for MemoryColumn {
+    fn count() -> usize {
         4
     }
 }
@@ -269,7 +270,6 @@ impl MemoryColumn {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::components::TraceType;
     use num_traits::Zero;
 
     #[test]
@@ -485,7 +485,7 @@ mod tests {
             .into_iter()
             .map(|col| CircleEvaluation::new(domain, col))
             .collect();
-        let expected_claim = Claim { log_size: expected_log_size, trace: TraceType::Memory };
+        let expected_claim = Claim::<MemoryColumn>::new(expected_log_size);
 
         assert_eq!(claim, expected_claim);
         for col_index in 0..expected_trace.len() {
diff --git a/crates/brainfuck_prover/src/components/mod.rs b/crates/brainfuck_prover/src/components/mod.rs
index b7919c0..ca3c260 100644
--- a/crates/brainfuck_prover/src/components/mod.rs
+++ b/crates/brainfuck_prover/src/components/mod.rs
@@ -1,6 +1,3 @@
-use instruction::table::InstructionColumn;
-use io::table::IoColumn;
-use memory::table::MemoryColumn;
 use stwo_prover::core::{
     backend::simd::SimdBackend,
     channel::Channel,
@@ -27,41 +24,21 @@ pub enum TraceError {
     EmptyTrace,
 }
 
-/// Represents the different trace types used in the Brainfuck STARK proving system.
-#[derive(Debug, Eq, PartialEq)]
-pub enum TraceType {
-    /// Memory access trace.
-    Memory,
-    /// Instruction execution trace.
-    Instruction,
-    /// Input/output trace.
-    Io,
-    /// Processor trace (register).
-    Processor,
-}
-
-impl TraceType {
-    /// 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 {
+pub struct Claim<T: TraceColumn> {
     /// Logarithmic size (`log2`) of the evaluated trace.
     pub log_size: u32,
-    /// Type of the associated trace.
-    pub trace: TraceType,
+    /// Marker for the trace type.
+    pub _marker: std::marker::PhantomData<T>,
 }
 
-impl Claim {
+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,
@@ -81,7 +58,7 @@ impl Claim {
     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 trace_log_sizes = vec![self.log_size; T::count()];
         let interaction_trace_log_sizes: Vec<u32> = vec![];
         TreeVec::new(vec![
             preprocessed_trace_log_sizes,
@@ -96,3 +73,9 @@ impl Claim {
         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;
+}

From c3719d2040e12ed1753163bd9314cb84166fc343 Mon Sep 17 00:00:00 2001
From: Thomas Coratger <thomas.coratger@gmail.com>
Date: Tue, 26 Nov 2024 18:19:01 +0100
Subject: [PATCH 4/5] fix marker visibility

---
 crates/brainfuck_prover/src/components/mod.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/crates/brainfuck_prover/src/components/mod.rs b/crates/brainfuck_prover/src/components/mod.rs
index ca3c260..cc44bc7 100644
--- a/crates/brainfuck_prover/src/components/mod.rs
+++ b/crates/brainfuck_prover/src/components/mod.rs
@@ -30,7 +30,7 @@ pub struct Claim<T: TraceColumn> {
     /// Logarithmic size (`log2`) of the evaluated trace.
     pub log_size: u32,
     /// Marker for the trace type.
-    pub _marker: std::marker::PhantomData<T>,
+    _marker: std::marker::PhantomData<T>,
 }
 
 impl<T: TraceColumn> Claim<T> {

From 92892c2d05e826b99202dc21bfa7d3465dfa35fc Mon Sep 17 00:00:00 2001
From: Thomas Coratger <thomas.coratger@gmail.com>
Date: Tue, 26 Nov 2024 18:23:57 +0100
Subject: [PATCH 5/5] fix

---
 .../brainfuck_prover/src/components/memory/component.rs  | 7 +------
 crates/brainfuck_prover/src/components/memory/table.rs   | 9 +++++++--
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/crates/brainfuck_prover/src/components/memory/component.rs b/crates/brainfuck_prover/src/components/memory/component.rs
index 1fd52ca..8df7cbc 100644
--- a/crates/brainfuck_prover/src/components/memory/component.rs
+++ b/crates/brainfuck_prover/src/components/memory/component.rs
@@ -1,9 +1,4 @@
-use super::table::MemoryColumn;
-use stwo_prover::core::{
-    channel::Channel,
-    fields::{qm31::SecureField, secure_column::SECURE_EXTENSION_DEGREE},
-    pcs::TreeVec,
-};
+use stwo_prover::core::{channel::Channel, fields::qm31::SecureField};
 
 /// The claim of the interaction phase 2 (with the logUp protocol).
 ///
diff --git a/crates/brainfuck_prover/src/components/memory/table.rs b/crates/brainfuck_prover/src/components/memory/table.rs
index 3912ddf..90ee38d 100644
--- a/crates/brainfuck_prover/src/components/memory/table.rs
+++ b/crates/brainfuck_prover/src/components/memory/table.rs
@@ -267,11 +267,16 @@ impl MemoryColumn {
             Self::D => 3,
         }
     }
+
+    /// Returns the total number of columns in the Memory table
+    pub const fn column_count() -> usize {
+        4
+    }
 }
 
 impl TraceColumn for MemoryColumn {
     fn count() -> usize {
-        4
+        Self::column_count()
     }
 }
 
@@ -284,7 +289,7 @@ impl TraceColumn for 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.