diff --git a/crates/prover/src/core/prover/mod.rs b/crates/prover/src/core/prover/mod.rs index fd68091de..7b584f1fe 100644 --- a/crates/prover/src/core/prover/mod.rs +++ b/crates/prover/src/core/prover/mod.rs @@ -172,8 +172,7 @@ pub fn prove>( // Check that the composition polynomial is not too big. // TODO(AlonH): Get traces log degree bounds from trace writer. - let composition_polynomial_log_degree_bound = - air.to_air_prover().composition_log_degree_bound(); + let composition_polynomial_log_degree_bound = air.composition_log_degree_bound(); if composition_polynomial_log_degree_bound + LOG_BLOWUP_FACTOR > MAX_CIRCLE_DOMAIN_LOG_SIZE { return Err(ProvingError::MaxCompositionDegreeExceeded { degree: composition_polynomial_log_degree_bound, @@ -192,7 +191,7 @@ pub fn prove>( evaluate_and_commit_on_trace(air, channel, &twiddles, trace)?; generate_proof( - air.to_air_prover(), + &air.to_air_prover(), channel, &interaction_elements, &twiddles, @@ -368,6 +367,7 @@ mod tests { use crate::trace_generation::registry::ComponentGenerationRegistry; use crate::trace_generation::{AirTraceGenerator, AirTraceVerifier, ComponentTraceGenerator}; + #[derive(Clone)] struct TestAir> { component: C, } @@ -393,8 +393,12 @@ mod tests { vec![] } - fn to_air_prover(&self) -> &impl AirProver { - self + fn to_air_prover(&self) -> impl AirProver { + self.clone() + } + + fn composition_log_degree_bound(&self) -> u32 { + self.component.max_constraint_log_degree_bound() } } diff --git a/crates/prover/src/examples/fibonacci/air.rs b/crates/prover/src/examples/fibonacci/air.rs index 40c8eea6f..2cb61ff1f 100644 --- a/crates/prover/src/examples/fibonacci/air.rs +++ b/crates/prover/src/examples/fibonacci/air.rs @@ -1,6 +1,6 @@ use itertools::{zip_eq, Itertools}; -use super::component::FibonacciComponent; +use super::component::{FibonacciComponent, FibonacciInput, FibonacciTraceGenerator}; use crate::core::air::{Air, AirProver, Component, ComponentProver}; use crate::core::backend::CpuBackend; use crate::core::channel::Blake2sChannel; @@ -8,8 +8,63 @@ use crate::core::fields::m31::BaseField; use crate::core::poly::circle::CircleEvaluation; use crate::core::poly::BitReversedOrder; use crate::core::{ColumnVec, InteractionElements}; -use crate::trace_generation::{AirTraceGenerator, AirTraceVerifier}; +use crate::trace_generation::registry::ComponentGenerationRegistry; +use crate::trace_generation::{AirTraceGenerator, AirTraceVerifier, ComponentTraceGenerator}; +pub struct FibonacciAirGenerator { + pub registry: ComponentGenerationRegistry, +} + +impl FibonacciAirGenerator { + pub fn new(inputs: &FibonacciInput) -> Self { + let mut component_generator = FibonacciTraceGenerator::new(); + component_generator.add_inputs(inputs); + let mut registry = ComponentGenerationRegistry::default(); + registry.register("fibonacci", component_generator); + Self { registry } + } +} + +impl AirTraceVerifier for FibonacciAirGenerator { + fn interaction_elements(&self, _channel: &mut Blake2sChannel) -> InteractionElements { + InteractionElements::default() + } +} + +impl AirTraceGenerator for FibonacciAirGenerator { + fn write_trace(&mut self) -> Vec> { + FibonacciTraceGenerator::write_trace("fibonacci", &mut self.registry) + } + + fn interact( + &self, + _trace: &ColumnVec>, + _elements: &InteractionElements, + ) -> Vec> { + vec![] + } + + fn to_air_prover(&self) -> impl AirProver { + let component_generator = self + .registry + .get_generator::("fibonacci"); + FibonacciAir { + component: component_generator.component(), + } + } + + fn composition_log_degree_bound(&self) -> u32 { + let component_generator = self + .registry + .get_generator::("fibonacci"); + assert!(component_generator.inputs_set(), "Fibonacci input not set."); + component_generator + .component() + .max_constraint_log_degree_bound() + } +} + +#[derive(Clone)] pub struct FibonacciAir { pub component: FibonacciComponent, } @@ -41,8 +96,12 @@ impl AirTraceGenerator for FibonacciAir { vec![] } - fn to_air_prover(&self) -> &impl AirProver { - self + fn to_air_prover(&self) -> impl AirProver { + self.clone() + } + + fn composition_log_degree_bound(&self) -> u32 { + self.component.max_constraint_log_degree_bound() } } @@ -52,6 +111,7 @@ impl AirProver for FibonacciAir { } } +#[derive(Clone)] pub struct MultiFibonacciAir { pub components: Vec, } @@ -90,8 +150,16 @@ impl AirTraceGenerator for MultiFibonacciAir { vec![] } - fn to_air_prover(&self) -> &impl AirProver { - self + fn to_air_prover(&self) -> impl AirProver { + self.clone() + } + + fn composition_log_degree_bound(&self) -> u32 { + self.components + .iter() + .map(|component| component.max_constraint_log_degree_bound()) + .max() + .unwrap() } } diff --git a/crates/prover/src/examples/fibonacci/component.rs b/crates/prover/src/examples/fibonacci/component.rs index f683ae8b3..88ac52a0e 100644 --- a/crates/prover/src/examples/fibonacci/component.rs +++ b/crates/prover/src/examples/fibonacci/component.rs @@ -18,7 +18,7 @@ use crate::core::prover::BASE_TRACE; use crate::core::utils::bit_reverse_index; use crate::core::{ColumnVec, InteractionElements, LookupValues}; use crate::trace_generation::registry::ComponentGenerationRegistry; -use crate::trace_generation::ComponentTraceGenerator; +use crate::trace_generation::{ComponentGen, ComponentTraceGenerator}; #[derive(Clone)] pub struct FibonacciComponent { @@ -126,17 +126,60 @@ impl Component for FibonacciComponent { } } -impl ComponentTraceGenerator for FibonacciComponent { - type Component = Self; - type Inputs = (); +#[derive(Copy, Clone)] +pub struct FibonacciInput { + pub log_size: u32, + pub claim: BaseField, +} - fn add_inputs(&mut self, _inputs: &Self::Inputs) {} +#[derive(Clone)] +pub struct FibonacciTraceGenerator { + input: Option, +} + +impl ComponentGen for FibonacciTraceGenerator {} + +impl FibonacciTraceGenerator { + #[allow(clippy::new_without_default)] + pub fn new() -> Self { + Self { input: None } + } + + pub fn inputs_set(&self) -> bool { + self.input.is_some() + } +} + +impl ComponentTraceGenerator for FibonacciTraceGenerator { + type Component = FibonacciComponent; + type Inputs = FibonacciInput; + + fn add_inputs(&mut self, inputs: &Self::Inputs) { + assert!(!self.inputs_set(), "Fibonacci input already set."); + self.input = Some(*inputs); + } fn write_trace( - _component_id: &str, - _registry: &mut ComponentGenerationRegistry, + component_id: &str, + registry: &mut ComponentGenerationRegistry, ) -> ColumnVec> { - vec![] + let trace_generator = registry.get_generator_mut::(component_id); + assert!(trace_generator.inputs_set(), "Fibonacci input not set."); + let trace_domain = CanonicCoset::new(trace_generator.input.unwrap().log_size); + let mut trace = Vec::with_capacity(trace_domain.size()); + + // Fill trace with fibonacci squared. + let mut a = BaseField::one(); + let mut b = BaseField::one(); + for _ in 0..trace_domain.size() { + trace.push(a); + let tmp = a.square() + b.square(); + a = b; + b = tmp; + } + + // Returns as a CircleEvaluation. + vec![CircleEvaluation::new_canonical_ordered(trace_domain, trace)] } fn write_interaction_trace( @@ -148,7 +191,8 @@ impl ComponentTraceGenerator for FibonacciComponent { } fn component(&self) -> Self::Component { - self.clone() + assert!(self.inputs_set(), "Fibonacci input not set."); + FibonacciComponent::new(self.input.unwrap().log_size, self.input.unwrap().claim) } } diff --git a/crates/prover/src/examples/fibonacci/mod.rs b/crates/prover/src/examples/fibonacci/mod.rs index 8fac01513..d2fe94b6a 100644 --- a/crates/prover/src/examples/fibonacci/mod.rs +++ b/crates/prover/src/examples/fibonacci/mod.rs @@ -17,6 +17,7 @@ use crate::core::vcs::hasher::Hasher; pub mod air; mod component; +#[derive(Clone)] pub struct Fibonacci { pub air: FibonacciAir, } @@ -97,7 +98,8 @@ impl MultiFibonacci { pub fn prove(&self) -> Result { let channel = &mut Blake2sChannel::new(Blake2sHasher::hash(BaseField::into_slice(&self.claims))); - prove(&self.air, channel, self.get_trace()) + let trace = self.get_trace(); + prove(&self.air, channel, trace) } pub fn verify(&self, proof: StarkProof) -> Result<(), VerificationError> { @@ -120,15 +122,22 @@ mod tests { use super::{Fibonacci, MultiFibonacci}; use crate::core::air::accumulation::PointEvaluationAccumulator; use crate::core::air::{AirExt, AirProverExt, Component, ComponentTrace}; + use crate::core::channel::{Blake2sChannel, Channel}; use crate::core::circle::CirclePoint; use crate::core::fields::m31::BaseField; use crate::core::fields::qm31::SecureField; + use crate::core::fields::IntoSlice; use crate::core::pcs::TreeVec; use crate::core::poly::circle::CanonicCoset; - use crate::core::prover::{VerificationError, BASE_TRACE}; + use crate::core::prover::{prove, VerificationError, BASE_TRACE}; use crate::core::queries::Queries; use crate::core::utils::bit_reverse; + use crate::core::vcs::blake2_hash::Blake2sHasher; + use crate::core::vcs::hasher::Hasher; use crate::core::{InteractionElements, LookupValues}; + use crate::examples::fibonacci::air::FibonacciAirGenerator; + use crate::examples::fibonacci::component::FibonacciInput; + use crate::trace_generation::AirTraceGenerator; use crate::{m31, qm31}; pub fn generate_test_queries(n_queries: usize, trace_length: usize) -> Vec { @@ -232,6 +241,21 @@ mod tests { fib.verify(proof).unwrap(); } + #[test] + fn test_fib_prove_2() { + const FIB_LOG_SIZE: u32 = 5; + const CLAIM: BaseField = m31!(443693538); + let mut fib_trace_generator = FibonacciAirGenerator::new(&FibonacciInput { + log_size: FIB_LOG_SIZE, + claim: CLAIM, + }); + + let trace = fib_trace_generator.write_trace(); + let channel = + &mut Blake2sChannel::new(Blake2sHasher::hash(BaseField::into_slice(&[CLAIM]))); + prove(&fib_trace_generator, channel, trace).unwrap(); + } + #[test] fn test_prove_invalid_trace_value() { const FIB_LOG_SIZE: u32 = 5; diff --git a/crates/prover/src/examples/poseidon/mod.rs b/crates/prover/src/examples/poseidon/mod.rs index 5d5d65255..947579aeb 100644 --- a/crates/prover/src/examples/poseidon/mod.rs +++ b/crates/prover/src/examples/poseidon/mod.rs @@ -41,6 +41,7 @@ const EXTERNAL_ROUND_CONSTS: [[BaseField; N_STATE]; 2 * N_HALF_FULL_ROUNDS] = const INTERNAL_ROUND_CONSTS: [BaseField; N_PARTIAL_ROUNDS] = [BaseField::from_u32_unchecked(1234); N_PARTIAL_ROUNDS]; +#[derive(Clone)] pub struct PoseidonComponent { pub log_n_instances: u32, } @@ -55,6 +56,7 @@ impl PoseidonComponent { } } +#[derive(Clone)] pub struct PoseidonAir { pub component: PoseidonComponent, } @@ -80,8 +82,12 @@ impl AirTraceGenerator for PoseidonAir { vec![] } - fn to_air_prover(&self) -> &impl AirProver { - self + fn to_air_prover(&self) -> impl AirProver { + self.clone() + } + + fn composition_log_degree_bound(&self) -> u32 { + self.component.max_constraint_log_degree_bound() } } diff --git a/crates/prover/src/examples/wide_fibonacci/component.rs b/crates/prover/src/examples/wide_fibonacci/component.rs index 4e0483826..26078c6bb 100644 --- a/crates/prover/src/examples/wide_fibonacci/component.rs +++ b/crates/prover/src/examples/wide_fibonacci/component.rs @@ -169,6 +169,7 @@ impl WideFibComponent { } } +#[derive(Clone)] pub struct WideFibAir { pub component: WideFibComponent, } diff --git a/crates/prover/src/examples/wide_fibonacci/constraint_eval.rs b/crates/prover/src/examples/wide_fibonacci/constraint_eval.rs index a75d28f67..0a61ff1e4 100644 --- a/crates/prover/src/examples/wide_fibonacci/constraint_eval.rs +++ b/crates/prover/src/examples/wide_fibonacci/constraint_eval.rs @@ -48,8 +48,12 @@ impl AirTraceGenerator for WideFibAir { .write_interaction_trace(&trace.iter().collect(), elements) } - fn to_air_prover(&self) -> &impl AirProver { - self + fn to_air_prover(&self) -> impl AirProver { + self.clone() + } + + fn composition_log_degree_bound(&self) -> u32 { + self.component.max_constraint_log_degree_bound() } } diff --git a/crates/prover/src/examples/wide_fibonacci/simd.rs b/crates/prover/src/examples/wide_fibonacci/simd.rs index b9d325b61..4721c3e80 100644 --- a/crates/prover/src/examples/wide_fibonacci/simd.rs +++ b/crates/prover/src/examples/wide_fibonacci/simd.rs @@ -50,6 +50,7 @@ impl SimdWideFibComponent { } // TODO(AlonH): Remove this once the Cpu and Simd implementations are aligned. +#[derive(Clone)] pub struct SimdWideFibAir { pub component: SimdWideFibComponent, } @@ -75,8 +76,12 @@ impl AirTraceGenerator for SimdWideFibAir { vec![] } - fn to_air_prover(&self) -> &impl AirProver { - self + fn to_air_prover(&self) -> impl AirProver { + self.clone() + } + + fn composition_log_degree_bound(&self) -> u32 { + self.component.max_constraint_log_degree_bound() } } diff --git a/crates/prover/src/trace_generation/mod.rs b/crates/prover/src/trace_generation/mod.rs index a11b5da5b..7019cc94f 100644 --- a/crates/prover/src/trace_generation/mod.rs +++ b/crates/prover/src/trace_generation/mod.rs @@ -49,11 +49,18 @@ pub trait AirTraceVerifier { } pub trait AirTraceGenerator: AirTraceVerifier { + fn composition_log_degree_bound(&self) -> u32; + + // TODO(AlonH): Remove default implementation once all the components are implemented. + fn write_trace(&mut self) -> Vec> { + vec![] + } + fn interact( &self, trace: &ColumnVec>, elements: &InteractionElements, ) -> Vec>; - fn to_air_prover(&self) -> &impl AirProver; + fn to_air_prover(&self) -> impl AirProver; }