import Spa.Analysis.Reaching

/-!
# Finding loop-invariant assignments (LICM groundwork)

This wires the **reaching-definitions** analysis (`Spa/Analysis/Reaching.lean`)
to the AST to *find* — not yet move — assignments inside a `while` loop whose
right-hand side depends only on definitions made *outside* the loop. These are
the candidates a later LICM pass could hoist.

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. This is the first-order check ("all
reaching definitions outside the loop"); transitive/iterated invariance and the
actual hoisting are out of scope here.
-/

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)

/-- 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