Target 0xMiden Plonky3 repo and use AirScriptAir and AirScriptBuilder traits

CHANGELOG.md

@@ -13,6 +13,7 @@
 - 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).
 
 ## 0.4.0 (2025-06-20)
 

air-script/Cargo.toml

@@ -30,18 +30,18 @@ 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 }
+p3-air = { package = "p3-air", git="https://github.com/0xMiden/Plonky3", rev = "2e6ff23090b5bd5138695ceb757181443a523c6d", default-features = false }
+p3-challenger = { package = "p3-challenger", git="https://github.com/0xMiden/Plonky3", rev = "2e6ff23090b5bd5138695ceb757181443a523c6d", default-features = false }
+p3-circle = { package = "p3-circle", git="https://github.com/0xMiden/Plonky3", rev = "2e6ff23090b5bd5138695ceb757181443a523c6d", default-features = false }
+p3-commit = { package = "p3-commit", git="https://github.com/0xMiden/Plonky3", rev = "2e6ff23090b5bd5138695ceb757181443a523c6d", default-features = false }
+p3-field = { package = "p3-field", git="https://github.com/0xMiden/Plonky3", rev = "2e6ff23090b5bd5138695ceb757181443a523c6d", default-features = false }
+p3-fri = { package = "p3-fri", git="https://github.com/0xMiden/Plonky3", rev = "2e6ff23090b5bd5138695ceb757181443a523c6d", default-features = false }
+p3-matrix = { package = "p3-matrix", git="https://github.com/0xMiden/Plonky3", rev = "2e6ff23090b5bd5138695ceb757181443a523c6d", default-features = false }
+p3-merkle-tree = { package = "p3-merkle-tree", git="https://github.com/0xMiden/Plonky3", rev = "2e6ff23090b5bd5138695ceb757181443a523c6d", default-features = false }
+p3-goldilocks = { package = "p3-goldilocks", git="https://github.com/0xMiden/Plonky3", rev = "2e6ff23090b5bd5138695ceb757181443a523c6d", default-features = false }
+p3-sha256 = { package = "p3-sha256", git="https://github.com/0xMiden/Plonky3", rev = "2e6ff23090b5bd5138695ceb757181443a523c6d", default-features = false }
+p3-symmetric = { package = "p3-symmetric", git="https://github.com/0xMiden/Plonky3", rev = "2e6ff23090b5bd5138695ceb757181443a523c6d", default-features = false }
+p3-uni-stark = { package = "p3-uni-stark", git="https://github.com/0xMiden/Plonky3", rev = "2e6ff23090b5bd5138695ceb757181443a523c6d", 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/tests/binary/binary_plonky3.rs

@@ -1,17 +1,20 @@
-use p3_air::{Air, AirBuilder, AirBuilderWithPublicValues, BaseAir, BaseAirWithPublicValues};
+use p3_air::{Air, BaseAir, BaseAirWithPublicValues, AirBuilder, ExtensionBuilder};
 use p3_matrix::Matrix;
-use p3_field::PrimeCharacteristicRing;
-use crate::helpers::{AirBuilderWithPeriodicColumns, BaseAirWithPeriodicColumns};
-
-pub const NUM_COLUMNS: usize = 2;
+use p3_field::{Field, PrimeCharacteristicRing};
+use crate::helpers::{AirScriptAir, AirScriptBuilder};
 
+pub const MAIN_WIDTH: usize = 2;
+pub const AUX_WIDTH: usize = 0;
+pub const NUM_PERIODIC_VALUES: usize = 0;
+pub const PERIOD: usize = 0;
 pub const NUM_PUBLIC_VALUES: usize = 16;
+pub const NUM_ALPHA_CHALLENGES: usize = 0;
 
 pub struct BinaryAir;
 
 impl<F> BaseAir<F> for BinaryAir {
     fn width(&self) -> usize {
-        NUM_COLUMNS
+        MAIN_WIDTH
     }
 }
 
@@ -21,24 +24,39 @@ impl<F> BaseAirWithPublicValues<F> for BinaryAir {
     }
 }
 
-impl<F: PrimeCharacteristicRing> BaseAirWithPeriodicColumns<F> for BinaryAir {
-    fn get_periodic_columns(&self) -> Vec<Vec<F>> {
-        vec![
-        ]
+impl<F: Field, AB: AirScriptBuilder<F = F>> AirScriptAir<F, AB> for BinaryAir {
+    fn aux_width(&self) -> usize {
+        AUX_WIDTH
+    }
+
+    fn num_alpha_challenges(&self) -> usize {
+        NUM_ALPHA_CHALLENGES
+    }
+
+    fn periodic_table(&self) -> Vec<Vec<F>> {
+        vec![]
     }
-}
 
-impl<AB: AirBuilderWithPublicValues + AirBuilderWithPeriodicColumns> Air<AB> for BinaryAir {
     fn eval(&self, builder: &mut AB) {
-        let main = builder.main();
         let public_values: [_; NUM_PUBLIC_VALUES] = builder.public_values().try_into().expect("Wrong number of public values");
-        let periodic_values = builder.periodic_columns();
+        let periodic_values: [_; NUM_PERIODIC_VALUES] = builder.periodic_evals().try_into().expect("Wrong number of periodic values");
+        let main = builder.main();
         let (main_current, main_next) = (
             main.row_slice(0).unwrap(),
             main.row_slice(1).unwrap(),
         );
-        builder.when_first_row().assert_zero::<_>(main_current[0].into() - public_values[0].into());
-        builder.assert_zero::<_>(main_current[0].into() * main_current[0].into() - main_current[0].into());
-        builder.assert_zero::<_>(main_current[1].into() * main_current[1].into() - main_current[1].into());
+
+        // Main boundary constraints
+        builder.when_first_row().assert_zero(main_current[0].clone().into() - public_values[0].into());
+
+        // Main integrity/transition constraints
+        builder.assert_zero(main_current[0].clone().into() * main_current[0].clone().into() - main_current[0].clone().into());
+        builder.assert_zero(main_current[1].clone().into() * main_current[1].clone().into() - main_current[1].clone().into());
+    }
+}
+
+impl<AB: AirScriptBuilder> Air<AB> for BinaryAir {
+    fn eval(&self, builder: &mut AB) {
+        <Self as AirScriptAir<AB::F, AB>>::eval(self, builder);
     }
 }

air-script/tests/binary/test_air_plonky3.rs

@@ -13,20 +13,20 @@ use p3_symmetric::{CompressionFunctionFromHasher, SerializingHasher};
 use p3_uni_stark::StarkConfig;
 
 use crate::{
-    binary::binary_plonky3::{BinaryAir, NUM_COLUMNS},
-    generate_air_plonky3_test,
-    helpers::check_constraints_with_periodic_columns,
+    binary::binary_plonky3::{BinaryAir, MAIN_WIDTH},
+    generate_air_plonky3_test_with_airscript_traits,
+    helpers::check_constraints_with_airscript_traits,
 };
 
 pub fn generate_trace_rows<F: PrimeField64>(inputs: Vec<u32>) -> RowMajorMatrix<F> {
     let num_rows = 32;
-    let trace_length = num_rows * NUM_COLUMNS;
+    let trace_length = num_rows * MAIN_WIDTH;
 
     let mut long_trace = F::zero_vec(trace_length);
 
-    let mut trace = RowMajorMatrix::new(long_trace, NUM_COLUMNS);
+    let mut trace = RowMajorMatrix::new(long_trace, MAIN_WIDTH);
 
-    let (prefix, rows, suffix) = unsafe { trace.values.align_to_mut::<[F; NUM_COLUMNS]>() };
+    let (prefix, rows, suffix) = unsafe { trace.values.align_to_mut::<[F; MAIN_WIDTH]>() };
     assert!(prefix.is_empty(), "Alignment should match");
     assert!(suffix.is_empty(), "Alignment should match");
     assert_eq!(rows.len(), num_rows);
@@ -52,4 +52,4 @@ fn generate_inputs() -> Vec<u32> {
     vec![1; 16]
 }
 
-generate_air_plonky3_test!(test_air_plonky3, BinaryAir);
+generate_air_plonky3_test_with_airscript_traits!(test_air_plonky3, BinaryAir);

air-script/tests/bitwise/bitwise_plonky3.rs

@@ -1,17 +1,20 @@
-use p3_air::{Air, AirBuilder, AirBuilderWithPublicValues, BaseAir, BaseAirWithPublicValues};
+use p3_air::{Air, BaseAir, BaseAirWithPublicValues, AirBuilder, ExtensionBuilder};
 use p3_matrix::Matrix;
-use p3_field::PrimeCharacteristicRing;
-use crate::helpers::{AirBuilderWithPeriodicColumns, BaseAirWithPeriodicColumns};
-
-pub const NUM_COLUMNS: usize = 14;
+use p3_field::{Field, PrimeCharacteristicRing};
+use crate::helpers::{AirScriptAir, AirScriptBuilder};
 
+pub const MAIN_WIDTH: usize = 14;
+pub const AUX_WIDTH: usize = 0;
+pub const NUM_PERIODIC_VALUES: usize = 2;
+pub const PERIOD: usize = 8;
 pub const NUM_PUBLIC_VALUES: usize = 16;
+pub const NUM_ALPHA_CHALLENGES: usize = 0;
 
 pub struct BitwiseAir;
 
 impl<F> BaseAir<F> for BitwiseAir {
     fn width(&self) -> usize {
-        NUM_COLUMNS
+        MAIN_WIDTH
     }
 }
 
@@ -21,41 +24,57 @@ impl<F> BaseAirWithPublicValues<F> for BitwiseAir {
     }
 }
 
-impl<F: PrimeCharacteristicRing> BaseAirWithPeriodicColumns<F> for BitwiseAir {
-    fn get_periodic_columns(&self) -> Vec<Vec<F>> {
+impl<F: Field, AB: AirScriptBuilder<F = F>> AirScriptAir<F, AB> for BitwiseAir {
+    fn aux_width(&self) -> usize {
+        AUX_WIDTH
+    }
+
+    fn num_alpha_challenges(&self) -> usize {
+        NUM_ALPHA_CHALLENGES
+    }
+
+    fn periodic_table(&self) -> Vec<Vec<F>> {
         vec![
             vec![F::from_u64(1), F::from_u64(0), F::from_u64(0), F::from_u64(0), F::from_u64(0), F::from_u64(0), F::from_u64(0), F::from_u64(0)],
             vec![F::from_u64(1), F::from_u64(1), F::from_u64(1), F::from_u64(1), F::from_u64(1), F::from_u64(1), F::from_u64(1), F::from_u64(0)],
         ]
     }
-}
 
-impl<AB: AirBuilderWithPublicValues + AirBuilderWithPeriodicColumns> Air<AB> for BitwiseAir {
     fn eval(&self, builder: &mut AB) {
-        let main = builder.main();
         let public_values: [_; NUM_PUBLIC_VALUES] = builder.public_values().try_into().expect("Wrong number of public values");
-        let periodic_values = builder.periodic_columns();
+        let periodic_values: [_; NUM_PERIODIC_VALUES] = builder.periodic_evals().try_into().expect("Wrong number of periodic values");
+        let main = builder.main();
         let (main_current, main_next) = (
             main.row_slice(0).unwrap(),
             main.row_slice(1).unwrap(),
         );
-        builder.when_first_row().assert_zero::<_>(main_current[13].into());
-        builder.assert_zero::<_>(main_current[0].into() * main_current[0].into() - main_current[0].into());
-        builder.when_transition().assert_zero::<_>(periodic_values[1].into() * (main_next[0].into() - main_current[0].into()));
-        builder.assert_zero::<_>(main_current[3].into() * main_current[3].into() - main_current[3].into());
-        builder.assert_zero::<_>(main_current[4].into() * main_current[4].into() - main_current[4].into());
-        builder.assert_zero::<_>(main_current[5].into() * main_current[5].into() - main_current[5].into());
-        builder.assert_zero::<_>(main_current[6].into() * main_current[6].into() - main_current[6].into());
-        builder.assert_zero::<_>(main_current[7].into() * main_current[7].into() - main_current[7].into());
-        builder.assert_zero::<_>(main_current[8].into() * main_current[8].into() - main_current[8].into());
-        builder.assert_zero::<_>(main_current[9].into() * main_current[9].into() - main_current[9].into());
-        builder.assert_zero::<_>(main_current[10].into() * main_current[10].into() - main_current[10].into());
-        builder.assert_zero::<_>(periodic_values[0].into() * (main_current[1].into() - (main_current[3].into() + main_current[4].into().double() + AB::Expr::from_u64(4) * main_current[5].into() + AB::Expr::from_u64(8) * main_current[6].into())));
-        builder.assert_zero::<_>(periodic_values[0].into() * (main_current[2].into() - (main_current[7].into() + main_current[8].into().double() + AB::Expr::from_u64(4) * main_current[9].into() + AB::Expr::from_u64(8) * main_current[10].into())));
-        builder.when_transition().assert_zero::<_>(periodic_values[1].into() * (main_next[1].into() - (main_current[1].into() * AB::Expr::from_u64(16) + main_current[3].into() + main_current[4].into().double() + AB::Expr::from_u64(4) * main_current[5].into() + AB::Expr::from_u64(8) * main_current[6].into())));
-        builder.when_transition().assert_zero::<_>(periodic_values[1].into() * (main_next[2].into() - (main_current[2].into() * AB::Expr::from_u64(16) + main_current[7].into() + main_current[8].into().double() + AB::Expr::from_u64(4) * main_current[9].into() + AB::Expr::from_u64(8) * main_current[10].into())));
-        builder.assert_zero::<_>(periodic_values[0].into() * main_current[11].into());
-        builder.when_transition().assert_zero::<_>(periodic_values[1].into() * (main_current[12].into() - main_next[11].into()));
-        builder.assert_zero::<_>((AB::Expr::ONE - main_current[0].into()) * (main_current[12].into() - (main_current[11].into() * AB::Expr::from_u64(16) + main_current[3].into() * main_current[7].into() + main_current[4].into().double() * main_current[8].into() + AB::Expr::from_u64(4) * main_current[5].into() * main_current[9].into() + AB::Expr::from_u64(8) * main_current[6].into() * main_current[10].into())) + main_current[0].into() * (main_current[12].into() - (main_current[11].into() * AB::Expr::from_u64(16) + main_current[3].into() + main_current[7].into() - main_current[3].into().double() * main_current[7].into() + (main_current[4].into() + main_current[8].into() - main_current[4].into().double() * main_current[8].into()).double() + AB::Expr::from_u64(4) * (main_current[5].into() + main_current[9].into() - main_current[5].into().double() * main_current[9].into()) + AB::Expr::from_u64(8) * (main_current[6].into() + main_current[10].into() - main_current[6].into().double() * main_current[10].into()))));
+
+        // Main boundary constraints
+        builder.when_first_row().assert_zero(main_current[13].clone().into());
+
+        // Main integrity/transition constraints
+        builder.assert_zero(main_current[0].clone().into() * main_current[0].clone().into() - main_current[0].clone().into());
+        builder.when_transition().assert_zero_ext(periodic_values[1].into() * (AB::ExprEF::from(main_next[0].clone().into()) - AB::ExprEF::from(main_current[0].clone().into())));
+        builder.assert_zero(main_current[3].clone().into() * main_current[3].clone().into() - main_current[3].clone().into());
+        builder.assert_zero(main_current[4].clone().into() * main_current[4].clone().into() - main_current[4].clone().into());
+        builder.assert_zero(main_current[5].clone().into() * main_current[5].clone().into() - main_current[5].clone().into());
+        builder.assert_zero(main_current[6].clone().into() * main_current[6].clone().into() - main_current[6].clone().into());
+        builder.assert_zero(main_current[7].clone().into() * main_current[7].clone().into() - main_current[7].clone().into());
+        builder.assert_zero(main_current[8].clone().into() * main_current[8].clone().into() - main_current[8].clone().into());
+        builder.assert_zero(main_current[9].clone().into() * main_current[9].clone().into() - main_current[9].clone().into());
+        builder.assert_zero(main_current[10].clone().into() * main_current[10].clone().into() - main_current[10].clone().into());
+        builder.assert_zero_ext(periodic_values[0].into() * (AB::ExprEF::from(main_current[1].clone().into()) - (AB::ExprEF::from(main_current[3].clone().into()) + AB::ExprEF::from(main_current[4].clone().into()).double() + AB::ExprEF::from_u64(4) * AB::ExprEF::from(main_current[5].clone().into()) + AB::ExprEF::from_u64(8) * AB::ExprEF::from(main_current[6].clone().into()))));
+        builder.assert_zero_ext(periodic_values[0].into() * (AB::ExprEF::from(main_current[2].clone().into()) - (AB::ExprEF::from(main_current[7].clone().into()) + AB::ExprEF::from(main_current[8].clone().into()).double() + AB::ExprEF::from_u64(4) * AB::ExprEF::from(main_current[9].clone().into()) + AB::ExprEF::from_u64(8) * AB::ExprEF::from(main_current[10].clone().into()))));
+        builder.when_transition().assert_zero_ext(periodic_values[1].into() * (AB::ExprEF::from(main_next[1].clone().into()) - (AB::ExprEF::from(main_current[1].clone().into()) * AB::ExprEF::from_u64(16) + AB::ExprEF::from(main_current[3].clone().into()) + AB::ExprEF::from(main_current[4].clone().into()).double() + AB::ExprEF::from_u64(4) * AB::ExprEF::from(main_current[5].clone().into()) + AB::ExprEF::from_u64(8) * AB::ExprEF::from(main_current[6].clone().into()))));
+        builder.when_transition().assert_zero_ext(periodic_values[1].into() * (AB::ExprEF::from(main_next[2].clone().into()) - (AB::ExprEF::from(main_current[2].clone().into()) * AB::ExprEF::from_u64(16) + AB::ExprEF::from(main_current[7].clone().into()) + AB::ExprEF::from(main_current[8].clone().into()).double() + AB::ExprEF::from_u64(4) * AB::ExprEF::from(main_current[9].clone().into()) + AB::ExprEF::from_u64(8) * AB::ExprEF::from(main_current[10].clone().into()))));
+        builder.assert_zero_ext(periodic_values[0].into() * AB::ExprEF::from(main_current[11].clone().into()));
+        builder.when_transition().assert_zero_ext(periodic_values[1].into() * (AB::ExprEF::from(main_current[12].clone().into()) - AB::ExprEF::from(main_next[11].clone().into())));
+        builder.assert_zero((AB::Expr::ONE - main_current[0].clone().into()) * (main_current[12].clone().into() - (main_current[11].clone().into() * AB::Expr::from_u64(16) + main_current[3].clone().into() * main_current[7].clone().into() + main_current[4].clone().into().double() * main_current[8].clone().into() + AB::Expr::from_u64(4) * main_current[5].clone().into() * main_current[9].clone().into() + AB::Expr::from_u64(8) * main_current[6].clone().into() * main_current[10].clone().into())) + main_current[0].clone().into() * (main_current[12].clone().into() - (main_current[11].clone().into() * AB::Expr::from_u64(16) + main_current[3].clone().into() + main_current[7].clone().into() - main_current[3].clone().into().double() * main_current[7].clone().into() + (main_current[4].clone().into() + main_current[8].clone().into() - main_current[4].clone().into().double() * main_current[8].clone().into()).double() + AB::Expr::from_u64(4) * (main_current[5].clone().into() + main_current[9].clone().into() - main_current[5].clone().into().double() * main_current[9].clone().into()) + AB::Expr::from_u64(8) * (main_current[6].clone().into() + main_current[10].clone().into() - main_current[6].clone().into().double() * main_current[10].clone().into()))));
+    }
+}
+
+impl<AB: AirScriptBuilder> Air<AB> for BitwiseAir {
+    fn eval(&self, builder: &mut AB) {
+        <Self as AirScriptAir<AB::F, AB>>::eval(self, builder);
     }
 }

air-script/tests/bitwise/test_air_plonky3.rs

@@ -13,20 +13,20 @@ use p3_symmetric::{CompressionFunctionFromHasher, SerializingHasher};
 use p3_uni_stark::StarkConfig;
 
 use crate::{
-    bitwise::bitwise_plonky3::{BitwiseAir, NUM_COLUMNS},
-    generate_air_plonky3_test,
-    helpers::check_constraints_with_periodic_columns,
+    bitwise::bitwise_plonky3::{BitwiseAir, MAIN_WIDTH},
+    generate_air_plonky3_test_with_airscript_traits,
+    helpers::check_constraints_with_airscript_traits,
 };
 
 pub fn generate_trace_rows<F: PrimeField64>(inputs: Vec<u32>) -> RowMajorMatrix<F> {
     let num_rows = 32;
-    let trace_length = num_rows * NUM_COLUMNS;
+    let trace_length = num_rows * MAIN_WIDTH;
 
     let mut long_trace = F::zero_vec(trace_length);
 
-    let mut trace = RowMajorMatrix::new(long_trace, NUM_COLUMNS);
+    let mut trace = RowMajorMatrix::new(long_trace, MAIN_WIDTH);
 
-    let (prefix, rows, suffix) = unsafe { trace.values.align_to_mut::<[F; NUM_COLUMNS]>() };
+    let (prefix, rows, suffix) = unsafe { trace.values.align_to_mut::<[F; MAIN_WIDTH]>() };
     assert!(prefix.is_empty(), "Alignment should match");
     assert!(suffix.is_empty(), "Alignment should match");
     assert_eq!(rows.len(), num_rows);
@@ -90,4 +90,4 @@ fn generate_inputs() -> Vec<u32> {
     vec![1; 16]
 }
 
-generate_air_plonky3_test!(test_air_plonky3, BitwiseAir);
+generate_air_plonky3_test_with_airscript_traits!(test_air_plonky3, BitwiseAir);