fib2_winterfell.rs

use winter_math::{FieldElement, fields::f64::BaseElement as Felt};
use winter_prover::{Prover, Trace};

/* +---+---+   Fibonacci: 2 columns per row. 
    | a | b |   Single column fibonacci is not possible with this constraint system as that would require looking back more than previous time step.
    +---+---+   Transition constraints:
    | c | d |   c = a + b
    +---+---+   d = b + c  */
fn main() {
    let (f0, f1) = (Felt::new(1), Felt::new(1));
    let nrows = 1024; // This has to be a power of 2
    let trace = build_fib2_trace(f0, f1, nrows);

    let num_queries = 32;
    let blowup_factor = 4;
    let grinding_factor = 21;
    let field_extension = winter_air::FieldExtension::None;
    let fri_folding_factor = 8;
    let fri_remainder_max_degree = 63;
    let options = winter_prover::ProofOptions::new(num_queries, blowup_factor, grinding_factor, field_extension, fri_folding_factor, fri_remainder_max_degree);
    let prover = Fib2Prover { options };
    let proof = prover.prove(trace).unwrap();
    
    println!("Proof size: {:.2} KB", proof.to_bytes().len() as f64 / 1024.0);
}

fn build_fib2_trace(f0: Felt, f1: Felt, nrows: usize) -> winter_prover::TraceTable<Felt> {
    let mut col0 = vec![Felt::new(0); nrows];
    let mut col1 = vec![Felt::new(0); nrows];
    (col0[0], col1[0]) = (f0, f1);
    for i in 1..nrows {
        col0[i] = col1[i-1] + col0[i-1];
        col1[i] = col0[i]   + col1[i-1];
    }
    winter_prover::TraceTable::init(vec![col0, col1])
}

struct Fib2Params {
    input: (Felt, Felt),
    output: Felt
}

impl winter_math::ToElements<Felt> for Fib2Params {
    fn to_elements(&self) -> Vec<Felt> {
        vec![self.input.0, self.input.1, self.output]
    }
}

struct Fib2Air {
    context: winter_air::AirContext<Felt>,
    pub_inputs: Fib2Params
}

impl winter_air::Air for Fib2Air {
    type BaseField = Felt;
    type PublicInputs = Fib2Params;

    // Boundary constraints
    fn get_assertions(&self) -> Vec<winter_air::Assertion<Self::BaseField>> { // let num_assertions = 3
        vec![winter_air::Assertion::single(0, 0, self.pub_inputs.input.0),
             winter_air::Assertion::single(1, 0, self.pub_inputs.input.1),
             winter_air::Assertion::single(1, self.trace_length() - 1, self.pub_inputs.output)]
    }

    fn evaluate_transition<E: FieldElement<BaseField = Self::BaseField>>(&self, frame: &winter_air::EvaluationFrame<E>, _periodic_values: &[E], result: &mut [E]) {
        result[0] = frame.current()[0] + frame.current()[1] - frame.next()[0];  // winter_air::TransitionConstraintDegree::new(1)
        result[1] = frame.current()[1] + frame.next()[0]    - frame.next()[1];  // winter_air::TransitionConstraintDegree::new(1)
    }

    fn new(trace_info: winter_air::TraceInfo, pub_inputs: Self::PublicInputs, options: winter_air::ProofOptions) -> Self {
        let constraint1_deg = winter_air::TransitionConstraintDegree::new(1);
        let constraint2_deg = winter_air::TransitionConstraintDegree::new(1);
        let transition_constraint_degrees = vec![constraint1_deg, constraint2_deg];
        let num_assertions = 3; // self.get_assertions().len()
        let context = winter_air::AirContext::new(trace_info, transition_constraint_degrees, num_assertions, options);
        Self { context, pub_inputs }
    }

    fn context(&self) -> &winter_air::AirContext<Self::BaseField> {
        &self.context
    }
}

struct Fib2Prover {
    options: winter_prover::ProofOptions
}

impl winter_prover::Prover for Fib2Prover {
    type BaseField = Felt;
    type Air = Fib2Air;
    type Trace = winter_prover::TraceTable<Self::BaseField>;
    type HashFn = winter_crypto::hashers::Blake3_192<Self::BaseField>;
    type RandomCoin = winter_crypto::DefaultRandomCoin<Self::HashFn>;
    
    fn get_pub_inputs(&self, trace: &Self::Trace) -> <<Self as winter_prover::Prover>::Air as winter_air::Air>::PublicInputs {
        let input = (trace.get(0, 0), trace.get(1, 0));
        let output = trace.get(1, trace.length() - 1);
        Fib2Params { input, output }
    }

    fn options(&self) -> &winter_air::ProofOptions {
        &self.options
    }

    type TraceLde<E: FieldElement<BaseField = Self::BaseField>> = winter_prover::DefaultTraceLde<E, Self::HashFn>;
    fn new_trace_lde<E: FieldElement<BaseField = Self::BaseField>>(&self, trace_info: &winter_air::TraceInfo, main_trace: &winter_prover::matrix::ColMatrix<Self::BaseField>, domain: &winter_prover::StarkDomain<Self::BaseField>) -> (Self::TraceLde<E>, winter_prover::TracePolyTable<E>) {
        Self::TraceLde::new(trace_info, main_trace, domain)
    }
    
    type ConstraintEvaluator<'a, E: FieldElement<BaseField = Self::BaseField>> = winter_prover::DefaultConstraintEvaluator<'a, Self::Air, E>;
    fn new_evaluator<'a, E: FieldElement<BaseField = Self::BaseField>>(&self, air: &'a Self::Air, aux_rand_elements: winter_air::AuxTraceRandElements<E>, composition_coefficients: winter_air::ConstraintCompositionCoefficients<E>) -> Self::ConstraintEvaluator<'a, E> {
        Self::ConstraintEvaluator::new(air, aux_rand_elements, composition_coefficients)
    }
}