Allow support for computed indices

CHANGELOG.md

@@ -9,6 +9,7 @@
 - Update documentation and tests thereof (#437).
 - Add a constant propagation pass after other mir passes (#439).
 - Removed `TraceSegmentId::index()` and replaced segment indexing with `TraceShape<T>`/`FullTraceShape<T>` across ACE codegen, MIR-to-AIR pass, and constraints (#442).
+- 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,61 @@
+def ComputedIndicesAir
+
+const MDS = [
+  [1, 2],
+  [2, 3],
+  [3, 4]
+];
+
+trace_columns {
+    main: [a[2], s[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;
+}
+
+# Note: 
+# In this test, we aim to test that computed indices work well even if:
+# - The value of the index can only be known late during the compilation process (during MIR's constant propagation)
+# - The 
+integrity_constraints {
+
+    # vec_1 is a list_comprehension that depends on the state
+    # vec_1 = [
+    #    1 * s[0] + 2 * s[1],
+    #    2 * s[0] + 3 * s[1],
+    #    3 * s[0] + 4 * s[1]
+    # ];
+    let vec_1 = apply_mds(s);
+
+    let state_2 = [2, 0];
+    # vec_2 is a list_comprehension that will not get constant-folded early, but will produce constant values
+    let vec_2 = apply_mds(state_2);
+
+    # x will get the value 2 * 2 - 4 = 0
+    let x = double(2) - vec_2[1];
+
+    # y will get the value vec_2[0] = 2
+    let y = vec_2[x];
+
+    # z will then be vec_1[2] = 3 * s[0] + 4 * s[1]
+    let z = vec_1[y];
+
+    # we enforce 3 * s[0] + 4 * s[1] = 0
+    enf z = 0;
+}
+
+# We use apply_mds function to produce a list comprehension that will not get constant-folded during AST
+fn apply_mds(state: felt[2]) -> felt[3] {
+    return [sum([s * m for (s, m) in (state, mds_row)]) for mds_row in MDS];
+}