feat: preprocessing api, cleanup, and testing (#195)

core/src/air/trace.rs

@@ -2,14 +2,30 @@ use p3_air::BaseAir;
 use p3_field::Field;
 use p3_matrix::dense::RowMajorMatrix;
 
-use crate::runtime::ExecutionRecord;
+use crate::runtime::{ExecutionRecord, Program};
 
+/// An AIR that is part of a Risc-V AIR arithmetization.
 pub trait MachineAir<F: Field>: BaseAir<F> {
+    /// A unique identifier for this AIR as part of a machine.
     fn name(&self) -> String;
 
+    /// Generate the trace for a given execution record.
+    ///
+    /// The mutable borrow of `record` allows a `MachineAir` to store additional information in the
+    /// record, such as inserting events for other AIRs to process.
     fn generate_trace(&self, record: &mut ExecutionRecord) -> RowMajorMatrix<F>;
 
     fn shard(&self, input: &ExecutionRecord, outputs: &mut Vec<ExecutionRecord>);
 
     fn include(&self, record: &ExecutionRecord) -> bool;
+
+    /// The number of preprocessed columns in the trace.
+    fn preprocessed_width(&self) -> usize {
+        0
+    }
+
+    #[allow(unused_variables)]
+    fn preprocessed_trace(&self, program: &Program) -> Option<RowMajorMatrix<F>> {
+        None
+    }
 }

core/src/bytes/trace.rs

@@ -2,7 +2,10 @@ use p3_field::Field;
 use p3_matrix::dense::RowMajorMatrix;
 
 use super::{air::BYTE_MULT_INDICES, ByteChip};
-use crate::{air::MachineAir, runtime::ExecutionRecord};
+use crate::{
+    air::MachineAir,
+    runtime::{ExecutionRecord, Program},
+};
 
 pub const NUM_ROWS: usize = 1 << 16;
 
@@ -33,4 +36,16 @@ impl<F: Field> MachineAir<F> for ByteChip {
 
         trace
     }
+
+    fn preprocessed_width(&self) -> usize {
+        10
+    }
+
+    fn preprocessed_trace(&self, _program: &Program) -> Option<RowMajorMatrix<F>> {
+        let values = (0..10 * NUM_ROWS)
+            .map(|i| F::from_canonical_usize(i))
+            .collect();
+
+        Some(RowMajorMatrix::new(values, 10))
+    }
 }

core/src/lib.rs

@@ -37,7 +37,7 @@ use p3_matrix::dense::RowMajorMatrix;
 use runtime::{Program, Runtime};
 use serde::de::DeserializeOwned;
 use serde::Serialize;
-use stark::{MainData, OpeningProof, ProgramVerificationError, Proof};
+use stark::{OpeningProof, ProgramVerificationError, Proof, ShardMainData};
 use stark::{RiscvStark, StarkGenericConfig};
 use std::fs;
 use utils::{prove_core, BabyBearBlake3, StarkUtils};
@@ -96,7 +96,7 @@ impl CurtaProver {
         OpeningProof<SC>: Send + Sync,
         <SC::Pcs as Pcs<SC::Val, RowMajorMatrix<SC::Val>>>::Commitment: Send + Sync,
         <SC::Pcs as Pcs<SC::Val, RowMajorMatrix<SC::Val>>>::ProverData: Send + Sync,
-        MainData<SC>: Serialize + DeserializeOwned,
+        ShardMainData<SC>: Serialize + DeserializeOwned,
         <SC as StarkGenericConfig>::Val: p3_field::PrimeField32,
     {
         let program = Program::from(elf);
@@ -121,8 +121,9 @@ impl CurtaVerifier {
     ) -> Result<(), ProgramVerificationError> {
         let config = BabyBearBlake3::new();
         let mut challenger = config.challenger();
-        let (machine, prover_data) = RiscvStark::init(config);
-        machine.verify(&mut challenger, &proof.proof)
+        let machine = RiscvStark::new(config);
+        let (_, vk) = machine.setup(&Program::from(elf));
+        machine.verify(&vk, &proof.proof, &mut challenger)
     }
 
     /// Verify a proof generated by `CurtaProver` with a custom config.
@@ -138,12 +139,14 @@ impl CurtaVerifier {
         OpeningProof<SC>: Send + Sync,
         <SC::Pcs as Pcs<SC::Val, RowMajorMatrix<SC::Val>>>::Commitment: Send + Sync,
         <SC::Pcs as Pcs<SC::Val, RowMajorMatrix<SC::Val>>>::ProverData: Send + Sync,
-        MainData<SC>: Serialize + DeserializeOwned,
+        ShardMainData<SC>: Serialize + DeserializeOwned,
         <SC as StarkGenericConfig>::Val: p3_field::PrimeField32,
     {
         let mut challenger = config.challenger();
-        let (machine, prover_data) = RiscvStark::init(config);
-        machine.verify(&mut challenger, &proof.proof)
+        let machine = RiscvStark::new(config);
+
+        let (_, vk) = machine.setup(&Program::from(elf));
+        machine.verify(&vk, &proof.proof, &mut challenger)
     }
 }
 

core/src/memory/global.rs

@@ -219,7 +219,7 @@ mod tests {
         let mut runtime = Runtime::new(program);
         runtime.run();
 
-        let (machine, _prover_data) = RiscvStark::init(BabyBearPoseidon2::new());
+        let machine = RiscvStark::new(BabyBearPoseidon2::new());
         debug_interactions_with_all_chips(
             &machine.chips(),
             &runtime.record,
@@ -234,7 +234,7 @@ mod tests {
         let mut runtime = Runtime::new(program);
         runtime.run();
 
-        let (machine, _prover_data) = RiscvStark::init(BabyBearPoseidon2::new());
+        let machine = RiscvStark::new(BabyBearPoseidon2::new());
         debug_interactions_with_all_chips(
             &machine.chips(),
             &runtime.record,

core/src/runtime/record.rs

@@ -28,6 +28,9 @@ pub struct ExecutionRecord {
     /// A trace of the CPU events which get emitted during execution.
     pub cpu_events: Vec<CpuEvent>,
 
+    /// Multiplicity counts for each instruction in the program.
+    pub instruction_counts: HashMap<u32, usize>,
+
     /// A trace of the ADD, and ADDI events.
     pub add_events: Vec<AluEvent>,
 

core/src/stark/chip.rs

@@ -6,7 +6,7 @@ use p3_util::log2_ceil_usize;
 use crate::{
     air::{CurtaAirBuilder, MachineAir, MultiTableAirBuilder},
     lookup::{Interaction, InteractionBuilder},
-    runtime::ExecutionRecord,
+    runtime::{ExecutionRecord, Program},
 };
 
 use super::{
@@ -156,6 +156,14 @@ where
     fn include(&self, record: &ExecutionRecord) -> bool {
         self.air.include(record)
     }
+
+    fn preprocessed_trace(&self, program: &Program) -> Option<RowMajorMatrix<F>> {
+        <A as MachineAir<F>>::preprocessed_trace(&self.air, program)
+    }
+
+    fn preprocessed_width(&self) -> usize {
+        self.air.preprocessed_width()
+    }
 }
 
 // Implement AIR directly on Chip, evaluating both execution and permutation constraints.
@@ -201,6 +209,14 @@ impl<'a, SC: StarkGenericConfig> MachineAir<SC::Val> for ChipRef<'a, SC> {
     fn include(&self, record: &ExecutionRecord) -> bool {
         <dyn StarkAir<SC> as MachineAir<SC::Val>>::include(self.air, record)
     }
+
+    fn preprocessed_trace(&self, program: &Program) -> Option<RowMajorMatrix<SC::Val>> {
+        <dyn StarkAir<SC> as MachineAir<SC::Val>>::preprocessed_trace(self.air, program)
+    }
+
+    fn preprocessed_width(&self) -> usize {
+        <dyn StarkAir<SC> as MachineAir<SC::Val>>::preprocessed_width(self.air)
+    }
 }
 
 impl<'a, 'b, SC: StarkGenericConfig> Air<ProverConstraintFolder<'b, SC>> for ChipRef<'a, SC> {