feat: add Program component constraints (#172)

* fix: remove ProgramElements to use InstructionElements instead

* feat: add Program constraints

* feat: add Program component to entrypoints

* chore: cleanup program table

* chore: cleanup program component

crates/brainfuck_prover/src/brainfuck_air/mod.rs

@@ -14,7 +14,12 @@ use crate::components::{
         component::{MemoryComponent, MemoryEval},
         table::{MemoryElements, MemoryTable},
     },
-    InstructionClaim, InteractionClaim, IoClaim, MemoryClaim,
+    program::{
+        self,
+        component::{ProgramComponent, ProgramEval},
+        table::ProgramTable,
+    },
+    InstructionClaim, InteractionClaim, IoClaim, MemoryClaim, ProgramClaim,
 };
 use brainfuck_vm::machine::Machine;
 use stwo_prover::{
@@ -54,6 +59,7 @@ pub struct BrainfuckClaim {
     pub output: IoClaim,
     pub memory: MemoryClaim,
     pub instruction: InstructionClaim,
+    pub program: ProgramClaim,
 }
 
 impl BrainfuckClaim {
@@ -62,6 +68,7 @@ impl BrainfuckClaim {
         self.output.mix_into(channel);
         self.memory.mix_into(channel);
         self.instruction.mix_into(channel);
+        self.program.mix_into(channel);
     }
 
     pub fn log_sizes(&self) -> TreeVec<Vec<u32>> {
@@ -71,6 +78,7 @@ impl BrainfuckClaim {
                 self.output.log_sizes(),
                 self.memory.log_sizes(),
                 self.instruction.log_sizes(),
+                self.program.log_sizes(),
             ]
             .into_iter(),
         );
@@ -113,6 +121,7 @@ pub struct BrainfuckInteractionClaim {
     output: InteractionClaim,
     memory: InteractionClaim,
     instruction: InteractionClaim,
+    program: InteractionClaim,
 }
 
 impl BrainfuckInteractionClaim {
@@ -122,6 +131,7 @@ impl BrainfuckInteractionClaim {
         self.output.mix_into(channel);
         self.memory.mix_into(channel);
         self.instruction.mix_into(channel);
+        self.program.mix_into(channel);
     }
 }
 
@@ -143,6 +153,7 @@ pub struct BrainfuckComponents {
     output: OutputComponent,
     memory: MemoryComponent,
     instruction: InstructionComponent,
+    program: ProgramComponent,
 }
 
 impl BrainfuckComponents {
@@ -165,16 +176,19 @@ impl BrainfuckComponents {
             InputEval::new(&claim.input, interaction_elements.input_lookup_elements.clone()),
             (interaction_claim.input.claimed_sum, None),
         );
+
         let output = OutputComponent::new(
             tree_span_provider,
             OutputEval::new(&claim.output, interaction_elements.output_lookup_elements.clone()),
             (interaction_claim.output.claimed_sum, None),
         );
+
         let memory = MemoryComponent::new(
             tree_span_provider,
             MemoryEval::new(&claim.memory, interaction_elements.memory_lookup_elements.clone()),
             (interaction_claim.memory.claimed_sum, None),
         );
+
         let instruction = InstructionComponent::new(
             tree_span_provider,
             InstructionEval::new(
@@ -184,12 +198,21 @@ impl BrainfuckComponents {
             (interaction_claim.instruction.claimed_sum, None),
         );
 
-        Self { input, output, memory, instruction }
+        let program = ProgramComponent::new(
+            tree_span_provider,
+            ProgramEval::new(
+                &claim.program,
+                interaction_elements.instruction_lookup_elements.clone(),
+            ),
+            (interaction_claim.program.claimed_sum, None),
+        );
+
+        Self { input, output, memory, instruction, program }
     }
 
     /// Returns the `ComponentProver` of each components, used by the prover.
     pub fn provers(&self) -> Vec<&dyn ComponentProver<SimdBackend>> {
-        vec![&self.input, &self.output, &self.memory, &self.instruction]
+        vec![&self.input, &self.output, &self.memory, &self.instruction, &self.program]
     }
 
     /// Returns the `Component` of each components, used by the verifier.
@@ -217,7 +240,7 @@ const LOG_MAX_ROWS: u32 = 20;
 ///
 /// Ideally, we should cover all possible log sizes, between
 /// 1 and `LOG_MAX_ROW`
-const IS_FIRST_LOG_SIZES: [u32; 9] = [15, 13, 10, 9, 8, 7, 6, 5, 4];
+const IS_FIRST_LOG_SIZES: [u32; 10] = [15, 13, 11, 10, 9, 8, 7, 6, 5, 4];
 
 /// Generate a STARK proof of the given Brainfuck program execution.
 ///
@@ -265,17 +288,21 @@ pub fn prove_brainfuck(
     let (memory_trace, memory_claim) = MemoryTable::from(&vm_trace).trace_evaluation().unwrap();
     let (instruction_trace, instruction_claim) =
         InstructionTable::from((&vm_trace, inputs.program())).trace_evaluation().unwrap();
+    let (program_trace, program_claim) =
+        ProgramTable::from(inputs.program()).trace_evaluation().unwrap();
 
     tree_builder.extend_evals(input_trace.clone());
     tree_builder.extend_evals(output_trace.clone());
     tree_builder.extend_evals(memory_trace.clone());
     tree_builder.extend_evals(instruction_trace.clone());
+    tree_builder.extend_evals(program_trace.clone());
 
     let claim = BrainfuckClaim {
         input: input_claim,
         output: output_claim,
         memory: memory_claim,
         instruction: instruction_claim,
+        program: program_claim,
     };
 
     // Mix the claim into the Fiat-Shamir channel.
@@ -321,16 +348,25 @@ pub fn prove_brainfuck(
         )
         .unwrap();
 
+    let (program_interaction_trace_eval, program_interaction_claim) =
+        program::table::interaction_trace_evaluation(
+            &program_trace,
+            &interaction_elements.instruction_lookup_elements,
+        )
+        .unwrap();
+
     tree_builder.extend_evals(input_interaction_trace_eval);
     tree_builder.extend_evals(output_interaction_trace_eval);
     tree_builder.extend_evals(memory_interaction_trace_eval);
     tree_builder.extend_evals(instruction_interaction_trace_eval);
+    tree_builder.extend_evals(program_interaction_trace_eval);
 
     let interaction_claim = BrainfuckInteractionClaim {
         input: input_interaction_claim,
         output: output_interaction_claim,
         memory: memory_interaction_claim,
         instruction: instruction_interaction_claim,
+        program: program_interaction_claim,
     };
 
     // Mix the interaction claim into the Fiat-Shamir channel.

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

@@ -1,23 +1,35 @@
-use super::table::ProgramElements;
-use crate::components::ProgramClaim;
-use stwo_prover::constraint_framework::{EvalAtRow, FrameworkComponent, FrameworkEval};
+use crate::components::{instruction::table::InstructionElements, ProgramClaim};
+use num_traits::One;
+use stwo_prover::{
+    constraint_framework::{
+        preprocessed_columns::PreprocessedColumn, EvalAtRow, FrameworkComponent, FrameworkEval,
+        RelationEntry,
+    },
+    core::fields::m31::BaseField,
+};
 
-/// Implementation of `Component` and `ComponentProver`
-/// for the `SimdBackend` from the constraint framework provided by Stwo
+/// Implementation of `Component` and `ComponentProver` for the [`ProgramComponent`].
+/// It targets the `SimdBackend` from the Stwo constraint framework, with a fallback
+/// on `CpuBakend` for small traces.
 pub type ProgramComponent = FrameworkComponent<ProgramEval>;
-
-/// The AIR for the Program component.
+/// The AIR for the [`ProgramComponent`].
 ///
-/// Constraints are defined through the `FrameworkEval`
+/// Constraints are defined through the [`FrameworkEval`]
 /// provided by the constraint framework of Stwo.
 pub struct ProgramEval {
+    /// The log size of the component's main trace height.
     log_size: u32,
-    _program_lookup_elements: ProgramElements,
+    /// The random elements used for the lookup protocol linking the instruction and
+    /// program components to the processor one.
+    instruction_lookup_elements: InstructionElements,
 }
 
 impl ProgramEval {
-    pub const fn new(claim: &ProgramClaim, program_lookup_elements: ProgramElements) -> Self {
-        Self { log_size: claim.log_size, _program_lookup_elements: program_lookup_elements }
+    pub const fn new(
+        claim: &ProgramClaim,
+        instruction_lookup_elements: InstructionElements,
+    ) -> Self {
+        Self { log_size: claim.log_size, instruction_lookup_elements }
     }
 }
 
@@ -26,17 +38,15 @@ impl FrameworkEval for ProgramEval {
     fn log_size(&self) -> u32 {
         self.log_size
     }
-
     /// The degree of the constraints is bounded by the size of the trace.
     ///
     /// Returns the ilog2 (upper) bound of the constraint degree for the component.
     fn max_constraint_log_degree_bound(&self) -> u32 {
         self.log_size + 1
     }
-
     /// Defines the AIR for the Program component.
     ///
-    /// Registers values from the current row are obtained through masks.
+    /// Registers' values from the current row are obtained through masks.
     /// When you apply a mask, you target the current column and then pass to the next
     /// one: the register order matters to correctly fetch them.
     ///
@@ -47,7 +57,118 @@ impl FrameworkEval for ProgramEval {
     /// Use `eval.add_to_relation` to define a global constraint for the logUp protocol.
     ///
     /// The logUp must be finalized with `eval.finalize_logup()`.
-    fn evaluate<E: EvalAtRow>(&self, _eval: E) -> E {
-        todo!()
+    fn evaluate<E: EvalAtRow>(&self, mut eval: E) -> E {
+        // Get the preprocessed column to check boundary constraints
+        let is_first = eval.get_preprocessed_column(PreprocessedColumn::IsFirst(self.log_size()));
+
+        // Get all the registers' columns
+        let ip = eval.next_trace_mask();
+        let ci = eval.next_trace_mask();
+        let ni = eval.next_trace_mask();
+        let d = eval.next_trace_mask();
+
+        // ┌──────────────────────────┐
+        // │   Boundary Constraints   │
+        // └──────────────────────────┘
+
+        // `ip` starts at 0.
+        eval.add_constraint(is_first * ip.clone());
+
+        // ┌─────────────────────────────┐
+        // │   Consistency Constraints   │
+        // └─────────────────────────────┘
+
+        // The dummy flag `d` is either 0 or 1.
+        eval.add_constraint(d.clone() * (d.clone() - BaseField::one().into()));
+
+        // If `d` is set, then `ci` equals 0
+        eval.add_constraint(d.clone() * ci.clone());
+
+        // If `d` is set, then `ni` equals 0
+        eval.add_constraint(d.clone() * ni.clone());
+
+        // ┌─────────────────────────────┐
+        // │   Interaction Constraints   │
+        // └─────────────────────────────┘
+
+        let num = E::F::one() - d;
+        eval.add_to_relation(&[RelationEntry::new(
+            &self.instruction_lookup_elements,
+            num.into(),
+            &[ip, ci, ni],
+        )]);
+
+        eval.finalize_logup();
+
+        eval
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use brainfuck_vm::{compiler::Compiler, test_helper::create_test_machine};
+    use stwo_prover::{
+        constraint_framework::{
+            assert_constraints, preprocessed_columns::gen_is_first, FrameworkEval,
+        },
+        core::{
+            pcs::TreeVec,
+            poly::circle::{CanonicCoset, CircleEvaluation},
+        },
+    };
+
+    use crate::components::{
+        instruction::table::InstructionElements,
+        program::{
+            component::ProgramEval,
+            table::{interaction_trace_evaluation, ProgramTable},
+        },
+    };
+
+    #[test]
+    fn test_program_constraints() {
+        const LOG_SIZE: u32 = 8;
+
+        // Get an execution trace from a valid Brainfuck program
+        let code = "+>,<[>+.<-]";
+        let mut compiler = Compiler::new(code);
+        let instructions = compiler.compile();
+        let (mut machine, _) = create_test_machine(&instructions, &[1]);
+        let () = machine.execute().expect("Failed to execute machine");
+
+        // Construct the IsFirst preprocessed column
+        let is_first_col = gen_is_first(LOG_SIZE);
+        let is_first_col_2 = gen_is_first(LOG_SIZE);
+        let preprocessed_trace = vec![is_first_col, is_first_col_2];
+
+        // Construct the main trace from the execution trace
+        let table = ProgramTable::from(machine.program());
+        let (instruction_trace, claim) = table.trace_evaluation().unwrap();
+
+        // Draw Interaction elements
+        let instruction_lookup_elements = InstructionElements::dummy();
+
+        // Generate interaction trace
+        let (interaction_trace, interaction_claim) =
+            interaction_trace_evaluation(&instruction_trace, &instruction_lookup_elements).unwrap();
+
+        // Create the trace evaluation TreeVec
+        let trace = TreeVec::new(vec![preprocessed_trace, instruction_trace, interaction_trace]);
+
+        // Interpolate the trace for the evaluation
+        let trace_polys = trace.map_cols(CircleEvaluation::interpolate);
+
+        // Get the Instruction AIR evaluator
+        let instruction_eval = ProgramEval::new(&claim, instruction_lookup_elements);
+
+        // Assert that the constraints are valid for a valid Brainfuck program.
+        assert_constraints(
+            &trace_polys,
+            CanonicCoset::new(LOG_SIZE),
+            |eval| {
+                instruction_eval.evaluate(eval);
+            },
+            (interaction_claim.claimed_sum, None),
+        );
     }
 }