Implement trace generator for fibonacci. (#674)

crates/prover/src/core/prover/mod.rs

@@ -172,8 +172,7 @@ pub fn prove<B: Backend + MerkleOps<MerkleHasher>>(
 
     // 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<B: Backend + MerkleOps<MerkleHasher>>(
         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<C: ComponentProver<CpuBackend>> {
         component: C,
     }
@@ -393,8 +393,12 @@ mod tests {
             vec![]
         }
 
-        fn to_air_prover(&self) -> &impl AirProver<CpuBackend> {
-            self
+        fn to_air_prover(&self) -> impl AirProver<CpuBackend> {
+            self.clone()
+        }
+
+        fn composition_log_degree_bound(&self) -> u32 {
+            self.component.max_constraint_log_degree_bound()
         }
     }
 

crates/prover/src/examples/fibonacci/air.rs

@@ -1,15 +1,70 @@
 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;
 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<CpuBackend> for FibonacciAirGenerator {
+    fn write_trace(&mut self) -> Vec<CircleEvaluation<CpuBackend, BaseField, BitReversedOrder>> {
+        FibonacciTraceGenerator::write_trace("fibonacci", &mut self.registry)
+    }
+
+    fn interact(
+        &self,
+        _trace: &ColumnVec<CircleEvaluation<CpuBackend, BaseField, BitReversedOrder>>,
+        _elements: &InteractionElements,
+    ) -> Vec<CircleEvaluation<CpuBackend, BaseField, BitReversedOrder>> {
+        vec![]
+    }
+
+    fn to_air_prover(&self) -> impl AirProver<CpuBackend> {
+        let component_generator = self
+            .registry
+            .get_generator::<FibonacciTraceGenerator>("fibonacci");
+        FibonacciAir {
+            component: component_generator.component(),
+        }
+    }
+
+    fn composition_log_degree_bound(&self) -> u32 {
+        let component_generator = self
+            .registry
+            .get_generator::<FibonacciTraceGenerator>("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<CpuBackend> for FibonacciAir {
         vec![]
     }
 
-    fn to_air_prover(&self) -> &impl AirProver<CpuBackend> {
-        self
+    fn to_air_prover(&self) -> impl AirProver<CpuBackend> {
+        self.clone()
+    }
+
+    fn composition_log_degree_bound(&self) -> u32 {
+        self.component.max_constraint_log_degree_bound()
     }
 }
 
@@ -52,6 +111,7 @@ impl AirProver<CpuBackend> for FibonacciAir {
     }
 }
 
+#[derive(Clone)]
 pub struct MultiFibonacciAir {
     pub components: Vec<FibonacciComponent>,
 }
@@ -90,8 +150,16 @@ impl AirTraceGenerator<CpuBackend> for MultiFibonacciAir {
         vec![]
     }
 
-    fn to_air_prover(&self) -> &impl AirProver<CpuBackend> {
-        self
+    fn to_air_prover(&self) -> impl AirProver<CpuBackend> {
+        self.clone()
+    }
+
+    fn composition_log_degree_bound(&self) -> u32 {
+        self.components
+            .iter()
+            .map(|component| component.max_constraint_log_degree_bound())
+            .max()
+            .unwrap()
     }
 }
 

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<CpuBackend> 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<FibonacciInput>,
+}
+
+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<CpuBackend> 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<CircleEvaluation<CpuBackend, BaseField, BitReversedOrder>> {
-        vec![]
+        let trace_generator = registry.get_generator_mut::<Self>(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<CpuBackend> 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)
     }
 }
 

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<StarkProof, ProvingError> {
         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<usize> {
@@ -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;

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<SimdBackend> for PoseidonAir {
         vec![]
     }
 
-    fn to_air_prover(&self) -> &impl AirProver<SimdBackend> {
-        self
+    fn to_air_prover(&self) -> impl AirProver<SimdBackend> {
+        self.clone()
+    }
+
+    fn composition_log_degree_bound(&self) -> u32 {
+        self.component.max_constraint_log_degree_bound()
     }
 }
 

crates/prover/src/examples/wide_fibonacci/component.rs

@@ -169,6 +169,7 @@ impl WideFibComponent {
     }
 }
 
+#[derive(Clone)]
 pub struct WideFibAir {
     pub component: WideFibComponent,
 }

crates/prover/src/examples/wide_fibonacci/constraint_eval.rs

@@ -48,8 +48,12 @@ impl AirTraceGenerator<CpuBackend> for WideFibAir {
             .write_interaction_trace(&trace.iter().collect(), elements)
     }
 
-    fn to_air_prover(&self) -> &impl AirProver<CpuBackend> {
-        self
+    fn to_air_prover(&self) -> impl AirProver<CpuBackend> {
+        self.clone()
+    }
+
+    fn composition_log_degree_bound(&self) -> u32 {
+        self.component.max_constraint_log_degree_bound()
     }
 }
 

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<SimdBackend> for SimdWideFibAir {
         vec![]
     }
 
-    fn to_air_prover(&self) -> &impl AirProver<SimdBackend> {
-        self
+    fn to_air_prover(&self) -> impl AirProver<SimdBackend> {
+        self.clone()
+    }
+
+    fn composition_log_degree_bound(&self) -> u32 {
+        self.component.max_constraint_log_degree_bound()
     }
 }