Add bus types and variable length public inputs to Plonky3 backend tests

air-script/Cargo.toml

@@ -29,9 +29,9 @@ 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 = "95fb1c965cc938eb1e04b10a6080138d4699b1ac", default-features = false }
-p3-field = { package = "p3-field", git="https://github.com/0xMiden/Plonky3", rev = "95fb1c965cc938eb1e04b10a6080138d4699b1ac", default-features = false }
-p3-miden-air = { package = "miden-air", git="https://github.com/0xMiden/Plonky3", rev = "95fb1c965cc938eb1e04b10a6080138d4699b1ac", default-features = false }
+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 }
 
 # MassaLabs fork
 miden-processor = { package = "miden-processor", git="https://github.com/massalabs/miden-vm", rev = "bc553af69a2543a0789830e8508b019694528181", default-features = false }
@@ -41,16 +41,16 @@ miden-air = { package = "miden-air", git="https://github.com/massalabs/miden-vm"
 expect-test = "1.4"
 
 # 0xMiden Plonky3 Fork
-p3-air = { package = "p3-air", git="https://github.com/0xMiden/Plonky3", rev = "95fb1c965cc938eb1e04b10a6080138d4699b1ac", default-features = false }
-p3-challenger = { package = "p3-challenger", git="https://github.com/0xMiden/Plonky3", rev = "95fb1c965cc938eb1e04b10a6080138d4699b1ac", default-features = false }
-p3-commit = { package = "p3-commit", git="https://github.com/0xMiden/Plonky3", rev = "95fb1c965cc938eb1e04b10a6080138d4699b1ac", default-features = false }
-p3-dft = { package = "p3-dft", git="https://github.com/0xMiden/Plonky3", rev = "95fb1c965cc938eb1e04b10a6080138d4699b1ac", default-features = false }
-p3-fri = { package = "p3-fri", git="https://github.com/0xMiden/Plonky3", rev = "95fb1c965cc938eb1e04b10a6080138d4699b1ac", default-features = false }
-p3-merkle-tree = { package = "p3-merkle-tree", git="https://github.com/0xMiden/Plonky3", rev = "95fb1c965cc938eb1e04b10a6080138d4699b1ac", default-features = false }
-p3-miden-prover = { package = "miden-prover", git="https://github.com/0xMiden/Plonky3", rev = "95fb1c965cc938eb1e04b10a6080138d4699b1ac", default-features = false }
-p3-goldilocks = { package = "p3-goldilocks", git="https://github.com/0xMiden/Plonky3", rev = "95fb1c965cc938eb1e04b10a6080138d4699b1ac", default-features = false }
-p3-sha256 = { package = "p3-sha256", git="https://github.com/0xMiden/Plonky3", rev = "95fb1c965cc938eb1e04b10a6080138d4699b1ac", default-features = false }
-p3-symmetric = { package = "p3-symmetric", git="https://github.com/0xMiden/Plonky3", rev = "95fb1c965cc938eb1e04b10a6080138d4699b1ac", default-features = false }
+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 }
 
 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

@@ -66,17 +66,31 @@ macro_rules! generate_air_plonky3_test_with_airscript_traits {
             let pcs = Pcs::new(dft, val_mmcs, fri_params);
             let config = MyConfig::new(pcs, challenger);
 
+            // Generate public inputs and convert them to Goldilocks field elements
             let inputs = generate_inputs();
-            let inputs_goldilocks: Vec<Val> = inputs
-                .iter()
-                .map(|&x| <Val as p3_field::PrimeCharacteristicRing>::from_u64(x))
-                .collect();
+            let inputs_goldilocks = crate::test_utils::pub_inputs_conversion_utils::convert_pub_inputs_to_goldilocks(&inputs);
+
+            // Generate variable-length public inputs as a `Vec<Vec<Vec<u64>>>`.
+            // The outer `Vec` represents multiple tables (one for each bus in the AIR)
+            // The middle `Vec` represents the rows of each table
+            // The innermost `Vec` represents the tuple values contained by the bus, that will be combined with randomness into a single field element.
+            let var_len_pub_inputs = generate_var_len_pub_inputs();
+            // Convert variable-length public inputs to Goldilocks field elements
+            let var_len_pub_inputs_goldilocks_vec = crate::test_utils::pub_inputs_conversion_utils::convert_var_len_pub_inputs_to_goldilocks(var_len_pub_inputs);
+            let var_len_pub_inputs_goldilocks_vec_slice = crate::test_utils::pub_inputs_conversion_utils::convert_inner_vec_to_slice(&var_len_pub_inputs_goldilocks_vec);
+            let var_len_pub_inputs_goldilocks = crate::test_utils::pub_inputs_conversion_utils::convert_mid_vec_to_slice(&var_len_pub_inputs_goldilocks_vec_slice);
 
             let trace = generate_trace_rows::<Val>(inputs);
 
             let proof = p3_miden_prover::prove(&config, &$air_name {}, &trace, &inputs_goldilocks);
-            p3_miden_prover::verify(&config, &$air_name {}, &proof, &inputs_goldilocks)
-                .expect("Verification failed");
+            p3_miden_prover::verify(
+                &config,
+                &$air_name {},
+                &proof,
+                &inputs_goldilocks,
+                &var_len_pub_inputs_goldilocks,
+            )
+            .expect("Verification failed");
         }
     };
 }

air-script/src/test_utils/mod.rs

@@ -2,5 +2,7 @@
 pub mod air_tester_macros;
 /// Code generation for tests/**/*.air files.
 pub mod codegen;
+/// Conversion utilities for test inputs (both public inputs and variable-length public inputs).
+pub mod pub_inputs_conversion_utils;
 /// Winterfell-specific traits
 pub mod winterfell_traits;

air-script/src/test_utils/pub_inputs_conversion_utils.rs

@@ -0,0 +1,46 @@
+type Val = p3_goldilocks::Goldilocks;
+
+/// Converts public inputs from u64 to Goldilocks field elements.
+pub(crate) fn convert_pub_inputs_to_goldilocks(pub_inputs: &[u64]) -> Vec<Val> {
+    pub_inputs
+        .iter()
+        .map(|&x| <Val as p3_field::PrimeCharacteristicRing>::from_u64(x))
+        .collect()
+}
+
+/// Converts variable-length public inputs from u64 to Goldilocks field elements.
+/// The input should be a Vec of tables (Vec<Vec<u64>>) in a RowMajor format.
+pub(crate) fn convert_var_len_pub_inputs_to_goldilocks(
+    var_len_pub_inputs: Vec<Vec<Vec<u64>>>,
+) -> Vec<Vec<Vec<Val>>> {
+    let mut var_len_pub_inputs_goldilocks_vec: Vec<Vec<Vec<Val>>> = vec![];
+    for arr in var_len_pub_inputs.iter() {
+        let mut goldilocks_arr: Vec<Vec<Val>> = vec![];
+        for slice in arr.iter() {
+            let goldilocks_slice: Vec<Val> = slice
+                .iter()
+                .map(|&x| <Val as p3_field::PrimeCharacteristicRing>::from_u64(x))
+                .collect();
+            goldilocks_arr.push(goldilocks_slice);
+        }
+        var_len_pub_inputs_goldilocks_vec.push(goldilocks_arr);
+    }
+    var_len_pub_inputs_goldilocks_vec
+}
+
+/// Converts the innermost vectors of variable-length public inputs to slices.
+pub(crate) fn convert_inner_vec_to_slice<'a>(
+    var_len_pub_inputs: &'a Vec<Vec<Vec<Val>>>,
+) -> Vec<Vec<&'a [Val]>> {
+    var_len_pub_inputs
+        .iter()
+        .map(|outer| outer.iter().map(|inner| inner.as_slice()).collect())
+        .collect()
+}
+
+/// Converts the middle vectors of variable-length public inputs to slices.
+pub(crate) fn convert_mid_vec_to_slice<'a>(
+    var_len_pub_inputs: &'a Vec<Vec<&'a [Val]>>,
+) -> Vec<&'a [&'a [Val]]> {
+    var_len_pub_inputs.iter().map(|v| v.as_slice()).collect()
+}

air-script/src/tests/binary/binary_plonky3.rs

@@ -2,7 +2,7 @@ use p3_field::{ExtensionField, Field, PrimeCharacteristicRing};
 use p3_matrix::Matrix;
 use p3_matrix::dense::RowMajorMatrixView;
 use p3_matrix::stack::VerticalPair;
-use p3_miden_air::{MidenAir, MidenAirBuilder, RowMajorMatrix};
+use p3_miden_air::{BusType, MidenAir, MidenAirBuilder, RowMajorMatrix};
 
 pub const MAIN_WIDTH: usize = 2;
 pub const AUX_WIDTH: usize = 0;
@@ -38,8 +38,6 @@ impl<F, EF> MidenAir<F, EF> for BinaryAir {
         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());
 
-        // Aux boundary constraints
-
         // Aux integrity/transition constraints
     }
 }

air-script/src/tests/binary/test_air_plonky3.rs

@@ -40,4 +40,8 @@ fn generate_inputs() -> Vec<u64> {
     vec![1; 16]
 }
 
+fn generate_var_len_pub_inputs<'a>() -> Vec<Vec<Vec<u64>>> {
+    vec![]
+}
+
 generate_air_plonky3_test_with_airscript_traits!(test_air_plonky3, BinaryAir);

air-script/src/tests/bitwise/bitwise_plonky3.rs

@@ -2,7 +2,7 @@ use p3_field::{ExtensionField, Field, PrimeCharacteristicRing};
 use p3_matrix::Matrix;
 use p3_matrix::dense::RowMajorMatrixView;
 use p3_matrix::stack::VerticalPair;
-use p3_miden_air::{MidenAir, MidenAirBuilder, RowMajorMatrix};
+use p3_miden_air::{BusType, MidenAir, MidenAirBuilder, RowMajorMatrix};
 
 pub const MAIN_WIDTH: usize = 14;
 pub const AUX_WIDTH: usize = 0;
@@ -67,8 +67,6 @@ where F: Field,
         builder.when_transition().assert_zero_ext(AB::ExprEF::from(periodic_values[1].clone().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()))));
 
-        // Aux boundary constraints
-
         // Aux integrity/transition constraints
     }
 }

air-script/src/tests/bitwise/test_air_plonky3.rs

@@ -78,4 +78,8 @@ fn generate_inputs() -> Vec<u64> {
     vec![1; 16]
 }
 
+fn generate_var_len_pub_inputs<'a>() -> Vec<Vec<Vec<u64>>> {
+    vec![]
+}
+
 generate_air_plonky3_test_with_airscript_traits!(test_air_plonky3, BitwiseAir);

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

@@ -2,7 +2,7 @@ use p3_field::{ExtensionField, Field, PrimeCharacteristicRing};
 use p3_matrix::Matrix;
 use p3_matrix::dense::RowMajorMatrixView;
 use p3_matrix::stack::VerticalPair;
-use p3_miden_air::{MidenAir, MidenAirBuilder, RowMajorMatrix};
+use p3_miden_air::{BusType, MidenAir, MidenAirBuilder, RowMajorMatrix};
 
 pub const MAIN_WIDTH: usize = 7;
 pub const AUX_WIDTH: usize = 2;
@@ -29,6 +29,13 @@ where F: Field,
         AUX_WIDTH
     }
 
+    fn bus_types(&self) -> Vec<BusType> {
+        vec![
+            BusType::Multiset,
+            BusType::Logup,
+        ]
+    }
+
     fn build_aux_trace(&self, _main: &RowMajorMatrix<F>, _challenges: &[EF]) -> Option<RowMajorMatrix<F>> {
         // Note: consider using Some(build_aux_trace_with_miden_vm::<F, EF>(_main, _challenges, module)) if you want to build the aux trace using Miden VM aux trace builders.
 
@@ -98,12 +105,6 @@ where F: Field,
         builder.assert_zero(main_current[2].clone().into() * main_current[2].clone().into() - main_current[2].clone().into());
         builder.assert_zero(main_current[3].clone().into() * main_current[3].clone().into() - main_current[3].clone().into());
 
-        // Aux boundary constraints
-        builder.when_first_row().assert_zero_ext(AB::ExprEF::from(aux_current[0].clone().into()) - AB::ExprEF::ONE);
-        builder.when_last_row().assert_zero_ext(AB::ExprEF::from(aux_current[0].clone().into()) - AB::ExprEF::ONE);
-        builder.when_first_row().assert_zero_ext(AB::ExprEF::from(aux_current[1].clone().into()));
-        builder.when_last_row().assert_zero_ext(AB::ExprEF::from(aux_current[1].clone().into()));
-
         // Aux integrity/transition constraints
         builder.when_transition().assert_zero_ext(((alpha.into() + beta_challenges[0].into() + (AB::ExprEF::from_u64(3) + AB::ExprEF::from(main_current[1].clone().into())) * beta_challenges[1].into() + AB::ExprEF::from(main_current[0].clone().into()) * beta_challenges[2].into()) * AB::ExprEF::from(main_current[2].clone().into()) + AB::ExprEF::ONE - AB::ExprEF::from(main_current[2].clone().into())) * ((alpha.into() + beta_challenges[0].into().double() + AB::ExprEF::from(main_current[1].clone().into()) * beta_challenges[1].into()) * (AB::ExprEF::ONE - AB::ExprEF::from(main_current[2].clone().into())) + AB::ExprEF::from(main_current[2].clone().into())) * AB::ExprEF::from(aux_current[0].clone().into()) - ((alpha.into() + beta_challenges[0].into() + (AB::ExprEF::from_u64(3) + AB::ExprEF::from(main_current[1].clone().into())) * beta_challenges[1].into() + AB::ExprEF::from(main_current[1].clone().into()) * beta_challenges[2].into()) * AB::ExprEF::from(main_current[3].clone().into()) + AB::ExprEF::ONE - AB::ExprEF::from(main_current[3].clone().into())) * ((alpha.into() + beta_challenges[0].into().double() + AB::ExprEF::from(main_current[0].clone().into()) * beta_challenges[1].into()) * (AB::ExprEF::ONE - AB::ExprEF::from(main_current[3].clone().into())) + AB::ExprEF::from(main_current[3].clone().into())) * AB::ExprEF::from(aux_next[0].clone().into()));
         builder.when_transition().assert_zero_ext((alpha.into() + AB::ExprEF::from_u64(3) * beta_challenges[0].into() + AB::ExprEF::from(main_current[0].clone().into()) * beta_challenges[1].into()) * (alpha.into() + AB::ExprEF::from_u64(3) * beta_challenges[0].into() + AB::ExprEF::from(main_current[0].clone().into()) * beta_challenges[1].into()) * (alpha.into() + AB::ExprEF::from_u64(3) * beta_challenges[0].into() + AB::ExprEF::from(main_current[1].clone().into()) * beta_challenges[1].into()) * AB::ExprEF::from(aux_current[1].clone().into()) + (alpha.into() + AB::ExprEF::from_u64(3) * beta_challenges[0].into() + AB::ExprEF::from(main_current[0].clone().into()) * beta_challenges[1].into()) * (alpha.into() + AB::ExprEF::from_u64(3) * beta_challenges[0].into() + AB::ExprEF::from(main_current[1].clone().into()) * beta_challenges[1].into()) * AB::ExprEF::from(main_current[4].clone().into()) + (alpha.into() + AB::ExprEF::from_u64(3) * beta_challenges[0].into() + AB::ExprEF::from(main_current[0].clone().into()) * beta_challenges[1].into()) * (alpha.into() + AB::ExprEF::from_u64(3) * beta_challenges[0].into() + AB::ExprEF::from(main_current[1].clone().into()) * beta_challenges[1].into()) * AB::ExprEF::from(main_current[5].clone().into()) - ((alpha.into() + AB::ExprEF::from_u64(3) * beta_challenges[0].into() + AB::ExprEF::from(main_current[0].clone().into()) * beta_challenges[1].into()) * (alpha.into() + AB::ExprEF::from_u64(3) * beta_challenges[0].into() + AB::ExprEF::from(main_current[0].clone().into()) * beta_challenges[1].into()) * (alpha.into() + AB::ExprEF::from_u64(3) * beta_challenges[0].into() + AB::ExprEF::from(main_current[1].clone().into()) * beta_challenges[1].into()) * AB::ExprEF::from(aux_next[1].clone().into()) + (alpha.into() + AB::ExprEF::from_u64(3) * beta_challenges[0].into() + AB::ExprEF::from(main_current[0].clone().into()) * beta_challenges[1].into()) * (alpha.into() + AB::ExprEF::from_u64(3) * beta_challenges[0].into() + AB::ExprEF::from(main_current[0].clone().into()) * beta_challenges[1].into()) * AB::ExprEF::from(main_current[6].clone().into())));

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

@@ -2,7 +2,7 @@ use p3_field::{ExtensionField, Field, PrimeCharacteristicRing};
 use p3_matrix::Matrix;
 use p3_matrix::dense::RowMajorMatrixView;
 use p3_matrix::stack::VerticalPair;
-use p3_miden_air::{MidenAir, MidenAirBuilder, RowMajorMatrix};
+use p3_miden_air::{BusType, MidenAir, MidenAirBuilder, RowMajorMatrix};
 
 pub const MAIN_WIDTH: usize = 1;
 pub const AUX_WIDTH: usize = 2;
@@ -29,6 +29,13 @@ where F: Field,
         AUX_WIDTH
     }
 
+    fn bus_types(&self) -> Vec<BusType> {
+        vec![
+            BusType::Multiset,
+            BusType::Logup,
+        ]
+    }
+
     fn build_aux_trace(&self, _main: &RowMajorMatrix<F>, _challenges: &[EF]) -> Option<RowMajorMatrix<F>> {
         // Note: consider using Some(build_aux_trace_with_miden_vm::<F, EF>(_main, _challenges, module)) if you want to build the aux trace using Miden VM aux trace builders.
 
@@ -95,11 +102,6 @@ where F: Field,
 
         // Main integrity/transition constraints
 
-        // Aux boundary constraints
-        builder.when_last_row().assert_zero_ext(AB::ExprEF::from(aux_current[0].clone().into()) - AB::ExprEF::ONE);
-        builder.when_first_row().assert_zero_ext(AB::ExprEF::from(aux_current[1].clone().into()));
-        builder.when_last_row().assert_zero_ext(AB::ExprEF::from(aux_current[1].clone().into()));
-
         // Aux integrity/transition constraints
         builder.when_transition().assert_zero_ext(((alpha.into() + beta_challenges[0].into()) * AB::ExprEF::from(main_current[0].clone().into()) + AB::ExprEF::ONE - AB::ExprEF::from(main_current[0].clone().into())) * AB::ExprEF::from(aux_current[0].clone().into()) - ((alpha.into() + beta_challenges[0].into()) * (AB::ExprEF::ONE - AB::ExprEF::from(main_current[0].clone().into())) + AB::ExprEF::from(main_current[0].clone().into())) * AB::ExprEF::from(aux_next[0].clone().into()));
         builder.when_transition().assert_zero_ext((alpha.into() + beta_challenges[0].into() + beta_challenges[1].into().double()) * (alpha.into() + beta_challenges[0].into() + beta_challenges[1].into().double()) * AB::ExprEF::from(aux_current[1].clone().into()) + (alpha.into() + beta_challenges[0].into() + beta_challenges[1].into().double()) * AB::ExprEF::from(main_current[0].clone().into()) - ((alpha.into() + beta_challenges[0].into() + beta_challenges[1].into().double()) * (alpha.into() + beta_challenges[0].into() + beta_challenges[1].into().double()) * AB::ExprEF::from(aux_next[1].clone().into()) + (alpha.into() + beta_challenges[0].into() + beta_challenges[1].into().double()).double()));