Add MidenVM constraints to plonky3 backend

Cargo.toml

@@ -9,6 +9,7 @@ members = [
     "codegen/winterfell",
     "codegen/ace",
     "codegen/plonky3",
+    "constraints",
 ]
 resolver = "2"
 

air-script/Cargo.toml

@@ -28,10 +28,11 @@ log = { version = "0.4", default-features = false }
 miden-diagnostics = { workspace = true }
 mir = { package = "air-mir", path = "../mir", version = "0.5" }
 
-# 0xMiden Plonky3 Fork
-p3-matrix = { package = "p3-matrix", git="https://github.com/0xMiden/Plonky3", rev = "632fb26a9141cc70f903e5378035ca8d5855c490", default-features = false }
-p3-field = { package = "p3-field", git="https://github.com/0xMiden/Plonky3", rev = "632fb26a9141cc70f903e5378035ca8d5855c490", default-features = false }
-p3-miden-air = { package = "miden-air", git="https://github.com/0xMiden/Plonky3", rev = "632fb26a9141cc70f903e5378035ca8d5855c490", default-features = false }
+# Upstream Plonky3 dependencies
+p3-field = { version = "0.4.2", default-features = false }
+p3-matrix = { version = "0.4.2", default-features = false }
+# Internal p3-miden crates
+p3-miden-air = { package = "p3-miden-air", git = "https://github.com/0xMiden/p3-miden", rev = "b6cb824d2a462bbcbbfb828c158de220aa7ca11f", default-features = false }
 
 # MassaLabs fork
 miden-processor = { package = "miden-processor", git="https://github.com/massalabs/miden-vm", rev = "bc553af69a2543a0789830e8508b019694528181", default-features = false }
@@ -40,17 +41,32 @@ miden-air = { package = "miden-air", git="https://github.com/massalabs/miden-vm"
 [dev-dependencies]
 expect-test = "1.4"
 
-# 0xMiden Plonky3 Fork
-p3-air = { package = "p3-air", git="https://github.com/0xMiden/Plonky3", rev = "632fb26a9141cc70f903e5378035ca8d5855c490", default-features = false }
-p3-challenger = { package = "p3-challenger", git="https://github.com/0xMiden/Plonky3", rev = "632fb26a9141cc70f903e5378035ca8d5855c490", default-features = false }
-p3-commit = { package = "p3-commit", git="https://github.com/0xMiden/Plonky3", rev = "632fb26a9141cc70f903e5378035ca8d5855c490", default-features = false }
-p3-dft = { package = "p3-dft", git="https://github.com/0xMiden/Plonky3", rev = "632fb26a9141cc70f903e5378035ca8d5855c490", default-features = false }
-p3-fri = { package = "p3-fri", git="https://github.com/0xMiden/Plonky3", rev = "632fb26a9141cc70f903e5378035ca8d5855c490", default-features = false }
-p3-merkle-tree = { package = "p3-merkle-tree", git="https://github.com/0xMiden/Plonky3", rev = "632fb26a9141cc70f903e5378035ca8d5855c490", default-features = false }
-p3-miden-prover = { package = "miden-prover", git="https://github.com/0xMiden/Plonky3", rev = "632fb26a9141cc70f903e5378035ca8d5855c490", default-features = false }
-p3-goldilocks = { package = "p3-goldilocks", git="https://github.com/0xMiden/Plonky3", rev = "632fb26a9141cc70f903e5378035ca8d5855c490", default-features = false }
-p3-sha256 = { package = "p3-sha256", git="https://github.com/0xMiden/Plonky3", rev = "632fb26a9141cc70f903e5378035ca8d5855c490", default-features = false }
-p3-symmetric = { package = "p3-symmetric", git="https://github.com/0xMiden/Plonky3", rev = "632fb26a9141cc70f903e5378035ca8d5855c490", default-features = false }
+# Upstream Plonky3 dependencies
+p3-air = { version = "0.4.2", default-features = false }
+p3-baby-bear = { version = "0.4.2", default-features = false }
+p3-challenger = { version = "0.4.2", default-features = false }
+p3-circle = { version = "0.4.2", default-features = false }
+p3-commit = { version = "0.4.2", default-features = false }
+p3-dft = { version = "0.4.2", default-features = false }
+p3-field-testing = { version = "0.4.2", default-features = false }
+p3-goldilocks = { version = "0.4.2", default-features = false }
+p3-interpolation = { version = "0.4.2", default-features = false }
+p3-keccak = { version = "0.4.2", default-features = false }
+p3-matrix = { version = "0.4.2", default-features = false }
+p3-maybe-rayon = { version = "0.4.2", default-features = false }
+p3-mds = { version = "0.4.2", default-features = false }
+p3-merkle-tree = { version = "0.4.2", default-features = false }
+p3-mersenne-31 = { version = "0.4.2", default-features = false }
+p3-poseidon = { version = "0.4.2", default-features = false }
+p3-poseidon2 = { version = "0.4.2", default-features = false }
+p3-sha256 = { version = "0.4.2", default-features = false }
+p3-symmetric = { version = "0.4.2", default-features = false }
+p3-uni-stark = { version = "0.4.2", default-features = false }
+p3-util = { version = "0.4.2", default-features = false }
+# Internal p3-miden crates
+p3-miden-fri = { package = "p3-miden-fri", git = "https://github.com/0xMiden/p3-miden", rev = "b6cb824d2a462bbcbbfb828c158de220aa7ca11f", default-features = false }
+p3-miden-prover = { package = "p3-miden-prover", git = "https://github.com/0xMiden/p3-miden", rev = "b6cb824d2a462bbcbbfb828c158de220aa7ca11f",  default-features = false }
+p3-miden-uni-stark = { package = "p3-miden-uni-stark", git = "https://github.com/0xMiden/p3-miden", rev = "b6cb824d2a462bbcbbfb828c158de220aa7ca11f",  default-features = false }
 
 winter-air = { package = "winter-air", version = "0.12", default-features = false }
 winter-math = { package = "winter-math", version = "0.12", default-features = false }

air-script/src/test_utils/air_tester_macros.rs

@@ -52,7 +52,7 @@ macro_rules! generate_air_plonky3_test_with_airscript_traits {
                 p3_challenger::HashChallenger<u8, ByteHash, 32>,
             >;
             type Dft = p3_dft::Radix2DitParallel<Val>;
-            type Pcs = p3_fri::TwoAdicFriPcs<Val, Dft, ValMmcs, ChallengeMmcs>;
+            type Pcs = p3_miden_fri::TwoAdicFriPcs<Val, Dft, ValMmcs, ChallengeMmcs>;
             type MyConfig = p3_miden_prover::StarkConfig<Pcs, Challenge, Challenger>;
 
             let byte_hash = ByteHash {};
@@ -62,7 +62,7 @@ macro_rules! generate_air_plonky3_test_with_airscript_traits {
             let challenge_mmcs = ChallengeMmcs::new(val_mmcs.clone());
             let challenger = Challenger::from_hasher(vec![], byte_hash);
             let dft = Dft::default();
-            let mut fri_params = p3_fri::create_recursive_miden_fri_params(challenge_mmcs);
+            let mut fri_params = p3_miden_fri::create_recursive_miden_fri_params(challenge_mmcs);
             let pcs = Pcs::new(dft, val_mmcs, fri_params);
             let config = MyConfig::new(pcs, challenger);
 

air-script/src/tests/buses/buses_complex_plonky3.rs

@@ -29,8 +29,8 @@ where F: Field,
         AUX_WIDTH
     }
 
-    fn bus_types(&self) -> Vec<BusType> {
-        vec![
+    fn bus_types(&self) -> &[BusType] {
+        &[
             BusType::Multiset,
             BusType::Logup,
         ]

air-script/src/tests/buses/buses_simple_plonky3.rs

@@ -29,8 +29,8 @@ where F: Field,
         AUX_WIDTH
     }
 
-    fn bus_types(&self) -> Vec<BusType> {
-        vec![
+    fn bus_types(&self) -> &[BusType] {
+        &[
             BusType::Multiset,
             BusType::Logup,
         ]

air-script/src/tests/buses/buses_varlen_boundary_both_plonky3.rs

@@ -29,8 +29,8 @@ where F: Field,
         AUX_WIDTH
     }
 
-    fn bus_types(&self) -> Vec<BusType> {
-        vec![
+    fn bus_types(&self) -> &[BusType] {
+        &[
             BusType::Multiset,
             BusType::Logup,
         ]

air-script/src/tests/buses/buses_varlen_boundary_first_plonky3.rs

@@ -29,8 +29,8 @@ where F: Field,
         AUX_WIDTH
     }
 
-    fn bus_types(&self) -> Vec<BusType> {
-        vec![
+    fn bus_types(&self) -> &[BusType] {
+        &[
             BusType::Multiset,
             BusType::Logup,
         ]

air-script/src/tests/buses/buses_varlen_boundary_last_plonky3.rs

@@ -29,8 +29,8 @@ where F: Field,
         AUX_WIDTH
     }
 
-    fn bus_types(&self) -> Vec<BusType> {
-        vec![
+    fn bus_types(&self) -> &[BusType] {
+        &[
             BusType::Multiset,
             BusType::Logup,
         ]

air-script/src/tests/comprehension_periodic_binding/comprehension_periodic_binding.air

@@ -0,0 +1,19 @@
+def ComprehensionPeriodicBindingTest
+
+use lib::test_comprehension;
+
+trace_columns {
+    main: [a, b],
+}
+
+public_inputs {
+    stack_inputs: [1],
+}
+
+boundary_constraints {
+    enf a.first = 0;
+}
+
+integrity_constraints {
+    enf test_comprehension([a, b]);
+}

air-script/src/tests/comprehension_periodic_binding/comprehension_periodic_binding.rs

@@ -0,0 +1,97 @@
+use winter_air::{Air, AirContext, Assertion, AuxRandElements, EvaluationFrame, ProofOptions as WinterProofOptions, TransitionConstraintDegree, TraceInfo};
+use winter_math::fields::f64::BaseElement as Felt;
+use winter_math::{ExtensionOf, FieldElement, ToElements};
+use winter_utils::{ByteWriter, Serializable};
+
+pub struct PublicInputs {
+    stack_inputs: [Felt; 1],
+}
+
+impl PublicInputs {
+    pub fn new(stack_inputs: [Felt; 1]) -> Self {
+        Self { stack_inputs }
+    }
+}
+
+impl Serializable for PublicInputs {
+    fn write_into<W: ByteWriter>(&self, target: &mut W) {
+        self.stack_inputs.write_into(target);
+    }
+}
+
+impl ToElements<Felt> for PublicInputs {
+    fn to_elements(&self) -> Vec<Felt> {
+        let mut elements = Vec::new();
+        elements.extend_from_slice(&self.stack_inputs);
+        elements
+    }
+}
+
+pub struct ComprehensionPeriodicBindingTest {
+    context: AirContext<Felt>,
+    stack_inputs: [Felt; 1],
+}
+
+impl ComprehensionPeriodicBindingTest {
+    pub fn last_step(&self) -> usize {
+        self.trace_length() - self.context().num_transition_exemptions()
+    }
+}
+
+impl Air for ComprehensionPeriodicBindingTest {
+    type BaseField = Felt;
+    type PublicInputs = PublicInputs;
+
+    fn context(&self) -> &AirContext<Felt> {
+        &self.context
+    }
+
+    fn new(trace_info: TraceInfo, public_inputs: PublicInputs, options: WinterProofOptions) -> Self {
+        let main_degrees = vec![TransitionConstraintDegree::with_cycles(1, vec![2, 2])];
+        let aux_degrees = vec![];
+        let num_main_assertions = 1;
+        let num_aux_assertions = 0;
+
+        let context = AirContext::new_multi_segment(
+            trace_info,
+            main_degrees,
+            aux_degrees,
+            num_main_assertions,
+            num_aux_assertions,
+            options,
+        )
+        .set_num_transition_exemptions(2);
+        Self { context, stack_inputs: public_inputs.stack_inputs }
+    }
+
+    fn get_periodic_column_values(&self) -> Vec<Vec<Felt>> {
+        vec![vec![Felt::ONE, Felt::new(2)], vec![Felt::new(3), Felt::new(4)]]
+    }
+
+    fn get_assertions(&self) -> Vec<Assertion<Felt>> {
+        let mut result = Vec::new();
+        result.push(Assertion::single(0, 0, Felt::ZERO));
+        result
+    }
+
+    fn get_aux_assertions<E: FieldElement<BaseField = Felt>>(&self, aux_rand_elements: &AuxRandElements<E>) -> Vec<Assertion<E>> {
+        let mut result = Vec::new();
+        result
+    }
+
+    fn evaluate_transition<E: FieldElement<BaseField = Felt>>(&self, frame: &EvaluationFrame<E>, periodic_values: &[E], result: &mut [E]) {
+        let main_current = frame.current();
+        let main_next = frame.next();
+        result[0] = main_next[0] - (main_current[0] * periodic_values[0] + main_current[1] * periodic_values[1]);
+    }
+
+    fn evaluate_aux_transition<F, E>(&self, main_frame: &EvaluationFrame<F>, aux_frame: &EvaluationFrame<E>, _periodic_values: &[F], aux_rand_elements: &AuxRandElements<E>, result: &mut [E])
+    where F: FieldElement<BaseField = Felt>,
+          E: FieldElement<BaseField = Felt> + ExtensionOf<F>,
+    {
+        let main_current = main_frame.current();
+        let main_next = main_frame.next();
+        let aux_current = aux_frame.current();
+        let aux_next = aux_frame.next();
+    }
+}

air-script/src/tests/comprehension_periodic_binding/comprehension_periodic_binding_plonky3.rs

@@ -0,0 +1,56 @@
+use p3_field::{ExtensionField, Field, PrimeCharacteristicRing};
+use p3_matrix::Matrix;
+use p3_matrix::dense::RowMajorMatrixView;
+use p3_matrix::stack::VerticalPair;
+use p3_miden_air::{BusType, MidenAir, MidenAirBuilder, RowMajorMatrix};
+
+pub const MAIN_WIDTH: usize = 2;
+pub const AUX_WIDTH: usize = 0;
+pub const NUM_PERIODIC_VALUES: usize = 2;
+pub const PERIOD: usize = 2;
+pub const NUM_PUBLIC_VALUES: usize = 1;
+pub const MAX_BETA_CHALLENGE_POWER: usize = 0;
+
+pub struct ComprehensionPeriodicBindingTest;
+
+impl<F, EF> MidenAir<F, EF> for ComprehensionPeriodicBindingTest
+where F: Field,
+      EF: ExtensionField<F>,
+{
+    fn width(&self) -> usize {
+        MAIN_WIDTH
+    }
+
+    fn num_public_values(&self) -> usize {
+        NUM_PUBLIC_VALUES
+    }
+
+    fn periodic_table(&self) -> Vec<Vec<F>> {
+        vec![
+            vec![F::from_u64(1), F::from_u64(2)],
+            vec![F::from_u64(3), F::from_u64(4)],
+        ]
+    }
+
+    fn eval<AB>(&self, builder: &mut AB)
+    where AB: MidenAirBuilder<F = F>,
+    {
+        let public_values: [_; NUM_PUBLIC_VALUES] = builder.public_values().try_into().expect("Wrong number of public values");
+        let periodic_values: [_; NUM_PERIODIC_VALUES] = builder.periodic_evals().try_into().expect("Wrong number of periodic values");
+        // Note: for now, we do not have any preprocessed values
+        // let preprocessed = builder.preprocessed();
+        let main = builder.main();
+        let (main_current, main_next) = (
+            main.row_slice(0).unwrap(),
+            main.row_slice(1).unwrap(),
+        );
+
+        // Main boundary constraints
+        builder.when_first_row().assert_zero(main_current[0].clone().into());
+
+        // Main integrity/transition constraints
+        builder.when_transition().assert_zero_ext(AB::ExprEF::from(main_next[0].clone().into()) - (AB::ExprEF::from(main_current[0].clone().into()) * AB::ExprEF::from(periodic_values[0].clone().into()) + AB::ExprEF::from(main_current[1].clone().into()) * AB::ExprEF::from(periodic_values[1].clone().into())));
+
+        // Aux integrity/transition constraints
+    }
+}

air-script/src/tests/comprehension_periodic_binding/lib.air

@@ -0,0 +1,16 @@
+mod lib
+
+periodic_columns {
+    k0: [1, 2],
+    k1: [3, 4],
+}
+
+ev test_comprehension([a, b]) {
+    # Create local variable holding periodic column references
+    let cols = [k0, k1];
+    let vals = [a, b];
+
+    # Iterate over the local variable - binding 'k' gets typed as PeriodicColumn
+    # but it's actually a local variable holding a value
+    enf a' = sum([x * k for (x, k) in (vals, cols)]);
+}

air-script/src/tests/comprehension_periodic_binding/mod.rs

@@ -0,0 +1,3 @@
+#[rustfmt::skip]
+#[allow(clippy::all)]
+mod comprehension_periodic_binding;

air-script/src/tests/cross_module_constants/constants_lib.air

@@ -0,0 +1,21 @@
+mod constants_lib
+
+# Constants used in comprehension iterables
+const WEIGHTS = [1, 2, 3, 4];
+
+# Pure function that uses constants in a comprehension
+fn weighted_sum(values: felt[4]) -> felt {
+    return sum([v * w for (v, w) in (values, WEIGHTS)]);
+}
+
+# Another function that calls the first
+fn compute_result(a: felt, b: felt, c: felt, d: felt) -> felt {
+    let values = [a, b, c, d];
+    return weighted_sum(values);
+}
+
+# Evaluator that uses the chain of functions
+ev apply_computation([cols[5]]) {
+    let result = compute_result(cols[0], cols[1], cols[2], cols[3]);
+    enf cols[4] = result;
+}