diff --git a/crates/brainfuck_prover/src/brainfuck_air/mod.rs b/crates/brainfuck_prover/src/brainfuck_air/mod.rs index e20233f..3210004 100644 --- a/crates/brainfuck_prover/src/brainfuck_air/mod.rs +++ b/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> { @@ -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> { - 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. diff --git a/crates/brainfuck_prover/src/components/program/component.rs b/crates/brainfuck_prover/src/components/program/component.rs index 243bbda..711d165 100644 --- a/crates/brainfuck_prover/src/components/program/component.rs +++ b/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; /// 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(&self, _eval: E) -> E { - todo!() + fn evaluate(&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), + ); + } +} diff --git a/crates/brainfuck_prover/src/components/program/table.rs b/crates/brainfuck_prover/src/components/program/table.rs index e202ae1..2329e8b 100644 --- a/crates/brainfuck_prover/src/components/program/table.rs +++ b/crates/brainfuck_prover/src/components/program/table.rs @@ -1,18 +1,16 @@ use super::component::InteractionClaim; -use crate::components::{ProgramClaim, TraceColumn, TraceError, TraceEval}; +use crate::components::{ + instruction::table::InstructionElements, ProgramClaim, TraceColumn, TraceError, TraceEval, +}; use brainfuck_vm::{machine::ProgramMemory, registers::Registers}; use num_traits::{One, Zero}; use stwo_prover::{ - constraint_framework::{ - logup::{LogupTraceGenerator, LookupElements}, - Relation, RelationEFTraitBound, - }, + constraint_framework::{logup::LogupTraceGenerator, Relation}, core::{ backend::{ simd::{column::BaseColumn, m31::LOG_N_LANES, qm31::PackedSecureField}, Column, }, - channel::Channel, fields::m31::BaseField, poly::circle::{CanonicCoset, CircleEvaluation}, }, @@ -222,58 +220,6 @@ impl TraceColumn for ProgramColumn { } } -/// The number of random elements necessary for the Program lookup arguments. -const PROGRAM_LOOKUP_ELEMENTS: usize = 3; - -/// The interaction elements are drawn for the extension column of the Program components. -/// -/// The logUp protocol uses these elements to combine the values of the different -/// registers of the main trace to create a random linear combination -/// of them, and use it in the denominator of the fractions in the logUp protocol. -/// -/// There are three lookup elements in the Program components: `ip`, `ci` and `ni`. -#[derive(Clone, Debug, Eq, PartialEq)] -pub struct ProgramElements(LookupElements); - -impl ProgramElements { - /// Provides dummy lookup elements. - pub fn dummy() -> Self { - Self(LookupElements::dummy()) - } - - /// Draw random elements from the Fiat-Shamir [`Channel`]. - /// - /// These elements are randomly secured, and will be use - /// to generate the interaction trace with the logUp protocol. - pub fn draw(channel: &mut impl Channel) -> Self { - Self(LookupElements::draw(channel)) - } -} - -impl> Relation for ProgramElements { - /// Combine multiple values from a basefield (e.g. [`BaseField`]) - /// and combine them to a value from an extension field (e.g. [`PackedSecureField`]) - /// - /// This is used when computing the interaction values from the main trace values. - fn combine(&self, values: &[F]) -> EF { - values - .iter() - .zip(self.0.alpha_powers) - .fold(EF::zero(), |acc, (value, power)| acc + EF::from(power) * value.clone()) - - self.0.z.into() - } - - /// Returns the name of the struct. - fn get_name(&self) -> &str { - stringify!(IoElements) - } - - /// Returns the number interaction elements. - fn get_size(&self) -> usize { - PROGRAM_LOOKUP_ELEMENTS - } -} - /// Creates the interaction trace from the main trace evaluation /// and the interaction elements for the Program component. /// @@ -296,7 +242,7 @@ impl> Relation for ProgramElements #[allow(clippy::similar_names)] pub fn interaction_trace_evaluation( main_trace_eval: &TraceEval, - lookup_elements: &ProgramElements, + lookup_elements: &InstructionElements, ) -> Result<(TraceEval, InteractionClaim), TraceError> { if main_trace_eval.is_empty() { return Err(TraceError::EmptyTrace) @@ -611,7 +557,7 @@ mod tests { #[test] fn test_empty_interaction_trace_evaluation() { let empty_eval = vec![]; - let lookup_elements = ProgramElements::dummy(); + let lookup_elements = InstructionElements::dummy(); let interaction_trace_eval = interaction_trace_evaluation(&empty_eval, &lookup_elements); assert!(matches!(interaction_trace_eval, Err(TraceError::EmptyTrace))); @@ -631,7 +577,7 @@ mod tests { let (trace_eval, claim) = program_table.trace_evaluation().unwrap(); - let lookup_elements = ProgramElements::dummy(); + let lookup_elements = InstructionElements::dummy(); let (interaction_trace_eval, interaction_claim) = interaction_trace_evaluation(&trace_eval, &lookup_elements).unwrap();