feat: add Program component constraints

crates/brainfuck_prover/src/brainfuck_air/mod.rs

@@ -14,7 +14,12 @@ use crate::components::{
         component::{MemoryComponent, MemoryEval},
         table::{MemoryElements, MemoryTable},
     },
-    InstructionClaim, IoClaim, MemoryClaim,
+    program::{
+        self,
+        component::{ProgramComponent, ProgramEval},
+        table::ProgramTable,
+    },
+    InstructionClaim, 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: io::component::InteractionClaim,
     memory: memory::component::InteractionClaim,
     instruction: instruction::component::InteractionClaim,
+    program: program::component::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,8 +1,14 @@
-use super::table::ProgramElements;
-use crate::components::ProgramClaim;
+use crate::components::{instruction::table::InstructionElements, ProgramClaim};
+use num_traits::One;
 use stwo_prover::{
-    constraint_framework::{EvalAtRow, FrameworkComponent, FrameworkEval},
-    core::{channel::Channel, fields::qm31::SecureField},
+    constraint_framework::{
+        preprocessed_columns::PreprocessedColumn, EvalAtRow, FrameworkComponent, FrameworkEval,
+        RelationEntry,
+    },
+    core::{
+        channel::Channel,
+        fields::{m31::BaseField, qm31::SecureField},
+    },
 };
 
 /// Implementation of `Component` and `ComponentProver`
@@ -15,12 +21,15 @@ pub type ProgramComponent = FrameworkComponent<ProgramEval>;
 /// provided by the constraint framework of Stwo.
 pub struct ProgramEval {
     log_size: u32,
-    _program_lookup_elements: ProgramElements,
+    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 }
     }
 }
 
@@ -50,8 +59,38 @@ 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
+        eval.add_constraint(is_first * ip.clone());
+
+        // Consistency constraints
+        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());
+
+        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
     }
 }
 
@@ -72,3 +111,72 @@ impl InteractionClaim {
         channel.mix_felts(&[self.claimed_sum]);
     }
 }
+
+#[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),
+        );
+    }
+}