Skip to content

Commit

Permalink
Implement trace generator for fibonacci.
Browse files Browse the repository at this point in the history
  • Loading branch information
alonh5 committed Jul 8, 2024
1 parent 4e36e51 commit 640151c
Show file tree
Hide file tree
Showing 9 changed files with 192 additions and 29 deletions.
14 changes: 9 additions & 5 deletions crates/prover/src/core/prover/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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,
}
Expand All @@ -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()
}
}

Expand Down
80 changes: 74 additions & 6 deletions crates/prover/src/examples/fibonacci/air.rs
Original file line number Diff line number Diff line change
@@ -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,
}
Expand Down Expand Up @@ -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()
}
}

Expand All @@ -52,6 +111,7 @@ impl AirProver<CpuBackend> for FibonacciAir {
}
}

#[derive(Clone)]
pub struct MultiFibonacciAir {
pub components: Vec<FibonacciComponent>,
}
Expand Down Expand Up @@ -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()
}
}

Expand Down
62 changes: 53 additions & 9 deletions crates/prover/src/examples/fibonacci/component.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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(
Expand All @@ -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)
}
}

Expand Down
28 changes: 26 additions & 2 deletions crates/prover/src/examples/fibonacci/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ use crate::core::vcs::hasher::Hasher;
pub mod air;
mod component;

#[derive(Clone)]
pub struct Fibonacci {
pub air: FibonacciAir,
}
Expand Down Expand Up @@ -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> {
Expand All @@ -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> {
Expand Down Expand Up @@ -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;
Expand Down
10 changes: 8 additions & 2 deletions crates/prover/src/examples/poseidon/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
}
Expand All @@ -55,6 +56,7 @@ impl PoseidonComponent {
}
}

#[derive(Clone)]
pub struct PoseidonAir {
pub component: PoseidonComponent,
}
Expand All @@ -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()
}
}

Expand Down
1 change: 1 addition & 0 deletions crates/prover/src/examples/wide_fibonacci/component.rs
Original file line number Diff line number Diff line change
Expand Up @@ -169,6 +169,7 @@ impl WideFibComponent {
}
}

#[derive(Clone)]
pub struct WideFibAir {
pub component: WideFibComponent,
}
Expand Down
8 changes: 6 additions & 2 deletions crates/prover/src/examples/wide_fibonacci/constraint_eval.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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()
}
}

Expand Down
9 changes: 7 additions & 2 deletions crates/prover/src/examples/wide_fibonacci/simd.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
}
Expand All @@ -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()
}
}

Expand Down
Loading

0 comments on commit 640151c

Please sign in to comment.