Update LICM/Reaching to node use NodeId

lean/Spa/Analysis/Reaching.lean

@@ -1,6 +1,5 @@
 import Spa.Analysis.Forward
 import Spa.Lattice.Finset
-import Spa.Language.Tagged.Graphs
 import Spa.Showable
 
 namespace Spa
@@ -13,20 +12,18 @@ instance {n : ℕ} : Showable (Finset (Fin n)) :=
         (fun i rest => if i ∈ s then show' i ++ ", " ++ rest else rest) ""
         ++ "}"⟩
 
-abbrev DefSet (prog : Program) : Type := Finset prog.NodeId
+abbrev DefSet (prog : Program) : Type := Finset prog.State
 
 namespace ReachingAnalysis
 
 variable (prog : Program)
 
-def genSet (s : prog.State) : DefSet prog := (prog.nodeIdOf s).elim {} (fun x => {x})
-
 def eval (s : prog.State) (vs : VariableValues (DefSet prog) prog) : VariableValues (DefSet prog) prog :=
   match prog.code s with
   | none => vs
   | some bs =>
     match bs with
-    | .assign k _ => FiniteMap.generalizedUpdate id (fun _ _ => genSet prog s) [k] vs
+    | .assign k _ => FiniteMap.generalizedUpdate id (fun _ _ => {s}) [k] vs
     | .noop => vs
 
 lemma eval_mono (s : prog.State) :
@@ -50,12 +47,11 @@ def output : String :=
 abbrev Run (prog : Program) : Type := List (prog.State × BasicStmt)
 
 @[aesop unsafe cases]
-inductive LastAssign (prog : Program) (x : String) : Run prog → prog.NodeId → Prop
+inductive LastAssign (prog : Program) (x : String) : Run prog → prog.State → Prop
-  | here (s : prog.State) (e : Expr) (hc : prog.code s = some (.assign x e))
+  | here (s : prog.State) (e : Expr) (rest : Run prog) :
-      (rest : Run prog) :
+      LastAssign prog x ((s, .assign x e) :: rest) s
-      LastAssign prog x ((s, .assign x e) :: rest) (prog.nodeIdOfNonempty s hc)
   | there (s : prog.State) (bs : BasicStmt) (hc : prog.code s = some bs)
-      (rest : Run prog) {n : prog.NodeId} :
+      (rest : Run prog) {n : prog.State} :
       (∀ e, bs ≠ .assign x e) → LastAssign prog x rest n →
       LastAssign prog x ((s, bs) :: rest) n
 
@@ -73,7 +69,7 @@ instance stateInterp : StateInterpretation (DefSet prog) prog where
   Post := @runOfTrace prog
 
   interp vs run := ∀ (x : String) (assigners : DefSet prog), (x, assigners) ∈ vs →
-    ∀ (n : prog.NodeId), LastAssign prog x run n → n ∈ assigners
+    ∀ (n : prog.State), LastAssign prog x run n → n ∈ assigners
   interp_sup := by
     intro vs₁ vs₂ run h x assigners hmem n hla
     obtain ⟨a₁, a₂, rfl, h₁, h₂⟩ := FiniteMap.mem_sup hmem
@@ -105,8 +101,7 @@ private lemma valid_step (s : prog.State) {ρ₁ ρ₂ : Env}
     by_cases hx : k = x
     · subst hx
       have hd := FiniteMap.generalizedUpdate_mem_eq (List.mem_singleton.mpr rfl) hmem
-      rcases hla
+      rcases hla <;> simp [hd] <;> aesop
-        <;> simp [Program.nodeIdOfNonempty, hd, genSet, Option.get] <;> aesop
     · have hmem' := FiniteMap.generalizedUpdate_not_mem_backward
         (fun hc => hx (List.mem_singleton.mp hc)) hmem
       aesop

lean/Spa/Transformation/Licm.lean

@@ -1,23 +1,27 @@
 import Spa.Analysis.Reaching
-import Spa.Language.Tagged.Graphs
 
 /-!
 # Finding loop-invariant assignments (LICM groundwork)
 
 This wires the **reaching-definitions** analysis (`Spa/Analysis/Reaching.lean`)
-to the **tagged AST** to *find* — not yet move — assignments inside a `while`
+to the AST to *find* — not yet move — assignments inside a `while` loop whose
-loop whose right-hand side depends only on definitions made *outside* the loop.
+right-hand side depends only on definitions made *outside* the loop. These are
-These are the candidates a later LICM pass could hoist.
+the candidates a later LICM pass could hoist.
 
-The pipeline, for each assignment immediately enclosed by a loop:
+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. locate its CFG state via the tagged-graph bridge (`Program.stateOfNodeId`);
+1. at an assignment, its CFG state is `Embed.singletonIndex` — the subtree's CFG
-2. read the reaching definitions at the assignment's *entry*
+   is a `singleton`, so its sole node is the state, and `nodes_eq` proves it
-   (`joinForKey s result` — the join over predecessors, i.e. before the
+   holds that very statement;
-   assignment itself runs);
+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. map each definition site back to its `RawId` (`Program.nodeIdOf`) and check
+4. check no definition site lies in the loop body's CFG range. Every embedding is
-   it is **not** inside the loop body (structural `subtreeIds` membership).
+   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
@@ -29,64 +33,74 @@ namespace Spa
 
 namespace LicmTransformation
 
-open Forward
+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 `whileLoop`'s tag (for reporting). -/
+  /-- The enclosing loop. -/
-  loopId : prog.NodeId
+  encl : Enclosing prog
-  /-- Every node id inside the loop body (the "is-child-of-loop" set). -/
+  /-- The assignment's CFG state. -/
-  bodyIds : List prog.NodeId
+  assignState : prog.State
-  /-- The assignment `BasicStmt`'s tag — what labels its CFG node. -/
-  assignId : prog.NodeId
   /-- The variables read by the assignment's RHS. -/
   rhsVars : List String
 
 /-- Collect every assignment together with its *immediately enclosing* loop.
-`enclosing` carries the current loop's tag and body id-set, or `none` outside any
+`enc` is `none` outside any loop, in which case assignments are skipped — only
-loop (in which case assignments are skipped — only in-loop assignments are
+in-loop assignments are candidates. -/
-candidates). -/
+def collectCandidates (prog : Program) (enc : Option (Enclosing prog)) :
-def collectCandidates (prog : Program) (enc : Option (prog.NodeId × List prog.NodeId)) :
+    (s : Stmt) → Embed s.cfg prog.cfg → List (Candidate prog)
-    Stmt.Tagged prog.NodeId → List (Candidate prog)
+  | .basic bs, e =>
-  | .basic _ bs =>
       match bs, enc with
-      | .assign t _ e, some (loopId, bodyIds) =>
+      | .assign _ ex, some l =>
-          [{ loopId := loopId, bodyIds := bodyIds, assignId := t,
+          [{ encl := l, assignState := e.singletonIndex,
-             rhsVars := e.erase.vars.sort (· ≤ ·) }]
+             rhsVars := ex.vars.sort (· ≤ ·) }]
       | _, _ => []
-  | .andThen _ a b => collectCandidates prog enc a ++ collectCandidates prog enc b
+  | .andThen s₁ s₂, e =>
-  | .ifElse _ _ a b => collectCandidates prog enc a ++ collectCandidates prog enc b
+      collectCandidates prog enc s₁ ((Embed.sequenceLeft s₁.cfg s₂.cfg).trans e) ++
-  | .whileLoop loopT _ body =>
+      collectCandidates prog enc s₂ ((Embed.sequenceRight s₁.cfg s₂.cfg).trans e)
-      collectCandidates prog (some (loopT, body.subtreeIds)) body
+  | .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 ⊥
 
-/-- The AST node ids marked as definition sites in a `DefSet`. With the
-`Finset`-of-AST-ids lattice these are just the elements of the set. -/
-def defSites (prog : Program) (d : DefSet prog) : List prog.NodeId :=
-  (List.finRange prog.size).filter (fun i => decide (i ∈ d))
-
 /-- Is the candidate assignment loop-invariant: do all reaching definitions of
-its RHS variables lie outside the loop body? Reaching sets are now keyed by AST
+its RHS variables lie outside the loop body? -/
-node id, so we compare against the loop-body ids directly (embedding the raw
-body ids into `p.NodeId`). -/
 def isInvariant (prog : Program) (c : Candidate prog) : Bool :=
-  match prog.stateOfNodeId c.assignId with
+  let entry := joinForKey c.assignState (result (DefSet prog) prog)
-  | none => false
-  | some s =>
-      let entry := joinForKey s (result (DefSet prog) prog)
   let combined : DefSet prog :=
     c.rhsVars.foldl (fun acc k => acc ⊔ lookupDef prog entry k) ⊥
-      (defSites prog combined).all (fun nid => ! decide (nid ∈ c.bodyIds))
+  -- `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 `(loopId, assignId)` pairs. -/
+/-- The loop-invariant assignments of `prog`, as `(loop, assignment)` state pairs. -/
-def licmCandidates (prog : Program) : List (prog.NodeId × prog.NodeId) :=
+def licmCandidates (prog : Program) : List (prog.State × prog.State) :=
-  (collectCandidates prog none prog.taggedFin).filterMap (fun c =>
+  (collectCandidates prog none prog.rootStmt prog.rootEmbed).filterMap (fun c =>
-    if isInvariant prog c then some (c.loopId, c.assignId) else none)
+    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 :=