Target 0xMiden Plonky3 repo and use AirScriptAir and AirScriptBuilder traits

.github/workflows/book.yml

@@ -42,7 +42,7 @@ jobs:
       - name: Install mdbook and plugins
         uses: taiki-e/install-action@v2
         with:
-          tool: mdbook, mdbook-linkcheck, mdbook-alerts, mdbook-katex, mdbook-mermaid
+          tool: mdbook@0.4.48, mdbook-linkcheck, mdbook-alerts, mdbook-katex, mdbook-mermaid
 
       - name: Build book
         run: mdbook build docs/

CHANGELOG.md

@@ -13,6 +13,8 @@
 - Fix regressions on MIR and list_comprehensions (#449).
 - Add Plonky3 codegen backend (#461).
 - Fixed a vector unrolling issue in nested match evaluations (#491).
+- In Plonky3 codegen, use AirScriptAir and AirScriptBuilder traits, and generate aux constraints (#508).
+- In Plonky3 codegen, use MidenAir and MidenAirBuilder from 0xMiden's Plonky3 fork instead of AirScriptAir and AirScriptBuilder (#515).
 
 ## 0.4.0 (2025-06-20)
 

air-script/Cargo.toml

@@ -30,18 +30,26 @@ mir = { package = "air-mir", path = "../mir", version = "0.5" }
 
 [dev-dependencies]
 expect-test = "1.4"
-p3-air = { package = "p3-air", version = "0.3", default-features = false }
-p3-challenger = { package = "p3-challenger", version = "0.3", default-features = false }
-p3-circle = { package = "p3-circle", version = "0.3", default-features = false }
-p3-commit = { package = "p3-commit", version = "0.3", default-features = false }
-p3-field = { package = "p3-field", version = "0.3", default-features = false }
-p3-fri = { package = "p3-fri", version = "0.3", default-features = false }
-p3-matrix = { package = "p3-matrix", version = "0.3", default-features = false }
-p3-merkle-tree = { package = "p3-merkle-tree", version = "0.3", default-features = false }
-p3-goldilocks = { package = "p3-goldilocks", version = "0.3", default-features = false }
-p3-sha256 = { package = "p3-sha256", version = "0.3", default-features = false }
-p3-symmetric = { package = "p3-symmetric", version = "0.3", default-features = false }
-p3-uni-stark = { package = "p3-uni-stark", version = "0.3.0", default-features = false }
+
+# 0xMiden Plonky3 Fork
+p3-air = { package = "p3-air", git="https://github.com/0xMiden/Plonky3", rev = "5352d47ee57bd1b5d6008a90b64cbb7fe00aa17e", default-features = false }
+p3-challenger = { package = "p3-challenger", git="https://github.com/0xMiden/Plonky3", rev = "5352d47ee57bd1b5d6008a90b64cbb7fe00aa17e", default-features = false }
+p3-commit = { package = "p3-commit", git="https://github.com/0xMiden/Plonky3", rev = "5352d47ee57bd1b5d6008a90b64cbb7fe00aa17e", default-features = false }
+p3-dft = { package = "p3-dft", git="https://github.com/0xMiden/Plonky3", rev = "5352d47ee57bd1b5d6008a90b64cbb7fe00aa17e", default-features = false }
+p3-field = { package = "p3-field", git="https://github.com/0xMiden/Plonky3", rev = "5352d47ee57bd1b5d6008a90b64cbb7fe00aa17e", default-features = false }
+p3-fri = { package = "p3-fri", git="https://github.com/0xMiden/Plonky3", rev = "5352d47ee57bd1b5d6008a90b64cbb7fe00aa17e", default-features = false }
+p3-matrix = { package = "p3-matrix", git="https://github.com/0xMiden/Plonky3", rev = "5352d47ee57bd1b5d6008a90b64cbb7fe00aa17e", default-features = false }
+p3-merkle-tree = { package = "p3-merkle-tree", git="https://github.com/0xMiden/Plonky3", rev = "5352d47ee57bd1b5d6008a90b64cbb7fe00aa17e", default-features = false }
+p3-miden-air = { package = "miden-air", git="https://github.com/0xMiden/Plonky3", rev = "5352d47ee57bd1b5d6008a90b64cbb7fe00aa17e", default-features = false }
+p3-goldilocks = { package = "p3-goldilocks", git="https://github.com/0xMiden/Plonky3", rev = "5352d47ee57bd1b5d6008a90b64cbb7fe00aa17e", default-features = false }
+p3-sha256 = { package = "p3-sha256", git="https://github.com/0xMiden/Plonky3", rev = "5352d47ee57bd1b5d6008a90b64cbb7fe00aa17e", default-features = false }
+p3-symmetric = { package = "p3-symmetric", git="https://github.com/0xMiden/Plonky3", rev = "5352d47ee57bd1b5d6008a90b64cbb7fe00aa17e", default-features = false }
+p3-uni-stark = { package = "p3-uni-stark", git="https://github.com/0xMiden/Plonky3", rev = "5352d47ee57bd1b5d6008a90b64cbb7fe00aa17e", default-features = false }
+
+# MassaLabs fork
+miden-processor = { package = "miden-processor", git="https://github.com/massalabs/miden-vm", rev = "bc553af69a2543a0789830e8508b019694528181", default-features = false }
+miden-air = { package = "miden-air", git="https://github.com/massalabs/miden-vm", rev = "bc553af69a2543a0789830e8508b019694528181", 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 }
 winter-utils = { package = "winter-utils", version = "0.12", default-features = false }

air-script/src/lib.rs

@@ -1,3 +1,8 @@
 pub use air_codegen_winter::CodeGenerator as WinterfellCodeGenerator;
 pub use air_ir::{Air, CompileError, compile};
 pub use air_parser::{parse, parse_file, transforms};
+
+#[cfg(test)]
+pub mod test_utils;
+#[cfg(test)]
+mod tests;

air-script/tests/helpers/macros.rs → air-script/src/test_utils/air_tester_macros.rs

@@ -36,49 +36,49 @@ macro_rules! generate_air_winterfell_test {
 /// * `test_name` - The identifier for the test function (e.g., `test_binary_air`)
 /// * `air_name` - The identifier for the AIR struct (e.g., `BinaryAir`)
 #[macro_export]
-macro_rules! generate_air_plonky3_test {
+macro_rules! generate_air_plonky3_test_with_airscript_traits {
     ($test_name:ident, $air_name:ident) => {
         #[test]
         fn $test_name() {
-            type Val = Goldilocks;
-            type Challenge = BinomialExtensionField<Val, 3>;
-
-            type ByteHash = Sha256;
-            type FieldHash = SerializingHasher<ByteHash>;
-            type MyCompress = CompressionFunctionFromHasher<ByteHash, 2, 32>;
-            type ValMmcs = MerkleTreeMmcs<Val, u8, FieldHash, MyCompress, 32>;
-            type ChallengeMmcs = ExtensionMmcs<Val, Challenge, ValMmcs>;
-            type Challenger = SerializingChallenger64<Val, HashChallenger<u8, ByteHash, 32>>;
-            type Pcs = CirclePcs<Val, ValMmcs, ChallengeMmcs>;
-            type MyConfig = StarkConfig<Pcs, Challenge, Challenger>;
+            type Val = p3_goldilocks::Goldilocks;
+            type Challenge = p3_field::extension::BinomialExtensionField<Val, 2>;
+            type ByteHash = p3_sha256::Sha256;
+            type FieldHash = p3_symmetric::SerializingHasher<ByteHash>;
+            type MyCompress = p3_symmetric::CompressionFunctionFromHasher<ByteHash, 2, 32>;
+            type ValMmcs = p3_merkle_tree::MerkleTreeMmcs<Val, u8, FieldHash, MyCompress, 32>;
+            type ChallengeMmcs = p3_commit::ExtensionMmcs<Val, Challenge, ValMmcs>;
+            type Challenger = p3_challenger::SerializingChallenger64<
+                Val,
+                p3_challenger::HashChallenger<u8, ByteHash, 32>,
+            >;
+            type Dft = p3_dft::Radix2DitParallel<Val>;
+            type Pcs = p3_fri::TwoAdicFriPcs<Val, Dft, ValMmcs, ChallengeMmcs>;
+            type MyConfig = p3_uni_stark::StarkConfig<Pcs, Challenge, Challenger>;
 
             let byte_hash = ByteHash {};
-            let field_hash = FieldHash::new(Sha256);
+            let field_hash = FieldHash::new(p3_sha256::Sha256);
             let compress = MyCompress::new(byte_hash);
             let val_mmcs = ValMmcs::new(field_hash, compress);
             let challenge_mmcs = ChallengeMmcs::new(val_mmcs.clone());
             let challenger = Challenger::from_hasher(vec![], byte_hash);
-            let fri_params = create_benchmark_fri_params(challenge_mmcs);
-            let pcs = Pcs {
-                mmcs: val_mmcs,
-                fri_params,
-                _phantom: PhantomData,
-            };
+            let dft = Dft::default();
+            let fri_params = p3_fri::create_benchmark_fri_params(challenge_mmcs);
+            let pcs = Pcs::new(dft, val_mmcs, fri_params);
             let config = MyConfig::new(pcs, challenger);
 
             let inputs = generate_inputs();
-            let inputs_goldilocks: Vec<Val> =
-                inputs.iter().map(|&x| Val::from_u32(x)).collect();
+            let inputs_goldilocks: Vec<Val> = inputs
+                .iter()
+                .map(|&x| <Val as p3_field::PrimeCharacteristicRing>::from_u32(x))
+                .collect();
 
             let trace = generate_trace_rows::<Val>(inputs);
 
-            check_constraints_with_periodic_columns(&$air_name {}, &trace, &inputs_goldilocks);
-
-            /*let prove_with_periodic_columns = prove_with_periodic_columns(&config, &BitwiseAir {}, trace, &inputs_goldilocks);
-            verify_with_periodic_columns(&config, &BitwiseAir {}, &prove_with_periodic_columns, &inputs_goldilocks).expect("Verification failed");*/
-
-            /*let proof = prove(&config, &BitwiseAir {}, trace, &inputs_goldilocks);
-            verify(&config, &BitwiseAir {}, &proof, &inputs_goldilocks).expect("Verification failed");*/
+            check_constraints_with_airscript_traits::<Val, Challenge, $air_name>(
+                &$air_name {},
+                &trace,
+                &inputs_goldilocks,
+            );
         }
     };
 }

air-script/tests/codegen/helpers.rs → air-script/src/test_utils/codegen.rs

@@ -1,11 +1,12 @@
 use std::sync::Arc;
 
 use air_ir::{CodeGenerator, CompileError};
-use air_script::compile;
 use miden_diagnostics::{
     CodeMap, DefaultEmitter, DiagnosticsHandler, term::termcolor::ColorChoice,
 };
 
+use crate::compile;
+
 pub enum Target {
     Winterfell,
     Plonky3,

air-script/src/test_utils/miden_vm_aux_trace_generator.rs

@@ -0,0 +1,116 @@
+use std::collections::BTreeMap;
+
+use miden_air::{Felt, FieldElement, trace::main_trace::MainTrace};
+use miden_processor::{
+    ColMatrix, Kernel, PrecompileTranscriptState, QuadExtension,
+    chiplets::{AceHints, AuxTraceBuilder as ChipletsAuxTraceBuilder},
+    decoder::AuxTraceBuilder as DecoderAuxTraceBuilder,
+    range::AuxTraceBuilder as RangeAuxTraceBuilder,
+    stack::AuxTraceBuilder as StackAuxTraceBuilder,
+};
+use p3_field::{ExtensionField, Field, PrimeField64};
+use p3_matrix::Matrix;
+use p3_miden_air::RowMajorMatrix;
+
+pub enum MidenModule {
+    Chiplets,
+    Decoder,
+    Stack,
+    Range,
+}
+
+impl MidenModule {
+    fn build_aux_columns<EF: FieldElement<BaseField = Felt>>(
+        &self,
+        main_trace: &MainTrace,
+        rand_elements: &[EF],
+    ) -> Vec<Vec<EF>> {
+        match self {
+            MidenModule::Chiplets => {
+                let kernel = Kernel::new(&[]).unwrap();
+                let ace_hints = AceHints::new(0, vec![]);
+                let final_transcript_state = PrecompileTranscriptState::default();
+                let aux_trace_builder =
+                    ChipletsAuxTraceBuilder::new(kernel, ace_hints, final_transcript_state);
+                aux_trace_builder.build_aux_columns(main_trace, rand_elements).to_vec()
+            },
+            MidenModule::Decoder => {
+                let aux_trace_builder = DecoderAuxTraceBuilder {};
+                aux_trace_builder.build_aux_columns(main_trace, rand_elements)
+            },
+            MidenModule::Stack => {
+                let aux_trace_builder = StackAuxTraceBuilder {};
+                aux_trace_builder.build_aux_columns(main_trace, rand_elements)
+            },
+            MidenModule::Range => {
+                let lookup_values = vec![];
+                let cycle_lookups = BTreeMap::new();
+                let values_start = 0;
+                let aux_trace_builder =
+                    RangeAuxTraceBuilder::new(lookup_values, cycle_lookups, values_start);
+                aux_trace_builder.build_aux_columns(main_trace, rand_elements)
+            },
+        }
+    }
+}
+
+/// Builds the Miden VM auxiliary trace using the provided main trace and challenges.
+pub fn build_aux_trace_with_miden_vm<F, EF>(
+    main: &RowMajorMatrix<F>,
+    challenges: &[EF],
+    module: MidenModule,
+) -> RowMajorMatrix<EF>
+where
+    F: Field + PrimeField64,
+    EF: ExtensionField<F>,
+{
+    // Convert main trace to Miden format: transposed to column-major and use `BaseElement` instead
+    // of `F``
+    let main_transposed = main.transpose();
+    let mut felt_columns_vec = Vec::new();
+    for row in main_transposed.rows() {
+        let col_felt = row.map(|x| Felt::new(x.as_canonical_u64())).collect::<Vec<_>>();
+        felt_columns_vec.push(col_felt);
+    }
+    let col_matrix = ColMatrix::new(felt_columns_vec);
+    let last_program_row = main.height().into();
+    let main_trace = MainTrace::new(col_matrix, last_program_row);
+
+    // Convert challenges to Miden format
+    let mut rand_elements = Vec::new();
+    for r in challenges {
+        let coeffs: Vec<Felt> = r
+            .as_basis_coefficients_slice()
+            .iter()
+            .map(|x| Felt::new(x.as_canonical_u64()))
+            .collect();
+        let r_fe: &[QuadExtension<Felt>] = FieldElement::slice_from_base_elements(&coeffs);
+        rand_elements.push(r_fe[0]);
+    }
+
+    // Build aux trace using Miden VM AuxTraceBuilder
+    let aux_trace_miden = module.build_aux_columns(&main_trace, &rand_elements);
+    let aux_width = aux_trace_miden.len();
+
+    // Convert aux trace back to RowMajorMatrix<EF>
+    let num_rows = main.height();
+    let trace_length = num_rows * aux_width;
+    let long_trace = EF::zero_vec(trace_length);
+    let mut aux_trace = RowMajorMatrix::new(long_trace, aux_width);
+
+    for j in 0..aux_width {
+        let col = aux_trace_miden.get(j).unwrap();
+        for i in 0..num_rows {
+            let value_felt = col[i];
+            let coeffs_f: Vec<F> = value_felt
+                .to_base_elements()
+                .iter()
+                .map(|x| F::from_canonical_checked(x.as_int()).unwrap())
+                .collect();
+            let value_ef = EF::from_basis_coefficients_iter(coeffs_f.iter().cloned()).unwrap();
+            aux_trace.row_mut(i)[j] = value_ef;
+        }
+    }
+
+    aux_trace
+}

air-script/src/test_utils/mod.rs

@@ -0,0 +1,10 @@
+/// Macros to generate Air tester structs and tests for both Plonky3 and Winterfell backends.
+pub mod air_tester_macros;
+/// Code generation for tests/**/*.air files.
+pub mod codegen;
+/// Miden VM auxiliary trace generator
+pub mod miden_vm_aux_trace_generator;
+/// Plonky3-specific Debug constraint builder implementation
+pub mod plonky3_traits;
+/// Winterfell-specific traits
+pub mod winterfell_traits;