1// Imports all basic types.
  2use plonky2::field::extension::{Extendable, FieldExtension};
  3use plonky2::field::packed::PackedField;
  4use plonky2::field::polynomial::PolynomialValues;
  5use plonky2::hash::hash_types::RichField;
  6use std::marker::PhantomData;
  7
  8// Imports to define the constraints of our STARK.
  9use starky::constraint_consumer::{ConstraintConsumer, RecursiveConstraintConsumer};
 10use starky::evaluation_frame::{StarkEvaluationFrame, StarkFrame};
 11use starky::stark::Stark;
 12
 13// Imports to define the recursive constraints of our STARK.
 14use plonky2::iop::ext_target::ExtensionTarget;
 15use plonky2::plonk::circuit_builder::CircuitBuilder;
 16use starky::util::trace_rows_to_poly_values;
 17
 18#[derive(Copy, Clone)]
 19pub struct FibonacciStark<F: RichField + Extendable<D>, const D: usize> {
 20    num_rows: usize,
 21    _phantom: PhantomData<F>,
 22}
 23// Define witness generation.
 24impl<F: RichField + Extendable<D>, const D: usize> FibonacciStark<F, D> {
 25    // The first public input is `x0`.
 26    const PI_INDEX_X0: usize = 0;
 27    // The second public input is `x1`.
 28    const PI_INDEX_X1: usize = 1;
 29    // The third public input is the second element of the last row,
 30    // which should be equal to the `num_rows`-th Fibonacci number.
 31    const PI_INDEX_RES: usize = 2;
 32
 33    pub fn new(num_rows: usize) -> Self {
 34        Self {
 35            num_rows,
 36            _phantom: PhantomData,
 37        }
 38    }
 39
 40    /// Generate the trace using `x0, x1, 0` as initial state values.
 41    pub fn generate_trace(&self, x0: F, x1: F) -> Vec<PolynomialValues<F>> {
 42        let trace_rows = (0..self.num_rows)
 43            .scan([x0, x1, F::ZERO], |acc, _| {
 44                let tmp = *acc;
 45                acc[0] = tmp[1];
 46                acc[1] = tmp[0] + tmp[1];
 47                acc[2] = tmp[2] + F::ONE;
 48                Some(tmp)
 49            })
 50            .collect::<Vec<_>>();
 51        // Transpose the row-wise trace for the prover.
 52        trace_rows_to_poly_values(trace_rows)
 53    }
 54
 55    /// performs the same logic as in `generate_trace`, but natively in rust
 56    pub fn compute_native(&self, x0: F, x1: F) -> F {
 57        (0..self.num_rows - 1)
 58            .fold((x0, x1), |acc, _| (acc.1, acc.0 + acc.1))
 59            .1
 60    }
 61}
 62
 63// Define constraints.
 64const COLUMNS: usize = 3;
 65const PUBLIC_INPUTS: usize = 3;
 66
 67impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for FibonacciStark<F, D> {
 68    type EvaluationFrame<FE, P, const D2: usize>
 69        = StarkFrame<P, P::Scalar, COLUMNS, PUBLIC_INPUTS>
 70    where
 71        FE: FieldExtension<D2, BaseField = F>,
 72        P: PackedField<Scalar = FE>;
 73
 74    type EvaluationFrameTarget =
 75        StarkFrame<ExtensionTarget<D>, ExtensionTarget<D>, COLUMNS, PUBLIC_INPUTS>;
 76
 77    // Define this STARK's constraints.
 78    fn eval_packed_generic<FE, P, const D2: usize>(
 79        &self,
 80        vars: &Self::EvaluationFrame<FE, P, D2>, // evaluation frame of a STARK table
 81        yield_constr: &mut ConstraintConsumer<P>,
 82    ) where
 83        FE: FieldExtension<D2, BaseField = F>,
 84        P: PackedField<Scalar = FE>,
 85    {
 86        let local_values = vars.get_local_values(); // ie. current row
 87        let next_values = vars.get_next_values(); // ie. next row
 88        let public_inputs = vars.get_public_inputs(); // ie. public inputs
 89
 90        // Check public inputs.
 91        yield_constr.constraint_first_row(local_values[0] - public_inputs[Self::PI_INDEX_X0]);
 92        yield_constr.constraint_first_row(local_values[1] - public_inputs[Self::PI_INDEX_X1]);
 93        yield_constr.constraint_last_row(local_values[1] - public_inputs[Self::PI_INDEX_RES]);
 94
 95        // Enforce the Fibonacci transition constraints.
 96        // x0' <- x1
 97        yield_constr.constraint_transition(next_values[0] - local_values[1]);
 98        // x1' <- x0 + x1
 99        yield_constr.constraint_transition(next_values[1] - local_values[0] - local_values[1]);
100    }
101
102    // Define the constraints to recursively verify this STARK.
103    fn eval_ext_circuit(
104        &self,
105        builder: &mut CircuitBuilder<F, D>,
106        vars: &Self::EvaluationFrameTarget,
107        yield_constr: &mut RecursiveConstraintConsumer<F, D>,
108    ) {
109        let local_values = vars.get_local_values();
110        let next_values = vars.get_next_values();
111        let public_inputs = vars.get_public_inputs();
112
113        // Check public inputs.
114        let pis_constraints = [
115            builder.sub_extension(local_values[0], public_inputs[Self::PI_INDEX_X0]),
116            builder.sub_extension(local_values[1], public_inputs[Self::PI_INDEX_X1]),
117            builder.sub_extension(local_values[1], public_inputs[Self::PI_INDEX_RES]),
118        ];
119
120        yield_constr.constraint_first_row(builder, pis_constraints[0]);
121        yield_constr.constraint_first_row(builder, pis_constraints[1]);
122        yield_constr.constraint_last_row(builder, pis_constraints[2]);
123
124        // Enforce the Fibonacci transition constraints.
125        // x0' <- x1
126        let first_col_constraint = builder.sub_extension(next_values[0], local_values[1]);
127        yield_constr.constraint_transition(builder, first_col_constraint);
128        // x1' <- x0 + x1
129        let second_col_constraint = {
130            let tmp = builder.sub_extension(next_values[1], local_values[0]);
131            builder.sub_extension(tmp, local_values[1])
132        };
133        yield_constr.constraint_transition(builder, second_col_constraint);
134    }
135
136    fn constraint_degree(&self) -> usize {
137        2
138    }
139}
140
141#[cfg(test)]
142mod tests {
143    use super::*;
144
145    use anyhow::Result;
146    use plonky2::field::types::Field;
147    use plonky2::util::timing::TimingTree;
148    use starky::prover::prove;
149    use starky::verifier::verify_stark_proof;
150
151    use crate::{C, D, F, STARK_CONFIG};
152
153    /// test that instantiates a new `FibonacciStark` instance, generates an associated STARK
154    /// trace, and generates a proof for it.
155    #[test]
156    fn test_fibonacci_stark() -> Result<()> {
157        let num_rows = 1 << 10;
158        let x0 = F::from_canonical_u32(2);
159        let x1 = F::from_canonical_u32(7);
160
161        let stark = FibonacciStark::<F, D>::new(num_rows);
162        let public_inputs = [x0, x1, stark.compute_native(x0, x1)];
163        let trace = stark.generate_trace(public_inputs[0], public_inputs[1]);
164
165        let proof = prove::<F, C, FibonacciStark<F, D>, D>(
166            stark,
167            &STARK_CONFIG,
168            trace,
169            &public_inputs,
170            None,
171            &mut TimingTree::default(),
172        )
173        .expect("We should have a valid proof!");
174
175        verify_stark_proof(stark, proof, &STARK_CONFIG, None)
176            .expect("We should be able to verify this proof!");
177        Ok(())
178    }
179}