See if we can use MapStmtExpr to simplify `transformExpr` in LiftImperativeExpression

[keyboardDrummer]

See if we can use MapStmtExpr to simplify transformExpr in LiftImperativeExpression

[keyboardDrummer-bot]

I investigated whether mapStmtExprM from MapStmtExpr.lean can simplify transformExpr in LiftImperativeExpressions.lean. My conclusion is that it cannot be used here in a straightforward way, because the two traversals have fundamentally different requirements:

Why mapStmtExprM doesn't fit transformExpr:

1. Right-to-left sibling processing: transformExpr processes PrimitiveOp and StaticCall arguments in reverse order (args.reverse.mapM transformExpr). This is essential because assignments create snapshots that affect how expressions to their left are substituted. mapStmtExprM processes children left-to-right.

2. Stateful interaction between siblings: The substitution map (subst) is updated as each sibling is processed. The result of processing one argument affects how the next (leftward) argument is transformed. mapStmtExprM applies f only after all children are rebuilt, so there's no way to thread state between sibling processing.

3. Custom recursion for IfThenElse: When branches contain assignments, transformExpr saves the outer state, processes each branch with a fresh empty state, then restores the outer state. This is incompatible with mapStmtExprM's uniform recursion.

4. Statement lifting as a side effect: transformExpr prepends statements during traversal. The order in which prepends are emitted depends on the right-to-left traversal order, which mapStmtExprM cannot replicate.

5. Selective non-recursion: Some nodes (like Assign) don't recurse into their value — they lift the original value directly. mapStmtExprM always recurses into all children.

Could transformStmt use it?

transformStmt also has custom per-constructor logic (handling prepends, resetting substitutions, special-casing StaticCall in Assign RHS), so it similarly can't delegate to mapStmtExprM.

Summary: The imperative-expression lifter's traversal is inherently order-dependent and stateful in ways that a generic bottom-up map cannot express. The current mutual recursion between transformExpr and transformStmt is the appropriate structure for this pass.

Let me know if you'd like me to explore a different angle — for example, whether a specialized right-to-left monadic traversal utility could reduce some of the boilerplate, or whether there's a subset of the logic you had in mind.

[keyboardDrummer]

How would we need to generalize mapStmtExpr to enable using it for LiftImperativeExpression?

• We could give it an option to do right to left mapping.
• Do we need an option to override the recursion into children entirely?

[keyboardDrummer]

@keyboardDrummer-bot please address my previous comment

[keyboardDrummer-bot]

Here's my analysis of what generalizations mapStmtExprM would need:

1. Right-to-left traversal option — yes, this would help

Adding a parameter like (reverseChildren : Bool := false) that reverses the order in which sibling lists (e.g. args) are traversed and then un-reverses the results. This directly addresses the args.reverse.mapM transformExpr pattern in transformExpr for PrimitiveOp, StaticCall, and Block.

2. Override recursion into children — yes, but it needs to be per-constructor

A single "don't recurse" flag isn't enough. What transformExpr actually needs is:

• Custom handling per constructor (not just post-processing): For Assign, IfThenElse (with assignments), StaticCall (imperative), Block, and Var (.Declare _), the pass doesn't want the generic recursion at all — it wants to run its own logic. For leaves and simple constructors, the generic recursion is fine.
• This suggests a "pre-hook" or "override" function that runs before recursion and can return either:
  • some result → skip recursion entirely, use this result
  • none → proceed with normal recursive traversal, then apply the post-hook f

So the signature would look something like:

def mapStmtExprM [Monad m]
    (pre : StmtExprMd → m (Option StmtExprMd))  -- override recursion
    (post : StmtExprMd → m StmtExprMd)           -- applied after rebuild
    (reverseChildren : Bool := false)
    (expr : StmtExprMd) : m StmtExprMd

When pre returns some e, the traversal returns e directly (the override is responsible for any sub-recursion it wants). When pre returns none, the generic recursion + post proceeds as today.

Would this actually simplify LiftImperativeExpressions?

Partially. The constructors that don't need custom logic (LiteralInt, LiteralBool, AsType, IsType, ReferenceEquals, Old, Fresh, etc.) are already trivial one-liners (| _ => return expr) in transformExpr. The real complexity lives in the custom cases (Assign, IfThenElse, StaticCall, Block, Var), which would all go into the pre hook and remain essentially unchanged.

The net savings would be:

• Eliminating the explicit PrimitiveOp arm (it becomes just the generic right-to-left traversal)
• Eliminating the leaf catch-all
• Not needing the termination_by / decreasing_by proof in transformExpr (it would piggyback on mapStmtExprM's termination proof)

The transformStmt function wouldn't benefit much since it has per-constructor logic for nearly every arm.

Summary: A pre/post hook design with a reverseChildren flag would let transformExpr delegate its "boring" cases to the generic traversal while keeping full control over the interesting ones. The main win is reduced boilerplate and not needing a separate termination proof. Want me to prototype this generalization?

[keyboardDrummer]

@keyboardDrummer-bot try to implement the above. I think it's important that the passes don't need to think about language features that they don't relate to, like PrimitiveOp in this case.

[keyboardDrummer-bot]

I've implemented this in PR #44. Here's a summary:

What was done:

• Added mapStmtExprChildrenM to MapStmtExpr.lean — maps a function over the immediate children of a node with optional right-to-left traversal
• Also added mapStmtExprPrePostM — a full traversal with pre/post hooks (useful for other passes)
• Refactored transformExpr in LiftImperativeExpressions.lean to use mapStmtExprChildrenM in the catch-all, eliminating the explicit PrimitiveOp arm

The catch-all is now:

| _ => mapStmtExprChildrenM transformExpr (reverseChildren := true) expr

New constructors added to StmtExpr will be traversed automatically without updating transformExpr.

Trade-off: transformExpr is marked partial because Lean's termination checker can't see through mapStmtExprChildrenM to verify that transformExpr is only called on structurally smaller values. The alternative (keeping the termination proof) requires listing every constructor with .attach proofs, which defeats the purpose of the generic catch-all.