agda-spa/lean/Spa/Transformation/Licm.lean

import Spa.Analysis.Reaching
import Spa.Language.Equivalence

/-!
# Loop-invariant code motion

This wires the **reaching-definitions** analysis (`Spa/Analysis/Reaching.lean`)
to the AST to find assignments inside a `while` loop whose right-hand side
depends only on definitions made outside the loop. `licmCandidates` reports
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
`Embed.trans` per descent). That embedding is what supplies program states:

1. at an assignment, its CFG state is `Embed.singletonIndex` — the subtree's CFG
   is a `singleton`, so its sole node is the state, and `nodes_eq` proves it
   holds that very statement;
2. read the reaching definitions at the assignment's *entry* (`joinForKey s
   result` — the join over predecessors, i.e. before the assignment runs);
3. union the definition sets of the RHS variables;
4. check no definition site lies in the loop body's CFG range. Every embedding is
   a constant index shift, so the body occupies the interval
   `[off, off + size)` (`GGraph.Embed.mem_range_iff`) and the test is two
   comparisons.

If every reaching definition of every RHS variable lies outside the loop, the
assignment is reported as loop-invariant. Hoisting additionally requires the
assignment to lead the body, its destination to be absent from the guard, and
no reassignment of that destination in the remaining body. The hoist is guarded
by the original condition, preserving zero-iteration behavior.

`LicmTransformation.hoistProgram_eval` in `Spa/Transformation/Licm/Correctness.lean`
proves preservation of terminating executions and observable final bindings.

Transitive invariance and motion of non-leading assignments are not implemented.
-/

namespace Spa

namespace LicmTransformation

open Forward GGraph

/-- The CFG footprint of an enclosing loop: its entry node (for reporting) and
the index interval its body occupies. -/
structure Enclosing (prog : Program) where
  /-- The loop's entry node, i.e. `GGraph.loopIn` embedded into the program. -/
  loopState : prog.State
  /-- Start of the body's index range. -/
  bodyOff : ℕ
  /-- Length of the body's index range. -/
  bodySize : ℕ

/-- Is this definition site inside the loop body's CFG range? -/
def Enclosing.covers {prog : Program} (l : Enclosing prog) (d : prog.State) : Bool :=
  decide (l.bodyOff ≤ d.val ∧ d.val < l.bodyOff + l.bodySize)

/-- An assignment found inside a loop, paired with the data needed to test its
invariance against that (immediately enclosing) loop. -/
structure Candidate (prog : Program) where
  /-- The enclosing loop. -/
  encl : Enclosing prog
  /-- The assignment's CFG state. -/
  assignState : prog.State
  /-- The variables read by the assignment's RHS. -/
  rhsVars : List String

/-- Collect every assignment together with its *immediately enclosing* loop.
`enc` is `none` outside any loop, in which case assignments are skipped — only
in-loop assignments are candidates. -/
def collectCandidates (prog : Program) (enc : Option (Enclosing prog)) :
    (s : Stmt) → Embed s.cfg prog.cfg → List (Candidate prog)
  | .basic bs, e =>
      match bs, enc with
      | .assign _ ex, some l =>
          [{ encl := l, assignState := e.singletonIndex,
             rhsVars := ex.vars.sort (· ≤ ·) }]
      | _, _ => []
  | .andThen s₁ s₂, e =>
      collectCandidates prog enc s₁ ((Embed.sequenceLeft s₁.cfg s₂.cfg).trans e) ++
      collectCandidates prog enc s₂ ((Embed.sequenceRight s₁.cfg s₂.cfg).trans e)
  | .ifElse _ s₁ s₂, e =>
      collectCandidates prog enc s₁ ((Embed.overlayLeft s₁.cfg s₂.cfg).trans e) ++
      collectCandidates prog enc s₂ ((Embed.overlayRight s₁.cfg s₂.cfg).trans e)
  | .whileLoop _ body, e =>
      let be := (Embed.loop body.cfg).trans e
      collectCandidates prog
        (some { loopState := e.f body.cfg.loopIn, bodyOff := be.off,
                bodySize := body.cfg.size }) body be

/-- Read the definition set assigned to variable `k`, or `⊥` if absent. -/
def lookupDef (prog : Program) (vs : VariableValues (DefSet prog) prog)
    (k : String) : DefSet prog :=
  if h : FiniteMap.MemKey k vs then (FiniteMap.locate h).1 else ⊥

/-- Is the candidate assignment loop-invariant: do all reaching definitions of
its RHS variables lie outside the loop body? -/
def isInvariant (prog : Program) (c : Candidate prog) : Bool :=
  let entry := joinForKey c.assignState (result (DefSet prog) prog)
  let combined : DefSet prog :=
    c.rhsVars.foldl (fun acc k => acc ⊔ lookupDef prog entry k) ⊥
  -- `Finset.toList` is noncomputable; the decidable bounded-∀ folds over the
  -- underlying multiset and keeps `lake exe` working.
  decide (∀ d ∈ combined, c.encl.covers d = false)

/-- The loop-invariant assignments of `prog`, as `(loop, assignment)` state pairs. -/
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
  | [] => "no loop-invariant assignments found"
  | cands =>
      "loop-invariant assignments (loop ↦ assignment):\n" ++
        String.intercalate "\n"
          (cands.map (fun p => s!"  loop #{p.1.val}: assignment #{p.2.val}"))

end LicmTransformation

end Spa