fix: program termination is enforced in AIR (#5)

Author: felicityin

crates/core/executor/src/executor.rs

@@ -75,6 +75,9 @@ impl Executor {
             self.record.cpu_memory_access.push(event);
         }
 
+        self.record.public_values.start_pc = self.record.cpu_events[0].pc;
+        self.record.public_values.next_pc = self.record.cpu_events.last().unwrap().next_pc;
+
         Ok(())
     }
 

crates/core/executor/src/record.rs

@@ -1,9 +1,10 @@
 use std::sync::Arc;
 
 use hashbrown::HashMap;
+use p3_field::FieldAlgebra;
 use serde::{Deserialize, Serialize};
 
-use bf_stark::MachineRecord;
+use bf_stark::{MachineRecord, PublicValues};
 
 use crate::events::*;
 use crate::program::Program;
@@ -31,6 +32,8 @@ pub struct ExecutionRecord {
     pub cpu_memory_access: Vec<MemoryEvent>,
     /// A trace of the byte lookups that are needed.
     pub byte_lookups: HashMap<ByteLookupEvent, usize>,
+    /// The public values.
+    pub public_values: PublicValues<u32>,
 }
 
 /// A memory access record.
@@ -85,4 +88,9 @@ impl MachineRecord for ExecutionRecord {
 
         self.cpu_memory_access.append(&mut other.cpu_memory_access);
     }
+
+    /// Retrieves the public values.  This method is needed for the `MachineRecord` trait, since
+    fn public_values<F: FieldAlgebra>(&self) -> Vec<F> {
+        self.public_values.to_vec()
+    }
 }

crates/core/machine/src/cpu/air.rs

@@ -1,9 +1,12 @@
 use core::borrow::Borrow;
-use p3_air::{Air, AirBuilder, BaseAir};
+use p3_air::{Air, AirBuilder, AirBuilderWithPublicValues, BaseAir};
 use p3_field::FieldAlgebra;
 use p3_matrix::Matrix;
 
-use bf_stark::air::{BaseAirBuilder, BfAirBuilder};
+use bf_stark::{
+    air::{BaseAirBuilder, BfAirBuilder},
+    PublicValues, PROOF_NUM_PV_ELTS,
+};
 
 use crate::{
     air::{BfCoreAirBuilder, MemoryAirBuilder, U8AirBuilder},
@@ -21,7 +24,7 @@ impl<F> BaseAir<F> for CpuChip {
 
 impl<AB> Air<AB> for CpuChip
 where
-    AB: BfCoreAirBuilder,
+    AB: BfCoreAirBuilder + AirBuilderWithPublicValues,
     AB::Var: Sized,
 {
     #[inline(never)]
@@ -31,6 +34,10 @@ where
         let local: &CpuCols<AB::Var> = (*local).borrow();
         let next: &CpuCols<AB::Var> = (*next).borrow();
 
+        let public_values_slice: [AB::PublicVar; PROOF_NUM_PV_ELTS] =
+            core::array::from_fn(|i| builder.public_values()[i]);
+        let public_values: &PublicValues<AB::PublicVar> = public_values_slice.as_slice().borrow();
+
         let clk =
             AB::Expr::from_canonical_u32(1u32 << 16) * local.clk_8bit_limb + local.clk_16bit_limb;
 
@@ -47,7 +54,7 @@ where
         self.eval_clk(builder, local, next, clk.clone());
 
         // Check that the pc is updated correctly.
-        self.eval_pc(builder, local, next);
+        self.eval_pc(builder, local, next, public_values);
 
         // Check that the is_real flag is correct.
         self.eval_is_real(builder, local, next);
@@ -128,6 +135,7 @@ impl CpuChip {
         builder: &mut AB,
         local: &CpuCols<AB::Var>,
         next: &CpuCols<AB::Var>,
+        public_values: &PublicValues<AB::PublicVar>,
     ) {
         builder.when_transition().when(next.is_real).assert_eq(local.next_pc, next.pc);
 
@@ -136,6 +144,19 @@ impl CpuChip {
             .when(local.is_real)
             .when_not(local.is_jump)
             .assert_eq(local.next_pc, local.pc + AB::Expr::from_canonical_u32(1));
+
+        // Verify the public value's next pc.  We need to handle two cases:
+        // 1. The last real row is a transition row.
+        // 2. The last real row is the last row.
+
+        // If the last real row is a transition row, verify the public value's next pc.
+        builder
+            .when_transition()
+            .when(local.is_real - next.is_real)
+            .assert_eq(public_values.next_pc, local.next_pc);
+
+        // If the last real row is the last row, verify the public value's next pc.
+        builder.when_last_row().when(local.is_real).assert_eq(public_values.next_pc, local.next_pc);
     }
 
     /// Constraints related to the is_real column.

crates/core/machine/src/utils/mod.rs

@@ -23,7 +23,7 @@ pub const fn indices_arr<const N: usize>() -> [usize; N] {
 }
 
 pub fn pad_to_power_of_two<const N: usize, T: Clone + Default>(values: &mut Vec<T>) {
-    debug_assert!(values.len() % N == 0);
+    debug_assert!(values.len().is_multiple_of(N));
     let mut n_real_rows = values.len() / N;
     if n_real_rows < 16 {
         n_real_rows = 16;
@@ -75,7 +75,7 @@ where
     P: Fn(usize, &mut [F]) + Send + Sync,
 {
     // Split the vector into `num_cpus` chunks, but at least `num_cpus` rows per chunk.
-    assert!(vec.len() % num_elements_per_event == 0);
+    assert!(vec.len().is_multiple_of(num_elements_per_event));
     let len = vec.len() / num_elements_per_event;
     let cpus = num_cpus::get();
     let ceil_div = len.div_ceil(cpus);

crates/stark/src/air/builder.rs

@@ -1,6 +1,6 @@
 use std::iter::once;
 
-use p3_air::{AirBuilder, FilteredAirBuilder, PermutationAirBuilder};
+use p3_air::{AirBuilder, AirBuilderWithPublicValues, FilteredAirBuilder, PermutationAirBuilder};
 use p3_field::{Field, FieldAlgebra};
 use p3_uni_stark::{
     ProverConstraintFolder, StarkGenericConfig, SymbolicAirBuilder, VerifierConstraintFolder,
@@ -28,7 +28,7 @@ pub trait MessageBuilder<M> {
 /// A trait which contains basic methods for building an AIR.
 pub trait BaseAirBuilder: AirBuilder + MessageBuilder<AirLookup<Self::Expr>> {
     /// Returns a sub-builder whose constraints are enforced only when `condition` is not one.
-    fn when_not<I: Into<Self::Expr>>(&mut self, condition: I) -> FilteredAirBuilder<Self> {
+    fn when_not<I: Into<Self::Expr>>(&mut self, condition: I) -> FilteredAirBuilder<'_, Self> {
         self.when_ne(condition, Self::F::ONE)
     }
 
@@ -248,7 +248,7 @@ pub trait MultiTableAirBuilder<'a>: PermutationAirBuilder {
 }
 
 /// A trait that contains the common helper methods for building machine AIRs.
-pub trait MachineAirBuilder: BaseAirBuilder {}
+pub trait MachineAirBuilder: BaseAirBuilder + AirBuilderWithPublicValues {}
 
 /// A trait which contains all helper methods for building machine AIRs.
 pub trait BfAirBuilder: MachineAirBuilder + ByteAirBuilder + InstructionAirBuilder {}
@@ -267,8 +267,8 @@ impl<AB: AirBuilder + MessageBuilder<AirLookup<AB::Expr>>> BaseAirBuilder for AB
 impl<AB: BaseAirBuilder> ByteAirBuilder for AB {}
 impl<AB: BaseAirBuilder> InstructionAirBuilder for AB {}
 
-impl<AB: BaseAirBuilder> MachineAirBuilder for AB {}
-impl<AB: BaseAirBuilder> BfAirBuilder for AB {}
+impl<AB: BaseAirBuilder + AirBuilderWithPublicValues> MachineAirBuilder for AB {}
+impl<AB: BaseAirBuilder + AirBuilderWithPublicValues> BfAirBuilder for AB {}
 
 impl<SC: StarkGenericConfig> EmptyMessageBuilder for ProverConstraintFolder<'_, SC> {}
 impl<SC: StarkGenericConfig> EmptyMessageBuilder for VerifierConstraintFolder<'_, SC> {}

crates/stark/src/air/mod.rs

@@ -3,7 +3,9 @@
 mod builder;
 mod lookup;
 mod machine;
+mod public_values;
 
 pub use builder::*;
 pub use lookup::*;
 pub use machine::*;
+pub use public_values::*;

crates/stark/src/air/public_values.rs

@@ -0,0 +1,75 @@
+use core::{fmt::Debug, mem::size_of};
+use std::borrow::{Borrow, BorrowMut};
+
+use p3_field::FieldAlgebra;
+use serde::{Deserialize, Serialize};
+
+use crate::PROOF_MAX_NUM_PVS;
+
+/// The number of non padded elements in the Ziren proofs public values vec.
+pub const PROOF_NUM_PV_ELTS: usize = size_of::<PublicValues<u8>>();
+
+/// Stores all of a shard proof's public values.
+#[derive(Serialize, Deserialize, Clone, Copy, Default, Debug)]
+#[repr(C)]
+pub struct PublicValues<T> {
+    /// The shard's start program counter.
+    pub start_pc: T,
+
+    /// The expected start program counter for the next shard.
+    pub next_pc: T,
+
+    /// This field is here to ensure that the size of the public values struct is a multiple of 8.
+    pub empty: [T; 6],
+}
+
+impl PublicValues<u32> {
+    /// Convert the public values into a vector of field elements.  This function will pad the
+    /// vector to the maximum number of public values.
+    #[must_use]
+    pub fn to_vec<F: FieldAlgebra>(&self) -> Vec<F> {
+        let mut ret = vec![F::ZERO; PROOF_MAX_NUM_PVS];
+
+        let field_values = PublicValues::<F> {
+            start_pc: F::from_canonical_u32(self.start_pc),
+            next_pc: F::from_canonical_u32(self.next_pc),
+            empty: [F::ZERO; 6],
+        };
+        let ret_ref_mut: &mut PublicValues<F> = ret.as_mut_slice().borrow_mut();
+        *ret_ref_mut = field_values;
+        ret
+    }
+
+    /// Resets the public values to zero.
+    #[must_use]
+    pub fn reset(&self) -> Self {
+        let mut copy = *self;
+        copy.start_pc = 0;
+        copy.next_pc = 0;
+        copy
+    }
+}
+
+impl<T: Clone> Borrow<PublicValues<T>> for [T] {
+    fn borrow(&self) -> &PublicValues<T> {
+        let size = std::mem::size_of::<PublicValues<u8>>();
+        debug_assert!(self.len() >= size);
+        let slice = &self[0..size];
+        let (prefix, shorts, _suffix) = unsafe { slice.align_to::<PublicValues<T>>() };
+        debug_assert!(prefix.is_empty(), "Alignment should match");
+        debug_assert_eq!(shorts.len(), 1);
+        &shorts[0]
+    }
+}
+
+impl<T: Clone> BorrowMut<PublicValues<T>> for [T] {
+    fn borrow_mut(&mut self) -> &mut PublicValues<T> {
+        let size = std::mem::size_of::<PublicValues<u8>>();
+        debug_assert!(self.len() >= size);
+        let slice = &mut self[0..size];
+        let (prefix, shorts, _suffix) = unsafe { slice.align_to_mut::<PublicValues<T>>() };
+        debug_assert!(prefix.is_empty(), "Alignment should match");
+        debug_assert_eq!(shorts.len(), 1);
+        &mut shorts[0]
+    }
+}

crates/stark/src/debug.rs

@@ -4,7 +4,10 @@ use std::{
     process::exit,
 };
 
-use p3_air::{Air, AirBuilder, ExtensionBuilder, PairBuilder, PermutationAirBuilder};
+use p3_air::{
+    Air, AirBuilder, AirBuilderWithPublicValues, ExtensionBuilder, PairBuilder,
+    PermutationAirBuilder,
+};
 use p3_field::{ExtensionField, Field, FieldAlgebra, PrimeField32};
 use p3_matrix::{
     dense::{RowMajorMatrix, RowMajorMatrixView},
@@ -27,6 +30,7 @@ pub fn debug_constraints<SC, A>(
     main: &RowMajorMatrix<Val<SC>>,
     perm: &RowMajorMatrix<SC::Challenge>,
     perm_challenges: &[SC::Challenge],
+    public_values: &[Val<SC>],
     cumulative_sum: &SC::Challenge,
 ) where
     SC: StarkGenericConfig,
@@ -84,6 +88,7 @@ pub fn debug_constraints<SC, A>(
             is_first_row: Val::<SC>::ZERO,
             is_last_row: Val::<SC>::ZERO,
             is_transition: Val::<SC>::ONE,
+            public_values,
         };
         if i == 0 {
             builder.is_first_row = Val::<SC>::ONE;
@@ -130,7 +135,7 @@ pub struct DebugConstraintBuilder<'a, F: Field, EF: ExtensionField<F>> {
     pub(crate) is_first_row: F,
     pub(crate) is_last_row: F,
     pub(crate) is_transition: F,
-    // pub(crate) public_values: &'a [F],
+    pub(crate) public_values: &'a [F],
 }
 
 impl<F, EF> ExtensionBuilder for DebugConstraintBuilder<'_, F, EF>
@@ -260,3 +265,13 @@ where
 }
 
 impl<F: Field, EF: ExtensionField<F>> EmptyMessageBuilder for DebugConstraintBuilder<'_, F, EF> {}
+
+impl<F: Field, EF: ExtensionField<F>> AirBuilderWithPublicValues
+    for DebugConstraintBuilder<'_, F, EF>
+{
+    type PublicVar = F;
+
+    fn public_values(&self) -> &[Self::PublicVar] {
+        self.public_values
+    }
+}