Allow support for computed indices #444

Leo-Besancon · 2025-08-14T09:54:46Z

This PR closes #245

Task list:

Allow computed indices in AST
- Replace e.g. AccessType::Index(usize) by AccessType::Index(Box<ScalarExpr>) in parser
- Adapt AST's SemanticAnalysis (many thing we cannot resolve early)
- Adapt AST's ConstantPropagation
Allow computed indices in MIR
- Update Accessor types to reference MIR nodes in AccessType, and update lowering from AST
- Update MIR's Unrolling pass
- Update MIR's ConstantPropagation pass
Add E2E tests, containing for example this list comprehension:

let y = [j * vec[j + 1] for j in 0..4];

Add parser tests
Add MIR tests
Potentially more complex E2E tests
Check that there are no regression in semantic analysis assumptions
We should integrate this PR with the new typing system, to ensure that indices are integers and not felts

Note: This PR contains a fix for Mir's Inlining that has been raised by the computed_indices_complex.rs E2E test. Because of this fix, we need to call the Inlining multiple times (until there are no more Call nodes in the graph). Maybe there is a better way to handle this (for instance, keeping track when we duplicate a Call node and add it to the working_stack), but I haven't thought too much about it for now.

Question: @adr1anh @bobbinth Should we also handle computed slice bounds? For now, slices are assumed constant (either an integer or a reference to a constant symbol) and are resolved early. I still handle them in the MIR, but the are constant by that point, so I should refactor the code to either ignore slices in the MIR, or handle them fully as computed expressions.
It may look something like this (though it may not be the best example to showcase usage..):

    let x = [ sum([2^i for i in j..j+3]) for j in 0..3 ]
    # x = [ sum(1, 2, 4), sum(2, 4, 8), sum(4, 8, 16) ]
    # x = [ 7, 14, 28]

Fails for now

Leo-Besancon · 2025-09-02T08:28:52Z

Hello @adr1anh, thank you for the review!

I think I've addressed all your comments for now.

I've also created this issue for one of the FIXME that should probably not be handled in this PR: #452

Al-Kindi-0 · 2025-09-23T12:59:24Z

@Leo-Besancon , is this ready? Or would you like an extra round of reviews?
It would be great to get this merged asap as it will simplify some of the work during the migration of VM constraints to the updated syntax.

Leo-Besancon · 2025-09-23T13:10:23Z

@Leo-Besancon , is this ready? Or would you like an extra round of reviews? It would be great to get this merged asap as it will simplify some of the work during the migration of VM constraints to the updated syntax.

Hi! This PR is ready on my side, but it needs additional reviews, especially as it touches various parts of the code (parser, AST and MIR).
The best order to go forward would be:

Review and merge this PR into next
Merge next into Separation between Felt, Bool, and Uint types #423
Add typing checks to computed indices in Separation between Felt, Bool, and Uint types #423
Review and merge Separation between Felt, Bool, and Uint types #423 into next

Then, the v0.5 Milestone would be done and the syntax isn't likely to change much (aside from the multi trace issue).

bobbinth

Looks good! Thank you! This is a very light review from me - I didn't really review the parser changes, and even the MIR part, I reviewed sporadically. I did leave a few questions/comments inline though.

bobbinth · 2025-09-29T07:33:53Z

mir/src/passes/mod.rs

+    match value.borrow().deref() {
+        Op::Value(Value {
+            value:
+                SpannedMirValue {
+                    value: MirValue::Constant(ConstantValue::Felt(c)),
+                    ..
+                },
+            ..
+        }) => Some(*c),
+        _ => None,
+    }


I think it would be cleaner to add a method on Value for this - something like Value::get_inner_constant(). Then here it would be:

match value.borrow().deref() { Op::Value(value) => value.get_inner_constant(), _ => None, }

We can also then re-use this code in the other get_inner_const() function (in constant propagation). In fact, I wonder if it may make sense to combine these two methods.

In fact, I wonder if it may make sense to combine these two methods.
I tried it but it seemed to break the added test (https://github.com/0xMiden/air-script/pull/444/files/88bb70ce0b281ab3593e9675ba7e07e6215f969b#diff-0ac8abbe17a8a361edfcb0a3a06c7108d5c316885eebc406d0669a3b5754a821)

For now Value's impl is reused in both constant propagation and translation from ast to Mir in 6e734dc

bobbinth · 2025-09-29T07:54:21Z

air-script/tests/computed_indices/computed_indices_complex.air

+fn apply_mds(state: felt[2]) -> felt[3] {
+    return [sum([s * m for (s, m) in (state, mds_row)]) for mds_row in MDS];
+}


This is the part that was not possible before, right?

This may not be the right place, but I imagine this is actually the only place where we'd ever use matrices. For the rare cases where we want to evaluate a matrix (only hash functions I think), we could just manually implement it. If so it might be worth removing matrices entirely. WDYT?

This is the part that was not possible before, right?

Yes when coupled with a function. I think it was ok without a function before because it was computed before semantic analysis.

If so it might be worth removing matrices entirely. WDYT?

I agree, its already in the todo list for the frontend refactor. Ill gather those in a separate issue.

This is the part that was not possible before, right?

Not exactly, the apply_mds function was already tested in https://github.com/0xMiden/air-script/blob/allow-computed-indices-upstream/air-script/tests/list_comprehension/list_comprehension_nested.air
In the list_comprehension_nested test, it is indeed applied to the state as you suggested below, so we can check the rust code.

What is new is to be able to index into the vec_1 and vec_2 list comprehensions with indices whose can only be computed late into the compilation process (for instance, because they rely on list comprehensions)

In this test (computed_indices_complex), I needed a vector that would not be folded as a constant during AST's constant propagation, otherwise I would not be able to test the computed indices targetting list comprehensions properly.

I'll add some comments to make what this test does clearer.

bobbinth · 2025-09-29T07:58:59Z

air-script/tests/computed_indices/computed_indices_complex.air

+    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);


I would maybe simplify this test and at the same time make it a bit more robust. Specifically:

Let's use trace columns as the state.

Maybe we can apply MDS once and then in Rust code we'd see the actual formula of how it is getting applied (I think).

I've adapted the air-script of the test in 82cd993 but it does not compile as is, I'm investigating why that is.
It seems we extract the constant from the wrong value, extracting the vec_1 instead of 1 or report the error incorrectly here:

air-script/air-script/tests/computed_indices/computed_indices_complex.air

Line 39 in 82cd993

let x = double(2) - vec_2[1];

I would maybe simplify this test and at the same time make it a bit more robust. Specifically:

* Let's use trace columns as the `state`. * Maybe we can apply MDS once and then in Rust code we'd see the actual formula of how it is getting applied (I think).

In the latest commit, I've made it so that one MDS is kept on a constant vector to be able to produce constant indices, and use it to index into the second MDS that does rely on the state.
I've also added comments on what is expected / happening in the test, please tell me if it's better!

bobbinth · 2025-09-29T08:03:16Z

mir/src/passes/mod.rs

+/// stage.
+pub fn handle_accessor_visit(
+    accessor: Link<Op>,
+    compute_indices: bool,


nit (feel free to ignore): I would maybe invert compute_indices into something like expect_constant_indices. May make the meaning a bit more clear (at least that's how we describe it in the doc comments).

Renamed in 90b5bdc and everywhere else in c8adbf5

bobbinth · 2025-09-29T08:12:38Z

parser/src/ast/trace.rs

+    /// The access type associated to this TraceBinding
+    pub access: Option<AccessType>,


What this filed mean here? That is, in what situations would a trace binding have access vs. when it could be None?

Im not sure why it is this way or if it could be without the Option, so I'm only interpreting here based on the code.

It seems it allows setting it after creation, potentially handling the difference between an unset access_type and AccessType::Default, distinguishing between an error (a mismatch) and a missing access_type.

Noting so we can revisit with @Leo-Besancon based on his feedback

What this filed mean here? That is, in what situations would a trace binding have access vs. when it could be None?

Indeed the Option is unnecessary as the None case can just be handled the same as the Default case. I've removed it in the latest commit

bobbinth · 2025-09-29T08:14:53Z

parser/src/ast/trace.rs

+            Some(AccessType::Default) => self.ty,
+            Some(AccessType::Slice(range_expr)) => Type::Vector(range_expr.to_slice_range().len()),
+            Some(AccessType::Index(_)) => Type::Felt,
+            Some(AccessType::Matrix(..)) => Type::Felt,


Under what circumstances could TraceBinding have type Matrix? Should this be unreachable!?

Made it unreachable in abbf525

Al-Kindi-0

LGTM (given my understanding of this part of the codebase)

Al-Kindi-0 · 2025-10-09T11:33:22Z

mir/src/passes/mod.rs

+}
+
+/// Handle the visit of an accessor node, used for both Unrolling and ConstantPropagation passes
+/// The `compute_indices` bool indicates whether we the indices need to be known constant at this


this needs rewording

Reworded in c8adbf5, let me know if it needs more adjustments.

rename handle_accessor_visit arg: compute_indices -> expect_constant_indices

Al-Kindi-0 · 2025-10-15T13:20:34Z

@Leo-Besancon , is this good to go?

Leo-Besancon · 2025-10-15T13:23:09Z

@Leo-Besancon , is this good to go?

Hi! Yes, if you guys are happy with our latest replies, it should be good to merge, I think we've addressed everything!

Al-Kindi-0 · 2025-10-15T15:15:02Z

Thank you!

Leo-Besancon changed the base branch from next to add_constant_propagation_mir_pass_upstream August 14, 2025 09:55

Leo-Besancon force-pushed the add_constant_propagation_mir_pass_upstream branch from 31ceb57 to 2dd6612 Compare August 18, 2025 05:56

Leo-Besancon force-pushed the allow-computed-indices-upstream branch from dd7f771 to 6fd045d Compare August 18, 2025 06:01

Leo-Besancon mentioned this pull request Aug 18, 2025

Add Constant Propagation pass after MIR passes #439

Merged

10 tasks

Leo-Besancon force-pushed the allow-computed-indices-upstream branch 4 times, most recently from ef562f2 to 33e2e4f Compare August 19, 2025 13:39

Leo-Besancon marked this pull request as ready for review August 19, 2025 14:03

adr1anh self-requested a review August 19, 2025 14:43

Leo-Besancon mentioned this pull request Aug 20, 2025

Consider moving MIR's constant propagation to the AIR #446

Open

Leo-Besancon added 13 commits August 20, 2025 07:53

feat(parser): Allow computed indices in AST

c56df8f

feat(mir): Allow computed indices in MIR

426ffd2

tests: Add computed indices E2E test

124ca12

tests(parser): Add computed indices parser tests

5456c4b

fix(parser): fix binding type access on non constant index

00df247

tests(mir): add mir tests for computed indices

b3bc934

chore: update CHANGELOG.md

52dbd20

tests: Update E2E test to have correct trace

2891862

chore: use chained if lets

1202dec

tests: add computed_indices_complex E2E test

6634e61

Fails for now

tests: Sort winterfell tests alphabetically

f9f7689

fix: duplicate_node for argument nodes in calls

d4a07ab

fix: Run Mir inlining until fixed point

8af9c70

Leo-Besancon force-pushed the allow-computed-indices-upstream branch from 33e2e4f to 8af9c70 Compare August 20, 2025 05:54

Leo-Besancon changed the base branch from add_constant_propagation_mir_pass_upstream to next August 20, 2025 05:54

Leo-Besancon added a commit to massalabs/air-script that referenced this pull request Aug 22, 2025

refactor: Improve first pass based on 0xMiden#444 review comments

214faf7

Leo-Besancon added a commit that referenced this pull request Aug 22, 2025

refactor: Improve first pass based on #444 review comments

157a905

Leo-Besancon added 2 commits August 24, 2025 19:37

Merge branch 'next' into allow-computed-indices-upstream

62f6ddd

refactor: Handle range as constant, remove their propagation into MIR

a150733

Leo-Besancon mentioned this pull request Sep 2, 2025

[Mir] Consider using direct computed indices for TraceAccess and PublicInputAccess #452

Open

Merge branch 'next' into allow-computed-indices-upstream

2b0de7d

bobbinth requested review from bitwalker and bobbinth September 5, 2025 15:06

bobbinth mentioned this pull request Sep 7, 2025

Add support for more general boundary constraints #433

Open

Leo-Besancon added 3 commits September 22, 2025 09:24

Merge branch 'next' into allow-computed-indices-upstream

85d4dc2

fix: small fixes after merge

b40cca5

refactor: refactor MIR's constant_propagation visit_node

88bb70c

bobbinth reviewed Sep 29, 2025

View reviewed changes

adr1anh requested a review from huitseeker October 9, 2025 07:58

Al-Kindi-0 approved these changes Oct 9, 2025

View reviewed changes

Soulthym added 4 commits October 9, 2025 16:53

feat(mir): Value::get_inner_const

6e734dc

fix(mir): rename handle_accessor_visit arg

90b5bdc

rename handle_accessor_visit arg: compute_indices -> expect_constant_indices

fix(parser): make AccessType::Matrix unreachable for TraceBinding

abbf525

fix(mir): improve docs and unify expect_constant_indices accross APIs.

c8adbf5

Soulthym force-pushed the allow-computed-indices-upstream branch from 9d75de1 to c8adbf5 Compare October 9, 2025 16:03

Soulthym and others added 6 commits October 9, 2025 18:39

test(mir): adapt computed_indices_complex to use the trace as state

82cd993

tests: rework computed indices test based on review, add comments

770b6b6

chore: lint fix

c48bfd4

Merge branch 'next' into allow-computed-indices-upstream

a988ae5

refactor: remove unneeded Option

d7b557f

Merge branch 'next' into allow-computed-indices-upstream

02610a4

Al-Kindi-0 merged commit 9d69e2d into next Oct 15, 2025
11 checks passed

Al-Kindi-0 deleted the allow-computed-indices-upstream branch October 15, 2025 15:15

		/// The access type associated to this TraceBinding
		pub access: Option<AccessType>,

Allow support for computed indices #444

Allow support for computed indices #444

Uh oh!

Conversation

Leo-Besancon commented Aug 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Leo-Besancon commented Sep 2, 2025

Uh oh!

Al-Kindi-0 commented Sep 23, 2025

Uh oh!

Leo-Besancon commented Sep 23, 2025

Uh oh!

bobbinth left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Soulthym Oct 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Soulthym Oct 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Al-Kindi-0 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Soulthym Oct 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Al-Kindi-0 commented Oct 15, 2025

Uh oh!

Leo-Besancon commented Oct 15, 2025

Uh oh!

Al-Kindi-0 commented Oct 15, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Leo-Besancon commented Aug 14, 2025 •

edited

Loading

Soulthym Oct 9, 2025 •

edited

Loading

Soulthym Oct 9, 2025 •

edited

Loading

Soulthym Oct 9, 2025 •

edited

Loading