Allow support for computed indices

CHANGELOG.md

@@ -3,12 +3,13 @@
 ## 0.5.0 (TBD)
 
 - Incremented MSRV to 1.89.
-- Add a constant propagation pass after other mir passes (#439).
 - Add a pass for common subexpression elimination in the constraint graph (#419).
+- Reorder unrolling passes to unroll list comprehensions before match statements (#431).
 - Refactored the unrolling pass in MIR (#434).
-- Fix regressions on MIR and list_comprehensions (#449).
 - Update documentation and tests thereof (#437).
-- Reorder unrolling passes to unroll list comprehensions before match statements (#431).
+- Add a constant propagation pass after other mir passes (#439).
+- Allow computed indices (#444).
+- Fix regressions on MIR and list_comprehensions (#449).
 
 ## 0.4.0 (2025-06-20)
 

air-script/tests/codegen/winterfell.rs

@@ -2,9 +2,6 @@ use expect_test::expect_file;
 
 use super::helpers::{Target, Test};
 
-// tests_wo_mir
-// ================================================================================================
-
 #[test]
 fn binary() {
     let generated_air = Test::new("tests/binary/binary.air".to_string())
@@ -16,21 +13,12 @@ fn binary() {
 }
 
 #[test]
-fn buses_simple() {
-    let generated_air = Test::new("tests/buses/buses_simple.air".to_string())
-        .transpile(Target::Winterfell)
-        .unwrap();
-
-    let expected = expect_file!["../buses/buses_simple.rs"];
-    expected.assert_eq(&generated_air);
-}
-#[test]
-fn buses_simple_with_evaluators() {
-    let generated_air = Test::new("tests/buses/buses_simple_with_evaluators.air".to_string())
+fn bitwise() {
+    let generated_air = Test::new("tests/bitwise/bitwise.air".to_string())
         .transpile(Target::Winterfell)
         .unwrap();
 
-    let expected = expect_file!["../buses/buses_simple.rs"];
+    let expected = expect_file!["../bitwise/bitwise.rs"];
     expected.assert_eq(&generated_air);
 }
 
@@ -45,22 +33,22 @@ fn buses_complex() {
 }
 
 #[test]
-fn buses_varlen_boundary_first() {
-    let generated_air = Test::new("tests/buses/buses_varlen_boundary_first.air".to_string())
+fn buses_simple() {
+    let generated_air = Test::new("tests/buses/buses_simple.air".to_string())
         .transpile(Target::Winterfell)
         .unwrap();
 
-    let expected = expect_file!["../buses/buses_varlen_boundary_first.rs"];
+    let expected = expect_file!["../buses/buses_simple.rs"];
     expected.assert_eq(&generated_air);
 }
 
 #[test]
-fn buses_varlen_boundary_last() {
-    let generated_air = Test::new("tests/buses/buses_varlen_boundary_last.air".to_string())
+fn buses_simple_with_evaluators() {
+    let generated_air = Test::new("tests/buses/buses_simple_with_evaluators.air".to_string())
         .transpile(Target::Winterfell)
         .unwrap();
 
-    let expected = expect_file!["../buses/buses_varlen_boundary_last.rs"];
+    let expected = expect_file!["../buses/buses_simple.rs"];
     expected.assert_eq(&generated_air);
 }
 
@@ -75,42 +63,53 @@ fn buses_varlen_boundary_both() {
 }
 
 #[test]
-fn periodic_columns() {
-    let generated_air = Test::new("tests/periodic_columns/periodic_columns.air".to_string())
+fn buses_varlen_boundary_first() {
+    let generated_air = Test::new("tests/buses/buses_varlen_boundary_first.air".to_string())
         .transpile(Target::Winterfell)
         .unwrap();
 
-    let expected = expect_file!["../periodic_columns/periodic_columns.rs"];
+    let expected = expect_file!["../buses/buses_varlen_boundary_first.rs"];
     expected.assert_eq(&generated_air);
 }
 
 #[test]
-fn pub_inputs() {
-    let generated_air = Test::new("tests/pub_inputs/pub_inputs.air".to_string())
+fn buses_varlen_boundary_last() {
+    let generated_air = Test::new("tests/buses/buses_varlen_boundary_last.air".to_string())
         .transpile(Target::Winterfell)
         .unwrap();
 
-    let expected = expect_file!["../pub_inputs/pub_inputs.rs"];
+    let expected = expect_file!["../buses/buses_varlen_boundary_last.rs"];
     expected.assert_eq(&generated_air);
 }
 
 #[test]
-fn system() {
-    let generated_air = Test::new("tests/system/system.air".to_string())
+fn computed_indices_complex() {
+    let generated_air =
+        Test::new("tests/computed_indices/computed_indices_complex.air".to_string())
+            .transpile(Target::Winterfell)
+            .unwrap();
+
+    let expected = expect_file!["../computed_indices/computed_indices_complex.rs"];
+    expected.assert_eq(&generated_air);
+}
+
+#[test]
+fn computed_indices_simple() {
+    let generated_air = Test::new("tests/computed_indices/computed_indices_simple.air".to_string())
         .transpile(Target::Winterfell)
         .unwrap();
 
-    let expected = expect_file!["../system/system.rs"];
+    let expected = expect_file!["../computed_indices/computed_indices_simple.rs"];
     expected.assert_eq(&generated_air);
 }
 
 #[test]
-fn bitwise() {
-    let generated_air = Test::new("tests/bitwise/bitwise.air".to_string())
+fn constant_in_range() {
+    let generated_air = Test::new("tests/constant_in_range/constant_in_range.air".to_string())
         .transpile(Target::Winterfell)
         .unwrap();
 
-    let expected = expect_file!["../bitwise/bitwise.rs"];
+    let expected = expect_file!["../constant_in_range/constant_in_range.rs"];
     expected.assert_eq(&generated_air);
 }
 
@@ -125,12 +124,21 @@ fn constants() {
 }
 
 #[test]
-fn constant_in_range() {
-    let generated_air = Test::new("tests/constant_in_range/constant_in_range.air".to_string())
-        .transpile(Target::Winterfell)
-        .unwrap();
+fn constraint_comprehension() {
+    let generated_air =
+        Test::new("tests/constraint_comprehension/constraint_comprehension.air".to_string())
+            .transpile(Target::Winterfell)
+            .unwrap();
 
-    let expected = expect_file!["../constant_in_range/constant_in_range.rs"];
+    let expected = expect_file!["../constraint_comprehension/constraint_comprehension.rs"];
+    expected.assert_eq(&generated_air);
+
+    let generated_air =
+        Test::new("tests/constraint_comprehension/cc_with_evaluators.air".to_string())
+            .transpile(Target::Winterfell)
+            .unwrap();
+
+    let expected = expect_file!["../constraint_comprehension/constraint_comprehension.rs"];
     expected.assert_eq(&generated_air);
 }
 
@@ -154,28 +162,6 @@ fn fibonacci() {
     expected.assert_eq(&generated_air);
 }
 
-#[test]
-fn functions_simple() {
-    let generated_air = Test::new("tests/functions/functions_simple.air".to_string())
-        .transpile(Target::Winterfell)
-        .unwrap();
-
-    let expected = expect_file!["../functions/functions_simple.rs"];
-    expected.assert_eq(&generated_air);
-}
-
-#[test]
-fn functions_simple_inlined() {
-    // make sure that the constraints generated using inlined functions are the same as the ones
-    // generated using regular functions
-    let generated_air = Test::new("tests/functions/inlined_functions_simple.air".to_string())
-        .transpile(Target::Winterfell)
-        .unwrap();
-
-    let expected = expect_file!["../functions/functions_simple.rs"];
-    expected.assert_eq(&generated_air);
-}
-
 #[test]
 fn functions_complex() {
     let generated_air = Test::new("tests/functions/functions_complex.air".to_string())
@@ -187,22 +173,24 @@ fn functions_complex() {
 }
 
 #[test]
-fn variables() {
-    let generated_air = Test::new("tests/variables/variables.air".to_string())
+fn functions_simple() {
+    let generated_air = Test::new("tests/functions/functions_simple.air".to_string())
         .transpile(Target::Winterfell)
         .unwrap();
 
-    let expected = expect_file!["../variables/variables.rs"];
+    let expected = expect_file!["../functions/functions_simple.rs"];
     expected.assert_eq(&generated_air);
 }
 
 #[test]
-fn trace_col_groups() {
-    let generated_air = Test::new("tests/trace_col_groups/trace_col_groups.air".to_string())
+fn functions_simple_inlined() {
+    // make sure that the constraints generated using inlined functions are the same as the ones
+    // generated using regular functions
+    let generated_air = Test::new("tests/functions/inlined_functions_simple.air".to_string())
         .transpile(Target::Winterfell)
         .unwrap();
 
-    let expected = expect_file!["../trace_col_groups/trace_col_groups.rs"];
+    let expected = expect_file!["../functions/functions_simple.rs"];
     expected.assert_eq(&generated_air);
 }
 
@@ -248,6 +236,26 @@ fn list_folding() {
     expected.assert_eq(&generated_air);
 }
 
+#[test]
+fn periodic_columns() {
+    let generated_air = Test::new("tests/periodic_columns/periodic_columns.air".to_string())
+        .transpile(Target::Winterfell)
+        .unwrap();
+
+    let expected = expect_file!["../periodic_columns/periodic_columns.rs"];
+    expected.assert_eq(&generated_air);
+}
+
+#[test]
+fn pub_inputs() {
+    let generated_air = Test::new("tests/pub_inputs/pub_inputs.air".to_string())
+        .transpile(Target::Winterfell)
+        .unwrap();
+
+    let expected = expect_file!["../pub_inputs/pub_inputs.rs"];
+    expected.assert_eq(&generated_air);
+}
+
 #[test]
 fn selectors() {
     let generated_air = Test::new("tests/selectors/selectors.air".to_string())
@@ -297,20 +305,31 @@ fn selectors_combine_with_list_comprehensions() {
 }
 
 #[test]
-fn constraint_comprehension() {
-    let generated_air =
-        Test::new("tests/constraint_comprehension/constraint_comprehension.air".to_string())
-            .transpile(Target::Winterfell)
-            .unwrap();
+fn system() {
+    let generated_air = Test::new("tests/system/system.air".to_string())
+        .transpile(Target::Winterfell)
+        .unwrap();
 
-    let expected = expect_file!["../constraint_comprehension/constraint_comprehension.rs"];
+    let expected = expect_file!["../system/system.rs"];
     expected.assert_eq(&generated_air);
+}
 
-    let generated_air =
-        Test::new("tests/constraint_comprehension/cc_with_evaluators.air".to_string())
-            .transpile(Target::Winterfell)
-            .unwrap();
+#[test]
+fn trace_col_groups() {
+    let generated_air = Test::new("tests/trace_col_groups/trace_col_groups.air".to_string())
+        .transpile(Target::Winterfell)
+        .unwrap();
 
-    let expected = expect_file!["../constraint_comprehension/constraint_comprehension.rs"];
+    let expected = expect_file!["../trace_col_groups/trace_col_groups.rs"];
+    expected.assert_eq(&generated_air);
+}
+
+#[test]
+fn variables() {
+    let generated_air = Test::new("tests/variables/variables.air".to_string())
+        .transpile(Target::Winterfell)
+        .unwrap();
+
+    let expected = expect_file!["../variables/variables.rs"];
     expected.assert_eq(&generated_air);
 }

air-script/tests/computed_indices/computed_indices_complex.air

@@ -0,0 +1,52 @@
+def ComputedIndicesAir
+
+const MDS = [
+  [1, 2],
+  [2, 3],
+  [3, 4]
+];
+
+trace_columns {
+    main: [a[2]],
+}
+
+public_inputs {
+    input: [1],
+}
+
+fn double(a: felt) -> felt {
+    let x = 3 * a;
+    let y = a;
+    return x - y;
+}
+
+boundary_constraints {
+    enf a[0].first = 0;
+}
+
+integrity_constraints {
+    let state_1 = [1, 1];
+    let state_2 = [2, 0];
+
+    # [1+2, 2+3, 3+4]
+    # [3, 5, 7]
+    let vec_1 = apply_mds(state_1);
+
+    # [2, 4, 6]
+    let vec_2 = apply_mds(state_2);
+
+    # 2 * 2 - 4
+    let x = double(2) - vec_2[1];
+
+    # vec_2[0] = 2
+    let y = vec_2[x];
+
+    # vec_1[2] = 7
+    let z = vec_1[y];
+
+    enf a[1] = double(z);
+}
+
+fn apply_mds(state: felt[2]) -> felt[3] {
+    return [sum([s * m for (s, m) in (state, mds_row)]) for mds_row in MDS];
+}