feat: Preprocessing + recursion (#450)

Cargo.toml

@@ -49,7 +49,7 @@ p3-maybe-rayon = { git = "https://github.com/Plonky3/Plonky3.git", branch = "sp1
 p3-bn254-fr = { git = "https://github.com/Plonky3/Plonky3.git", branch = "sp1" }
 
 # For local development.
-#
+
 # p3-air = { path = "../Plonky3/air" }
 # p3-field = { path = "../Plonky3/field" }
 # p3-commit = { path = "../Plonky3/commit" }
@@ -69,4 +69,4 @@ p3-bn254-fr = { git = "https://github.com/Plonky3/Plonky3.git", branch = "sp1" }
 # p3-symmetric = { path = "../Plonky3/symmetric" }
 # p3-uni-stark = { path = "../Plonky3/uni-stark" }
 # p3-maybe-rayon = { path = "../Plonky3/maybe-rayon" }
-# p3-bn254-fr = { path = "../Plonky3/bn254-fr" }
+# p3-bn254-fr = { path = "../Plonky3/bn254-fr" }

core/src/air/builder.rs

@@ -624,13 +624,13 @@ impl<'a, AB: AirBuilder + MessageBuilder<M>, M> MessageBuilder<M> for FilteredAi
 }
 
 impl<AB: AirBuilder + MessageBuilder<AirInteraction<AB::Expr>>> BaseAirBuilder for AB {}
-impl<AB: AirBuilder + MessageBuilder<AirInteraction<AB::Expr>>> ByteAirBuilder for AB {}
-impl<AB: AirBuilder + MessageBuilder<AirInteraction<AB::Expr>>> WordAirBuilder for AB {}
-impl<AB: AirBuilder + MessageBuilder<AirInteraction<AB::Expr>>> AluAirBuilder for AB {}
-impl<AB: AirBuilder + MessageBuilder<AirInteraction<AB::Expr>>> MemoryAirBuilder for AB {}
-impl<AB: AirBuilder + MessageBuilder<AirInteraction<AB::Expr>>> ProgramAirBuilder for AB {}
-impl<AB: AirBuilder + MessageBuilder<AirInteraction<AB::Expr>>> ExtensionAirBuilder for AB {}
-impl<AB: AirBuilder + MessageBuilder<AirInteraction<AB::Expr>>> SP1AirBuilder for AB {}
+impl<AB: BaseAirBuilder> ByteAirBuilder for AB {}
+impl<AB: BaseAirBuilder> WordAirBuilder for AB {}
+impl<AB: BaseAirBuilder> AluAirBuilder for AB {}
+impl<AB: BaseAirBuilder> MemoryAirBuilder for AB {}
+impl<AB: BaseAirBuilder> ProgramAirBuilder for AB {}
+impl<AB: BaseAirBuilder> ExtensionAirBuilder for AB {}
+impl<AB: BaseAirBuilder> SP1AirBuilder for AB {}
 
 impl<'a, SC: StarkGenericConfig> EmptyMessageBuilder for ProverConstraintFolder<'a, SC> {}
 impl<'a, SC: StarkGenericConfig> EmptyMessageBuilder for VerifierConstraintFolder<'a, SC> {}

core/src/air/machine.rs

@@ -2,14 +2,17 @@ use p3_air::BaseAir;
 use p3_field::Field;
 use p3_matrix::dense::RowMajorMatrix;
 
-use crate::{runtime::Program, stark::MachineRecord};
+use crate::stark::MachineRecord;
 
 pub use sp1_derive::MachineAir;
 
 /// An AIR that is part of a Risc-V AIR arithmetization.
 pub trait MachineAir<F: Field>: BaseAir<F> {
+    /// The execution record containing events for producing the air trace.
     type Record: MachineRecord;
 
+    type Program;
+
     /// A unique identifier for this AIR as part of a machine.
     fn name(&self) -> String;
 
@@ -25,15 +28,16 @@ pub trait MachineAir<F: Field>: BaseAir<F> {
         self.generate_trace(input, output);
     }
 
-    /// The number of preprocessed columns in the trace.
+    /// Whether this execution record contains events for this air.
+    fn included(&self, shard: &Self::Record) -> bool;
+
     fn preprocessed_width(&self) -> usize {
         0
     }
 
+    /// Generate the preprocessed trace given a specific program.
     #[allow(unused_variables)]
-    fn generate_preprocessed_trace(&self, program: &Program) -> Option<RowMajorMatrix<F>> {
+    fn generate_preprocessed_trace(&self, program: &Self::Program) -> Option<RowMajorMatrix<F>> {
         None
     }
-
-    fn included(&self, shard: &Self::Record) -> bool;
 }

core/src/alu/add_sub/mod.rs

@@ -12,7 +12,7 @@ use tracing::instrument;
 use crate::air::MachineAir;
 use crate::air::{SP1AirBuilder, Word};
 use crate::operations::AddOperation;
-use crate::runtime::{ExecutionRecord, Opcode};
+use crate::runtime::{ExecutionRecord, Opcode, Program};
 use crate::stark::MachineRecord;
 use crate::utils::pad_to_power_of_two;
 
@@ -51,6 +51,8 @@ pub struct AddSubCols<T> {
 impl<F: PrimeField> MachineAir<F> for AddSubChip {
     type Record = ExecutionRecord;
 
+    type Program = Program;
+
     fn name(&self) -> String {
         "AddSub".to_string()
     }

core/src/alu/bitwise/mod.rs

@@ -10,7 +10,7 @@ use tracing::instrument;
 use crate::air::MachineAir;
 use crate::air::{SP1AirBuilder, Word};
 use crate::bytes::{ByteLookupEvent, ByteOpcode};
-use crate::runtime::{ExecutionRecord, Opcode};
+use crate::runtime::{ExecutionRecord, Opcode, Program};
 use crate::utils::pad_to_power_of_two;
 
 /// The number of main trace columns for `BitwiseChip`.
@@ -45,6 +45,8 @@ pub struct BitwiseCols<T> {
 impl<F: PrimeField> MachineAir<F> for BitwiseChip {
     type Record = ExecutionRecord;
 
+    type Program = Program;
+
     fn name(&self) -> String {
         "Bitwise".to_string()
     }

core/src/alu/divrem/mod.rs

@@ -80,7 +80,7 @@ use crate::alu::AluEvent;
 use crate::bytes::{ByteLookupEvent, ByteOpcode};
 use crate::disassembler::WORD_SIZE;
 use crate::operations::{IsEqualWordOperation, IsZeroWordOperation};
-use crate::runtime::{ExecutionRecord, Opcode};
+use crate::runtime::{ExecutionRecord, Opcode, Program};
 use crate::utils::pad_to_power_of_two;
 
 /// The number of main trace columns for `DivRemChip`.
@@ -184,6 +184,8 @@ pub struct DivRemCols<T> {
 impl<F: PrimeField> MachineAir<F> for DivRemChip {
     type Record = ExecutionRecord;
 
+    type Program = Program;
+
     fn name(&self) -> String {
         "DivRem".to_string()
     }

core/src/alu/lt/mod.rs

@@ -12,7 +12,7 @@ use tracing::instrument;
 
 use crate::air::{SP1AirBuilder, Word};
 
-use crate::runtime::{ExecutionRecord, Opcode};
+use crate::runtime::{ExecutionRecord, Opcode, Program};
 use crate::utils::pad_to_power_of_two;
 
 /// The number of main trace columns for `LtChip`.
@@ -76,6 +76,8 @@ impl LtCols<u32> {
 impl<F: PrimeField> MachineAir<F> for LtChip {
     type Record = ExecutionRecord;
 
+    type Program = Program;
+
     fn name(&self) -> String {
         "Lt".to_string()
     }

core/src/alu/mul/mod.rs

@@ -47,7 +47,7 @@ use crate::air::{SP1AirBuilder, Word};
 use crate::alu::mul::utils::get_msb;
 use crate::bytes::{ByteLookupEvent, ByteOpcode};
 use crate::disassembler::WORD_SIZE;
-use crate::runtime::{ExecutionRecord, Opcode};
+use crate::runtime::{ExecutionRecord, Opcode, Program};
 use crate::stark::MachineRecord;
 use crate::utils::pad_to_power_of_two;
 
@@ -118,6 +118,8 @@ pub struct MulCols<T> {
 impl<F: PrimeField> MachineAir<F> for MulChip {
     type Record = ExecutionRecord;
 
+    type Program = Program;
+
     fn name(&self) -> String {
         "Mul".to_string()
     }

core/src/alu/sll/mod.rs

@@ -43,7 +43,7 @@ use tracing::instrument;
 use crate::air::MachineAir;
 use crate::air::{SP1AirBuilder, Word};
 use crate::disassembler::WORD_SIZE;
-use crate::runtime::{ExecutionRecord, Opcode};
+use crate::runtime::{ExecutionRecord, Opcode, Program};
 use crate::utils::pad_to_power_of_two;
 
 /// The number of main trace columns for `ShiftLeft`.
@@ -93,6 +93,8 @@ pub struct ShiftLeftCols<T> {
 impl<F: PrimeField> MachineAir<F> for ShiftLeft {
     type Record = ExecutionRecord;
 
+    type Program = Program;
+
     fn name(&self) -> String {
         "ShiftLeft".to_string()
     }

core/src/alu/sr/mod.rs

@@ -59,7 +59,7 @@ use crate::alu::sr::utils::{nb_bits_to_shift, nb_bytes_to_shift};
 use crate::bytes::utils::shr_carry;
 use crate::bytes::{ByteLookupEvent, ByteOpcode};
 use crate::disassembler::WORD_SIZE;
-use crate::runtime::{ExecutionRecord, Opcode};
+use crate::runtime::{ExecutionRecord, Opcode, Program};
 use crate::utils::pad_to_power_of_two;
 
 /// The number of main trace columns for `ShiftRightChip`.
@@ -125,6 +125,8 @@ pub struct ShiftRightCols<T> {
 impl<F: PrimeField> MachineAir<F> for ShiftRightChip {
     type Record = ExecutionRecord;
 
+    type Program = Program;
+
     fn name(&self) -> String {
         "ShiftRight".to_string()
     }