Use mapStmtExprChildrenM to simplify transformExpr

[keyboardDrummer-bot]

Addresses #40.

Summary

Adds mapStmtExprChildrenM to MapStmtExpr.lean — a utility that maps a function over the immediate children of a StmtExprMd node (one level, no recursion), with an optional reverseChildren flag for right-to-left traversal.

Uses it in LiftImperativeExpressions.lean to replace the explicit PrimitiveOp arm in transformExpr with a generic catch-all:

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

This means the pass no longer needs to enumerate language features it doesn't relate to. If a new StmtExpr constructor is added, it will be traversed automatically without updating transformExpr.

Changes

MapStmtExpr.lean:

• Add mapStmtExprChildrenM: one-level child mapping with optional right-to-left order
• Add mapStmtExprPrePostM: full traversal with pre/post hooks (useful for other passes)

LiftImperativeExpressions.lean:

• Remove explicit PrimitiveOp arm from transformExpr
• Replace catch-all | _ => return expr with mapStmtExprChildrenM transformExpr (reverseChildren := true) expr
• Add explicit | .Assert _ | .Assume _ => return expr since these are statement-like constructs lifted out of blocks and should not have substitutions applied
• Remove mutual block (no longer needed since transformStmt doesn't need mutual recursion with transformExpr)
• Mark transformExpr as partial since mapStmtExprChildrenM is opaque to the termination checker

Trade-off

The use of partial is a trade-off: we lose the termination proof for transformExpr in exchange for not needing to enumerate every constructor. The termination checker cannot see through mapStmtExprChildrenM to verify that transformExpr is only called on structurally smaller values. The original code avoided this by listing every constructor with .attach proofs.

Testing

All 591 test jobs pass.
"

[keyboardDrummer-bot]

I investigated the "Run internal benchmarks of Strata" CI failure. The benchmark job triggers an AWS CodeBuild build that runs internal tests.

Finding: This failure is not caused by the code changes in this PR. The same benchmark job is also failing on other unrelated PRs:

• PR Improve block formatting and diagnostic reporting #38 (formatting changes only) — benchmark failed
• PR Add EliminateReturnStatements and ContractPass #34 (contract pass) — benchmark failed

The benchmark infrastructure appears to have a pre-existing issue (likely AWS credentials, CodeBuild configuration, or the benchmark environment itself). The failure happens quickly (~100 seconds), suggesting it fails during setup or early build rather than during actual test execution.

Local verification:

• lake build ✅ (577 jobs, all successful)
• lake test ✅ (all tests pass; the one failure is an unrelated missing jar file for Java integration tests)
• All #guard_msgs tests pass, confirming the lift pass produces correct output

[keyboardDrummer]

@keyboardDrummer-bot are there other passes that can make use of the new utilities that this PR introduces?

[keyboardDrummer-bot]

Yes, several other passes could benefit from the new utilities introduced here:

mapStmtExprChildrenM (one-level child mapping)

InferHoleTypes.lean — inferExpr manually enumerates every constructor to propagate expectedType through children. Most arms just recurse with the same or a default expected type. The pass could handle only the constructors that need custom type propagation (.Hole, .PrimitiveOp, .StaticCall, .InstanceCall, .IfThenElse, .Block, .Assign, .While, .Return, .Quantifier) and use mapStmtExprChildrenM with a default-type wrapper for the catch-all. However, since almost every constructor needs a different expected type for its children, the benefit is moderate — it would mainly future-proof against new constructors being silently skipped.

HeapParameterization.lean — heapTransformExpr.recurse manually enumerates every constructor (~100 lines of boilerplate for the "just recurse into children" cases like .InstanceCall, .ReferenceEquals, .Quantifier, .Old, .Fresh, .Assigned, .Assume, .ProveBy, .ContractOf). The constructors that need custom logic (.Var (.Field ..), .StaticCall, .Assign, .PrimitiveOp, .AsType, .Block) could be handled explicitly, and the rest could fall through to mapStmtExprChildrenM. This would eliminate ~40 lines and make the pass resilient to new constructors.

FilterPrelude.lean — collectExprNames manually enumerates every constructor to collect referenced names. A generic "visit all children" utility would simplify this, though mapStmtExprChildrenM maps rather than folds, so a separate foldStmtExprChildrenM or using mapStmtExprChildrenM with a StateM that ignores the rebuilt tree would be needed.

mapStmtExprPrePostM (pre/post hook traversal)

HeapParameterization.lean — heapTransformExpr is a top-down recursive traversal with custom logic for specific constructors and generic recursion for the rest. It's a natural fit for mapStmtExprPrePostM where the pre hook handles .Var (.Field ..), .StaticCall, .Assign, .PrimitiveOp, .AsType, and .Block (returning some), and everything else falls through to generic recursion. The complication is the valueUsed parameter threaded through recursion, which doesn't fit the pre/post signature directly — it would need to be carried in a StateM or the utility would need a variant that passes context.

LiftImperativeExpressions.lean — transformStmt could potentially use mapStmtExprPrePostM with a pre that handles the statement-specific constructors (.Assign, .Block, .While, .IfThenElse, .Return, .StaticCall) and delegates expression-position nodes to transformExpr. This is already partially addressed by this PR's use of mapStmtExprChildrenM in transformExpr.

Summary

The highest-value opportunity is HeapParameterization.lean, which has the most boilerplate that mapStmtExprChildrenM could eliminate. The valueUsed threading is a complication but could be handled by making it part of the monad state. InferHoleTypes would benefit moderately for future-proofing. FilterPrelude would need a fold variant rather than a map.