Implement guarded LICM transformation

lean/Spa/Transformation/Licm.lean

@@ -1,12 +1,13 @@
 import Spa.Analysis.Reaching
+import Spa.Language.Equivalence
 
 /-!
-# Finding loop-invariant assignments (LICM groundwork)
+# Loop-invariant code motion
 
 This wires the **reaching-definitions** analysis (`Spa/Analysis/Reaching.lean`)
-to the AST to *find* — not yet move — assignments inside a `while` loop whose
+to the AST to find assignments inside a `while` loop whose right-hand side
-right-hand side depends only on definitions made *outside* the loop. These are
+depends only on definitions made outside the loop. `licmCandidates` reports
-the candidates a later LICM pass could hoist.
+these assignments; `hoistProgram` moves eligible leading assignments.
 
 The traversal recurses over the plain `Stmt`, threading a `GGraph.Embed` of the
 current subtree's CFG into the program's (`Program.rootEmbed`, then one
@@ -24,9 +25,12 @@ current subtree's CFG into the program's (`Program.rootEmbed`, then one
    comparisons.
 
 If every reaching definition of every RHS variable lies outside the loop, the
-assignment is reported as loop-invariant. This is the first-order check ("all
+assignment is reported as loop-invariant. Hoisting additionally requires the
-reaching definitions outside the loop"); transitive/iterated invariance and the
+assignment to lead the body, its destination to be absent from the guard, and
-actual hoisting are out of scope here.
+no reassignment of that destination in the remaining body. The hoist is guarded
+by the original condition, preserving zero-iteration behavior.
+
+Transitive invariance and motion of non-leading assignments are not implemented.
 -/
 
 namespace Spa
@@ -102,6 +106,50 @@ def licmCandidates (prog : Program) : List (prog.State × prog.State) :=
   (collectCandidates prog none prog.rootStmt prog.rootEmbed).filterMap (fun c =>
     if isInvariant prog c then some (c.encl.loopState, c.assignState) else none)
 
+/-- Candidate for the leading assignment of a loop body. -/
+def headCandidate (prog : Program) (cond : Expr) (x : String) (rhs : Expr) (tail : Stmt)
+    (e : Embed (Stmt.whileLoop cond (.andThen (.basic (.assign x rhs)) tail)).cfg prog.cfg) :
+    Candidate prog :=
+  let body := Stmt.andThen (.basic (.assign x rhs)) tail
+  let be := (Embed.loop body.cfg).trans e
+  let ae := (Embed.sequenceLeft (Stmt.basic (.assign x rhs)).cfg tail.cfg).trans be
+  { encl := { loopState := e.f body.cfg.loopIn, bodyOff := be.off, bodySize := body.cfg.size },
+    assignState := ae.singletonIndex, rhsVars := rhs.vars.sort (· ≤ ·) }
+
+/-- Guard the hoist so that a zero-iteration loop never evaluates the RHS. -/
+def hoistHead (cond : Expr) (x : String) (rhs : Expr) (tail : Stmt) : Stmt :=
+  .ifElse cond (.andThen (.basic (.assign x rhs)) (.whileLoop cond tail)) (.basic .noop)
+
+/-- Hoist a leading invariant assignment when its destination is neither
+reassigned in the remaining body nor read by the guard. -/
+def hoistLoop (prog : Program) (cond : Expr) (body : Stmt)
+    (e : Embed (Stmt.whileLoop cond body).cfg prog.cfg) : Option Stmt :=
+  match body with
+  | .andThen (.basic (.assign x rhs)) tail =>
+      if isInvariant prog (headCandidate prog cond x rhs tail e) &&
+          decide (x ∉ tail.writes ∧ x ∉ cond.vars) then
+        some (hoistHead cond x rhs tail)
+      else none
+  | _ => none
+
+/-- Apply guarded leading-assignment LICM throughout the source tree. When a
+loop is hoisted, keep its remaining body intact; further passes can reanalyze it. -/
+def hoistStmt (prog : Program) : (s : Stmt) → Embed s.cfg prog.cfg → Stmt
+  | .basic bs, _ => .basic bs
+  | .andThen a b, e =>
+      .andThen (hoistStmt prog a ((Embed.sequenceLeft a.cfg b.cfg).trans e))
+        (hoistStmt prog b ((Embed.sequenceRight a.cfg b.cfg).trans e))
+  | .ifElse cond a b, e =>
+      .ifElse cond (hoistStmt prog a ((Embed.overlayLeft a.cfg b.cfg).trans e))
+        (hoistStmt prog b ((Embed.overlayRight a.cfg b.cfg).trans e))
+  | .whileLoop cond body, e =>
+      match hoistLoop prog cond body e with
+      | some moved => moved
+      | none => .whileLoop cond (hoistStmt prog body ((Embed.loop body.cfg).trans e))
+
+/-- Run reaching definitions on the source program and perform guarded LICM. -/
+def hoistProgram (prog : Program) : Stmt := hoistStmt prog prog.rootStmt prog.rootEmbed
+
 /-- A human-readable report of the loop-invariant assignments. -/
 def output (prog : Program) : String :=
   match licmCandidates prog with