Fix a bug in assignment extraction logic, and enable heap mutations in expression positions #358

keyboardDrummer · 2026-01-30T13:17:14Z

Changes

Fix bug in assignment lifting pass that showed when doing subsequent assignments to the same variable.
Stop using globals in heap parameterization, and do heap parameterization before assignment lifting, so the assignment lifting pass can work on the heap assignments

Testing

For (1), add test-case to T2_ImpureExpressions
For (2), add test-case to T1_MutableFields

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

andrewmwells-amazon · 2026-01-30T15:20:24Z

Strata/Languages/Laurel/HeapParameterization.lean


 def heapTransformProcedure (proc : Procedure) : TransformM Procedure := do
-  let heapName := "$heap"
+  let heapInName := "heap_in"


We may want to use some names that aren't legal variable names for user-written code.

Sounds good, so then heap_in won't be a legal variable name?

Given that Laurel code is generally instantiated by being deserialized, the calling side will be able to choose any string they want for identifiers, including things with "$", regardless of what the Laurel parser allows.

If the calling side wants the Laurel to be parseable after printing, then they'll have to avoid using Laurel reserved keywords, in which we can include "heap_in".

But I see your point. It's easier to tell the calling side "don't use $ in identifiers" then to come with a list of disallowed words (although we'll have to do that anywhere for some syntactic things), and it's easier for us to check as well.

joscoh · 2026-01-30T15:11:38Z

Strata/Languages/Laurel/HeapParameterization.lean

+
+3. Procedure calls are transformed:
+   - Calls to heap-writing procedures in expressions:
+     `f()` => `(var freshVar: type; freshVar, heap := f(heap); freshVar)`


Based on what you wrote below, should this be heap, freshVar := f(heap)?

You're right. I'll fix the comment, and let me also add the optimization that's shown in the second - of this list, but which wasn't part of the implementation.

joscoh · 2026-01-30T15:14:38Z

Strata/Languages/Laurel/Laurel.lean

-  | Assign (target : StmtExpr) (value : StmtExpr) (md : Imperative.MetaData Core.Expression)
+  /- Assign is only allowed in an impure context.
+     For single target assignments, use a single-element list.
+     Multiple targets are only allowed when the value is a StaticCall to a procedure


Eventually, this should be enforced by a WF predicate or similar.

It would be nice if we could already construct the proof when the AST node is constructed, and then store node specific parts with the nodes, but given that it checking that the number of return values matches the number of target values requires resolving which procedure is called, is seems difficult to do that at construction time.

If we don't store the proof with the AST, would the WF proof then be a separate tree with the same shape as the AST, but with proofs instead of values?

joscoh · 2026-01-30T15:17:58Z

Strata/Languages/Laurel/LiftExpressionAssignments.lean

+      for target in targets do
+        match target with
+        | .Identifier varName =>
+            let snapshotName ← SequenceM.freshTempFor varName


How do you ensure there is not already a variable with this name?

I'll add $ as a prefix. Currently there is no check that prevents front-ends from using $ in identifiers, but we can add that as part of the WF-ness check.

Would be cool if we can find name conflict bugs when we try to prove the transformation maintains semantics equivalence.

I think an alternative approach would be to let the declarations have both a displayName (for debugging) and a generated numeric identifier that's unique across the program, then we wouldn't have to add $.

joscoh · 2026-01-30T15:18:54Z

Strata/Languages/Laurel/LiftExpressionAssignments.lean

 Returns the transformed expression with assignments replaced by variable references.
 -/
-def transformExpr (expr : StmtExpr) : SequenceM StmtExpr := do
+partial def transformExpr (expr : StmtExpr) : SequenceM StmtExpr := do


Why is this partial now?

The changes make it so that the termination proof is no longer automatic. I spent over 60 minutes trying to provide a proof, but didn't get very far. A big part is that I'm wholly unfamiliar with writing these Lean proofs.

When I add the termination metrics to the three involved recursive functions, it shows that it fails to prove some recursive calls inside processBlock. I have tried to add some have ... by sorrys there to try and get it to pass, but I end up with goals like this:

unsolved goals expr : StmtExpr metadata : Option Identifier head : StmtExpr tail : List StmtExpr h : sizeOf tail < sizeOf (head :: tail) ⊢ sizeOf tail < sizeOf (head :: tail)

Note the last two lines. I would have thought that a tactic like

decreasing_by all_goals simp_wf <;> omega

Would take care of such trivial proofs, but apparently it does not.

There are similar very simple cases where it fails, like:

h : sizeOf head < sizeOf remStmts ⊢ sizeOf head < sizeOf (x✝ :: tail)

I think that x✝ :: tail results from pattern matching on remStmts, but for some reason it can't figure out that h is equal to the goal.

Could you help me figure out how to write this proof? Should be simple given that each recursive call is on a smaller argument.

joscoh · 2026-01-30T15:21:02Z

Strata/Languages/Laurel/LiftExpressionAssignments.lean

+      | [] => return .Block [] metadata
+      | [last] => transformExpr last
+      | _ =>
+          -- Process all but the last statement


I am unclear on why you needed to change the let rec next ... structure. The head :: tail case there already handles all but the last statement (which matches the first case).

Was unnecessary, although I prefer the for over the recursive call. Undid the change.

joscoh · 2026-01-30T15:22:07Z

Strata/Languages/Laurel/LiftExpressionAssignments.lean

+                -- For assignments in block context, we need to set snapshots
+                -- so that subsequent expressions see the correct values
+                -- IMPORTANT: Use transformTarget for targets (no snapshot substitution)
+                -- and transformExpr for values (with snapshot substitution)


I don't quite understand what this comment means. What is a snapshot exactly?

The value of a mutable variable at a particular at the current execution point. I've updated the comment so it's hopefully more clear

joscoh · 2026-01-30T15:23:10Z

Strata/Languages/Laurel/LiftExpressionAssignments.lean

 Returns a list of statements (the original one may be split into multiple).
 -/
-def transformStmt (stmt : StmtExpr) : SequenceM (List StmtExpr) := do
+partial def transformStmt (stmt : StmtExpr) : SequenceM (List StmtExpr) := do


Same thing, we should avoid making things partial unless there is a good reason.

MikaelMayer

I think it would be worth it to prove the termination of the function you annotated as partial. We have other cases in the codebase which have proof templates for proving termination easily, I hope you can apply one of them here.

MikaelMayer · 2026-01-30T15:19:11Z

Strata/Languages/Laurel/HeapParameterization.lean

+          let postcond' ← heapTransformExpr heapInName postcond
+          let impl' ← impl.mapM (heapTransformExpr heapInName)
+          let modif' ← modif.mapM (heapTransformExpr heapInName)
+          pure (.Opaque postcond' impl' modif')


If you ever wanted to avoid parentheses, you could use the following idiom

Suggested change

pure (.Opaque postcond' impl' modif')

return .Opaque postcond' impl' modif'

return here would bind with the do of this line def heapTransformProcedure (proc : Procedure) : TransformM Procedure := do, so it changes the meaning, but good reminder, thanks!

MikaelMayer · 2026-01-30T15:21:23Z

Strata/Languages/Laurel/LaurelFormat.lean


 def formatStmtExpr (s:StmtExpr) : Format :=
-  match h: s with
+  match s with


Note: I don't know if it applies here but if you ever need to do termination proofs and naming the matches causes a spurious "unused variable" despite the variable being used in automatic proofs, you could also use '_' for the variable name.

Suggested change

match s with

match _: s with

MikaelMayer · 2026-01-30T15:24:18Z

Strata/Languages/Laurel/LiftExpressionAssignments.lean

 Returns a list of statements (the original one may be split into multiple).
 -/
-def transformStmt (stmt : StmtExpr) : SequenceM (List StmtExpr) := do
+partial def transformStmt (stmt : StmtExpr) : SequenceM (List StmtExpr) := do


Can you try to prove termination?

keyboardDrummer added 11 commits January 29, 2026 14:50

Allow outline MCP server

e7b71aa

Stop using globals in heap parameterization

b1cfa22

Fix termination proof

bb16c5c

Add test-case

43e419d

Fixes

e46e315

Workaround parsing issue

f322454

Fix related to subsequent assignments in expressions

821021a

Add partials

6beab5a

Separate counters for different variables

6e2ecbb

Fixes

040dc13

Remove dbg_trace

524e965

keyboardDrummer changed the title ~~Enable extrancting heap mutations from expressions~~ Fix a bug in assignment extraction logic, and enable heap mutations in expression positions Jan 30, 2026

Merge remote-tracking branch 'origin/main' into stopUsingGlobals

dd2a6c1

keyboardDrummer marked this pull request as ready for review January 30, 2026 13:55

keyboardDrummer requested a review from a team as a code owner January 30, 2026 13:55

keyboardDrummer enabled auto-merge January 30, 2026 13:57

keyboardDrummer requested a review from joscoh January 30, 2026 14:27

keyboardDrummer mentioned this pull request Jan 30, 2026

Support assignment in conditional parts of an expression #359

Open

andrewmwells-amazon previously approved these changes Jan 30, 2026

View reviewed changes

joscoh reviewed Jan 30, 2026

View reviewed changes

MikaelMayer reviewed Jan 30, 2026

View reviewed changes

keyboardDrummer added 2 commits February 2, 2026 11:18

Code review

c896d1f

Improve comment

dfaa58a

keyboardDrummer dismissed andrewmwells-amazon’s stale review via dfaa58a February 2, 2026 11:25

Merge branch 'main' into stopUsingGlobals

4582a97

keyboardDrummer requested a review from joscoh February 2, 2026 11:46

	pure (.Opaque postcond' impl' modif')
	return .Opaque postcond' impl' modif'

Fix a bug in assignment extraction logic, and enable heap mutations in expression positions #358

Are you sure you want to change the base?

Fix a bug in assignment extraction logic, and enable heap mutations in expression positions #358

Conversation

keyboardDrummer commented Jan 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Changes

Testing

Uh oh!

Choose a reason for hiding this comment

Uh oh!

keyboardDrummer Feb 2, 2026 • 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!

keyboardDrummer Feb 2, 2026 • 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!

keyboardDrummer Feb 2, 2026 • 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!

keyboardDrummer Feb 2, 2026 • 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!

MikaelMayer left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

keyboardDrummer Feb 2, 2026 • 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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

keyboardDrummer commented Jan 30, 2026 •

edited

Loading

keyboardDrummer Feb 2, 2026 •

edited

Loading

keyboardDrummer Feb 2, 2026 •

edited

Loading

keyboardDrummer Feb 2, 2026 •

edited

Loading

keyboardDrummer Feb 2, 2026 •

edited

Loading

keyboardDrummer Feb 2, 2026 •

edited

Loading