PR feedback

starky/src/fibonacci_stark.rs

@@ -12,26 +12,33 @@ use crate::vars::{StarkEvaluationTargets, StarkEvaluationVars};
 /// Toy STARK system used for testing.
 /// Computes a Fibonacci sequence with inital values `x0, x1` using the transition
 /// `x0 <- x1, x1 <- x0 + x1`.
-pub struct FibonacciStark<F: RichField + Extendable<D>, const D: usize> {
+struct FibonacciStark<F: RichField + Extendable<D>, const D: usize> {
     x0: F,
     x1: F,
+    num_rows: usize,
     _phantom: PhantomData<F>,
 }
 
 impl<F: RichField + Extendable<D>, const D: usize> FibonacciStark<F, D> {
-    const NUM_COLUMNS: usize = 2;
-    const NUM_ROWS: usize = 1 << 5;
+    // The first public input is `x0`.
+    const PI_INDEX_X0: usize = 0;
+    // The second public input is `x1`.
+    const PI_INDEX_X1: usize = 1;
+    // The third public input is the second element of the last row, which should be equal to the
+    // `(num_rows + 1)`-th Fibonacci number.
+    const PI_INDEX_RES: usize = 2;
 
-    fn new(x0: F, x1: F) -> Self {
+    fn new(num_rows: usize, x0: F, x1: F) -> Self {
         Self {
             x0,
             x1,
+            num_rows,
             _phantom: PhantomData,
         }
     }
 
-    fn generate_trace(&self) -> Vec<[F; Self::NUM_COLUMNS]> {
-        (0..Self::NUM_ROWS)
+    fn generate_trace(&self) -> Vec<[F; Self::COLUMNS]> {
+        (0..self.num_rows)
             .scan([self.x0, self.x1], |acc, _| {
                 let tmp = *acc;
                 acc[0] = tmp[1];
@@ -43,8 +50,8 @@ impl<F: RichField + Extendable<D>, const D: usize> FibonacciStark<F, D> {
 }
 
 impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for FibonacciStark<F, D> {
-    const COLUMNS: usize = Self::NUM_COLUMNS;
-    const PUBLIC_INPUTS: usize = 0;
+    const COLUMNS: usize = 2;
+    const PUBLIC_INPUTS: usize = 3;
 
     fn eval_packed_generic<FE, P, const D2: usize>(
         &self,
@@ -54,6 +61,9 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for FibonacciStar
         FE: FieldExtension<D2, BaseField = F>,
         P: PackedField<Scalar = FE>,
     {
+        yield_constr.one_first_row(vars.local_values[0] - vars.public_inputs[Self::PI_INDEX_X0]);
+        yield_constr.one_first_row(vars.local_values[1] - vars.public_inputs[Self::PI_INDEX_X1]);
+        yield_constr.one_last_row(vars.local_values[1] - vars.public_inputs[Self::PI_INDEX_RES]);
         // x0 <- x1
         yield_constr.one(vars.next_values[0] - vars.local_values[1]);
         // x1 <- x0 + x1
@@ -81,6 +91,10 @@ mod tests {
     use crate::fibonacci_stark::FibonacciStark;
     use crate::prover::prove;
 
+    fn fibonacci(n: usize, x0: usize, x1: usize) -> usize {
+        (0..n).fold((0, 1), |x, _| (x.1, x.0 + x.1)).1
+    }
+
     #[test]
     fn test_fibonacci_stark() -> Result<()> {
         const D: usize = 2;
@@ -89,9 +103,21 @@ mod tests {
         type S = FibonacciStark<F, D>;
 
         let config = StarkConfig::standard_fast_config();
-        let stark = S::new(F::ZERO, F::ONE);
+        let num_rows = 1 << 5;
+        let public_inputs = [
+            F::ZERO,
+            F::ONE,
+            F::from_canonical_usize(fibonacci(num_rows - 1, 0, 1)),
+        ];
+        let stark = S::new(num_rows, public_inputs[0], public_inputs[1]);
         let trace = stark.generate_trace();
-        prove::<F, C, S, D>(stark, config, trace, &mut TimingTree::default())?;
+        prove::<F, C, S, D>(
+            stark,
+            config,
+            trace,
+            public_inputs,
+            &mut TimingTree::default(),
+        )?;
 
         Ok(())
     }

starky/src/prover.rs

@@ -25,6 +25,7 @@ pub fn prove<F, C, S, const D: usize>(
     stark: S,
     config: StarkConfig,
     trace: Vec<[F; S::COLUMNS]>,
+    public_inputs: [F; S::PUBLIC_INPUTS],
     timing: &mut TimingTree,
 ) -> Result<StarkProof<F, C, D>>
 where
@@ -72,6 +73,7 @@ where
     let quotient_polys = compute_quotient_polys::<F, C, S, D>(
         &stark,
         &trace_commitment,
+        public_inputs,
         &alphas,
         degree_bits,
         rate_bits,
@@ -142,6 +144,7 @@ where
 fn compute_quotient_polys<F, C, S, const D: usize>(
     stark: &S,
     trace_commitment: &PolynomialBatch<F, C, D>,
+    public_inputs: [F; S::PUBLIC_INPUTS],
     alphas: &[F],
     degree_bits: usize,
     rate_bits: usize,
@@ -196,7 +199,7 @@ where
                                     trace_commitment,
                                     (i + 1) % (degree << rate_bits),
                                 ),
-                                public_inputs: &[F::ZERO; S::PUBLIC_INPUTS],
+                                public_inputs: &public_inputs,
                             };
                         stark.eval_packed_base(vars, &mut consumer);
                         // TODO: Fix this once we a genuine `PackedField`.

system_zero/src/system_zero.rs

@@ -83,27 +83,33 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for SystemZero<F,
 
 #[cfg(test)]
 mod tests {
+    use anyhow::Result;
     use log::Level;
+    use plonky2::field::field_types::Field;
     use plonky2::field::goldilocks_field::GoldilocksField;
     use plonky2::plonk::config::PoseidonGoldilocksConfig;
     use plonky2::util::timing::TimingTree;
     use starky::config::StarkConfig;
     use starky::prover::prove;
+    use starky::stark::Stark;
 
     use crate::system_zero::SystemZero;
 
     #[test]
     #[ignore] // TODO
-    fn run() {
+    fn run() -> Result<()> {
         type F = GoldilocksField;
         type C = PoseidonGoldilocksConfig;
         const D: usize = 2;
 
         type S = SystemZero<F, D>;
         let system = S::default();
+        let public_inputs = [F::ZERO; S::PUBLIC_INPUTS];
         let config = StarkConfig::standard_fast_config();
         let mut timing = TimingTree::new("prove", Level::Debug);
         let trace = system.generate_trace();
-        prove::<F, C, S, D>(system, config, trace, &mut timing).unwrap();
+        prove::<F, C, S, D>(system, config, trace, public_inputs, &mut timing)?;
+
+        Ok(())
     }
 }