Add function back in to Embedding

This is a special case of an embedding, but it has the nice
property for checking inclusion.

lean/Spa.lean

@@ -19,10 +19,5 @@ import Spa.Showable
 import Spa.Analysis.Utils
 import Spa.Analysis.Sign
 import Spa.Analysis.Constant
-import Spa.Language.Tagged.Id
-import Spa.Language.Tagged.Derive
-import Spa.Language.Tagged.Basic
-import Spa.Language.Tagged.Properties
-import Spa.Language.Tagged.Graphs
 import Spa.Analysis.Reaching
 import Spa.Transformation.Licm

lean/Spa/Analysis/Constant.lean

@@ -137,12 +137,11 @@ theorem analyze_correct {ρ : Env} (hrun : EvalStmt [] prog.rootStmt ρ) :
     ⟦ variablesAt prog.finalState (result ConstLattice prog) ⟧ ρ :=
   Forward.analyze_correct ConstLattice prog hrun
 
-theorem analyze_correct_at {ρf : Env} (hrun : EvalStmt [] prog.rootStmt ρf)
-    {s : prog.State} {ρin ρout : Env}
-    (hr : Reaches (prog.trace hrun) s ρin ρout) :
+theorem analyze_correct_at {s : prog.State} {ρin ρout : Env}
+    (hr : Reaches s ρin ρout) :
     ⟦ joinForKey s (result ConstLattice prog) ⟧ ρin
       ∧ ⟦ variablesAt s (result ConstLattice prog) ⟧ ρout :=
-  Forward.analyze_correct_at ConstLattice prog hrun hr
+  Forward.analyze_correct_at ConstLattice prog hr
 
 end ConstAnalysis
 

lean/Spa/Analysis/Forward.lean

@@ -90,75 +90,54 @@ lemma stepTrace {s₁ s₂ : prog.State} {ρ₁ ρ₂ : Env}
   rw [variablesAt_joinAll]
   exact hjoin
 
-/-- Soundness at *every* visited node: if the analysis result over-approximates the
-    incoming environment at the start of the trace, then at each node reached along the
-    way it over-approximates both the environment entering that node (via `joinForKey`)
-    and the environment leaving it (via `variablesAt`). The intermediate `variablesAt`
-    evidence used to be computed and discarded inside `walkTrace`; here it is returned. -/
-lemma walkTrace_reaches {s₁ s₂ s₃: prog.State} {ρ₁ ρ₂ ρ₃: Env}
-    {s : prog.State} {ρin ρout : Env}
-    {tr : Trace prog.cfg s₂ s₃ ρ₂ ρ₃}
-    (hr : Reaches tr s ρin ρout)
-    (trₗ : Traceₗ prog.cfg s₁ s₂ ρ₁ ρ₂)
-    (hjoin : ⟦ joinForKey s₂ (result L prog) ⟧ (S.Pre trₗ)) :
-    ⟦ joinForKey s (result L prog) ⟧ (S.Pre (trₗ ++ hr.pre))
-      ∧ ⟦ variablesAt s (result L prog) ⟧ (S.Post (trₗ ++ hr.post)) := by
-  induction hr with
-  | single_here hnode =>
-    simp [Reaches.pre, Reaches.post]
-    refine ⟨?_, ?_⟩ <;> try simpa [HAppend.hAppend]
-    exact stepTrace trₗ hjoin hnode
-  | edge_here hnode hedge rest =>
-    simp [Reaches.pre, Reaches.post]
-    refine ⟨?_, ?_⟩ <;> try simpa [HAppend.hAppend]
-    exact stepTrace trₗ hjoin hnode
-  | edge_there hnode hedge rest hr' ih =>
-    have hstep := stepTrace trₗ hjoin hnode
-    have hmem := FiniteMap.mem_valuesAt prog.states_nodup
-      (prog.mem_incoming_of_edge hedge) (variablesAt_mem _ (result L prog))
-    simpa [Reaches.pre, Reaches.post, HAppend.hAppend] using
-      ih ((trₗ ++ hnode).addEdge hedge)
-        (interp_foldr (S.post_pre (trₗ ++ hnode) hedge hstep) hmem)
+/-- Soundness propagates along an execution prefix: if the analysis is sound at
+    `s₂` for the run so far (`trₗ`), then it is sound wherever the further prefix
+    `mid` ends up. -/
+lemma walkPrefix : ∀ {s₂ s : prog.State} {ρ₂ ρin : Env}
+    (mid : Traceₗ prog.cfg s₂ s ρ₂ ρin) {s₁ : prog.State} {ρ₁ : Env}
+    (trₗ : Traceₗ prog.cfg s₁ s₂ ρ₁ ρ₂),
+    ⟦ joinForKey s₂ (result L prog) ⟧ (S.Pre trₗ) →
+    ⟦ joinForKey s (result L prog) ⟧ (S.Pre (trₗ ++ mid)) := by
+  intro s₂ s ρ₂ ρin mid
+  induction mid with
+  | nil => intro s₁ ρ₁ trₗ hjoin; simpa [HAppend.hAppend, Traceₗ.append] using hjoin
+  | cons hnode hedge rest ih =>
+      intro s₁ ρ₁ trₗ hjoin
+      have hstep := stepTrace trₗ hjoin hnode
+      have hmem := FiniteMap.mem_valuesAt prog.states_nodup
+        (prog.mem_incoming_of_edge hedge) (variablesAt_mem _ (result L prog))
+      simpa [HAppend.hAppend, Traceₗ.append] using
+        ih ((trₗ ++ hnode).addEdge hedge)
+          (interp_foldr (S.post_pre (trₗ ++ hnode) hedge hstep) hmem)
 
 omit [DecidableEq L] in
 /-- The final node of a trace is always reached, with the environment/state the trace
-    ends in. Used to recover the final-state soundness theorem from `walkTrace_reaches`. -/
-def reaches_final {s₁ s₂ : prog.State} {ρ₁ ρ₂ : Env}
-    (tr : Trace prog.cfg s₁ s₂ ρ₁ ρ₂) :
-    Σ ρin, Reaches tr s₂ ρin ρ₂ :=
-  match tr with
-  | .single hnode => ⟨_, .single_here hnode⟩
-  | .edge hnode hedge rest =>
-    let ⟨ρin, r'⟩ := reaches_final rest; ⟨ρin, .edge_there hnode hedge _ r'⟩
+    ends in. Used to recover the final-state soundness theorem from `walkPrefix`. -/
+def reaches_final {s : prog.State} {ρ : Env}
+    (tr : Trace prog.cfg prog.initialState s [] ρ) : Σ ρin, Reaches s ρin ρ :=
+  ⟨_, ⟨tr.split.2.1, tr.split.2.2⟩⟩
 
 omit [DecidableEq L] in
-/-- Reaching the final node covers the whole trace. -/
-@[simp] lemma reaches_final_post {s₁ s₂ : prog.State} {ρ₁ ρ₂ : Env}
-    (tr : Trace prog.cfg s₁ s₂ ρ₁ ρ₂) :
-    (reaches_final tr).2.post = tr := by
-  induction tr with
-  | single hnode => rfl
-  | edge hnode hedge rest ih => simp [reaches_final, Reaches.post, ih]
+@[simp] lemma reaches_final_post {s : prog.State} {ρ : Env}
+    (tr : Trace prog.cfg prog.initialState s [] ρ) :
+    (reaches_final tr).2.post = tr := Trace.split_append tr
 
 variable (L prog) in
-/-- Soundness at every program point reached during execution: for any node `s` visited
-    by the run `hrun` (witnessed by `hr`), the analysis result over-approximates both the
-    environment entering `s` and the one leaving it. The final-state theorem
-    `analyze_correct_state` is the special case where `s` is `prog.finalState`. -/
-theorem analyze_correct_at {ρf : Env} (hrun : EvalStmt [] prog.rootStmt ρf)
-    {s : prog.State} {ρin ρout : Env}
-    (hr : Reaches (prog.trace hrun) s ρin ρout) :
+/-- Soundness at every program point an execution actually visits: the analysis
+    over-approximates both the environment entering that point and the one leaving
+    it. -/
+theorem analyze_correct_at {s : prog.State} {ρin ρout : Env} (hr : Reaches s ρin ρout) :
     ⟦ joinForKey s (result L prog) ⟧ (S.Pre hr.pre)
-      ∧ ⟦ variablesAt s (result L prog) ⟧ (S.Post hr.post) := by
-  refine walkTrace_reaches hr (Traceₗ.single _ _ []) ?_
-  rw [joinForKey_initialState]
-  exact ValidStateEvaluator.botV_init
+      ∧ ⟦ variablesAt s (result L prog) ⟧ (S.Post hr.post) :=
+  have hpre := walkPrefix hr.pre Traceₗ.nil
+    (by rw [joinForKey_initialState]; exact ValidStateEvaluator.botV_init)
+  ⟨hpre, stepTrace hr.pre hpre hr.step⟩
 
 variable (L prog) in
 theorem analyze_correct'
     {ρ : Env} (hrun : EvalStmt [] prog.rootStmt ρ) :
     ⟦ variablesAt prog.finalState (result L prog) ⟧ (S.Post (prog.trace hrun)) := by
-  have h := (analyze_correct_at L prog hrun (reaches_final (prog.trace hrun)).2).2
+  have h := (analyze_correct_at L prog (reaches_final (prog.trace hrun)).2).2
   rwa [reaches_final_post] at h
 
 end

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
-  | here (s : prog.State) (e : Expr) (hc : prog.code s = some (.assign x e))
-      (rest : Run prog) :
-      LastAssign prog x ((s, .assign x e) :: rest) (prog.nodeIdOfNonempty s hc)
+inductive LastAssign (prog : Program) (x : String) : Run prog → prog.State → Prop
+  | here (s : prog.State) (e : Expr) (rest : Run prog) :
+      LastAssign prog x ((s, .assign x e) :: rest) s
   | 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
-        <;> simp [Program.nodeIdOfNonempty, hd, genSet, Option.get] <;> aesop
+      rcases hla <;> simp [hd] <;> aesop
     · have hmem' := FiniteMap.generalizedUpdate_not_mem_backward
         (fun hc => hx (List.mem_singleton.mp hc)) hmem
       aesop
@@ -123,12 +118,11 @@ theorem analyze_correct {ρ : Env} (hrun : EvalStmt [] prog.rootStmt ρ) :
       (runOfTrace prog (prog.trace hrun)) :=
   Forward.analyze_correct' (DefSet prog) prog hrun
 
-theorem analyze_correct_at {ρf : Env} (hrun : EvalStmt [] prog.rootStmt ρf)
-    {s : prog.State} {ρin ρout : Env}
-    (hr : Reaches (prog.trace hrun) s ρin ρout) :
+theorem analyze_correct_at {s : prog.State} {ρin ρout : Env}
+    (hr : Reaches s ρin ρout) :
     ⟦ joinForKey s (result (DefSet prog) prog) ⟧ (runOfTraceₗ prog hr.pre)
       ∧ ⟦ variablesAt s (result (DefSet prog) prog) ⟧ (runOfTrace prog hr.post) :=
-  Forward.analyze_correct_at (DefSet prog) prog hrun hr
+  Forward.analyze_correct_at (DefSet prog) prog hr
 
 end ReachingAnalysis
 

lean/Spa/Analysis/Sign.lean

@@ -111,13 +111,16 @@ namespace SignAnalysis
 
 variable (prog : Program)
 
+/-- The sign of an integer literal. -/
+def signOf (z : ℤ) : SignLattice :=
+  if z = 0 then .mk .zero else if 0 < z then .mk .plus else .mk .minus
+
 def eval : Expr → VariableValues SignLattice prog → SignLattice
   | .add e₁ e₂, vs => plus (eval e₁ vs) (eval e₂ vs)
   | .sub e₁ e₂, vs => minus (eval e₁ vs) (eval e₂ vs)
   | .var k, vs =>
     if h : FiniteMap.MemKey k vs then (FiniteMap.locate h).1 else .top
-  | .num 0, _ => .mk .zero
-  | .num (_ + 1), _ => .mk .plus
+  | .num z, _ => signOf z
 
 lemma eval_mono (e : Expr) : Monotone (eval prog e) := by
   induction e with
@@ -139,7 +142,7 @@ lemma eval_mono (e : Expr) : Monotone (eval prog e) := by
           dif_neg (fun hm => hk (FiniteMap.MemKey_iff.mp hm))]
   | num n =>
     intro vs₁ vs₂ _
-    cases n <;> exact le_refl _
+    exact le_refl _
 
 instance exprEvaluator : ExprEvaluator SignLattice prog :=
   ⟨eval prog, eval_mono prog⟩
@@ -159,6 +162,20 @@ private lemma int_neg_iff (z : ℤ) : (∃ n : ℕ, z = -((n : ℤ) + 1)) ↔ z
   · rintro ⟨n, rfl⟩; omega
   · intro h; exact ⟨(-z - 1).toNat, by omega⟩
 
+/-- `signOf` really does describe the literal it was computed from. -/
+lemma interp_signOf (z : ℤ) : ⟦signOf z⟧ (Value.int z) := by
+  unfold signOf
+  split
+  · case isTrue h => subst h; rfl
+  · rename_i hne
+    split
+    · case isTrue hpos =>
+        simp only [signInterpretation, interpSign, Value.int.injEq, int_pos_iff]
+        exact hpos
+    · case isFalse hnpos =>
+        simp only [signInterpretation, interpSign, Value.int.injEq, int_neg_iff]
+        omega
+
 lemma plus_valid {g₁ g₂ : SignLattice} {z₁ z₂ : ℤ}
     (h₁ : ⟦g₁⟧ (.int z₁)) (h₂ : ⟦g₂⟧ (.int z₂)) :
     ⟦plus g₁ g₂⟧ (.int (z₁ + z₂)) := by
@@ -184,9 +201,7 @@ instance eval_valid : ValidExprEvaluator SignLattice prog := by
   | num n =>
     intro _
     show ⟦eval prog (.num n) vs⟧ (.int n)
-    cases n with
-    | zero => rfl
-    | succ n' => exact ⟨n', congrArg Value.int (by norm_cast)⟩
+    exact interp_signOf n
   | var x v hxv =>
     intro hvs
     show ⟦eval prog (.var x) vs⟧ v
@@ -213,12 +228,11 @@ theorem analyze_correct {ρ : Env} (hrun : EvalStmt [] prog.rootStmt ρ) :
     ⟦ variablesAt prog.finalState (result SignLattice prog) ⟧ ρ :=
   Forward.analyze_correct SignLattice prog hrun
 
-theorem analyze_correct_at {ρf : Env} (hrun : EvalStmt [] prog.rootStmt ρf)
-    {s : prog.State} {ρin ρout : Env}
-    (hr : Reaches (prog.trace hrun) s ρin ρout) :
+theorem analyze_correct_at {s : prog.State} {ρin ρout : Env}
+    (hr : Reaches s ρin ρout) :
     ⟦ joinForKey s (result SignLattice prog) ⟧ ρin
       ∧ ⟦ variablesAt s (result SignLattice prog) ⟧ ρout :=
-  Forward.analyze_correct_at SignLattice prog hrun hr
+  Forward.analyze_correct_at SignLattice prog hr
 
 end SignAnalysis
 

lean/Spa/Language/Base.lean

@@ -5,9 +5,13 @@ import Mathlib.Data.Finset.Basic
 # Base Language
 
 This file defines the core object language for the program analysis and
-transformation. It's a very basic imperative language. The `Spa/Language/Tagged/Basic.lean`
-file provides an auto-derived version of the `Expr`, `BasicStmt`, and `Stmt` data
-types with unique IDs per condtructor, enabling in-AST pointers.
+transformation. It's a very basic imperative language.
+
+Program points are identified by their node in the control flow graph rather than
+by an identifier stored in the AST: a recursion over a `Stmt` threads a
+`Spa.GGraph.Embed` of the subtree's CFG into the whole program's (starting from
+`Spa.Program.rootEmbed`), which yields the CFG index of each basic statement
+along with a proof that the node carries it.
 
 -/
 
@@ -18,7 +22,7 @@ inductive Expr where
   | add (e₁ e₂ : Expr)
   | sub (e₁ e₂ : Expr)
   | var (x : String)
-  | num (n : ℕ)
+  | num (z : ℤ)
   deriving DecidableEq
 
 /-- A statement that cannot alter control flow (and thus, can be part of a basic block).

lean/Spa/Language/Graphs.lean

@@ -215,62 +215,103 @@ lemma wrap_outputs (g : GGraph (Option β)) :
 
 /-! ### Embeddings
 
-Each composition operator includes its operands into the result via an index
-translation that preserves node payloads and edges. `Embed` captures exactly
-those two facts, so anything defined from `nodes` and `edges` (traces, node
-labels, …) can be transported along an embedding once, instead of once per
-operator.
+To be able to reason compositionally about traces through the graphs,
+we need to be able to reason about how a trace within a sub-graph maps
+to the full graph. Fortunately, graphs are built using composition operators,
+and these composition operators always include their arguments as embedded
+subgraphs in the full result. Moreover, each embedding "just" offsets the
+existing node IDs by a given amount.
 
-`Embed` is deliberately a structure rather than a class: for `g ⤳ g`, both the
-left and the right inclusion inhabit the same type `Embed g (g ⤳ g)`, so
-instance resolution could silently pick the wrong copy. Embeddings into a
-composed graph are non-canonical by design; a named witness says which
-inclusion is meant. -/
+This section formalizes this fact by providing an `Embed` type that
+represents an offset-based embedding, and showing that such an embedding
+exists for all arguments given to graph composition operators. Furthermore,
+because of the offset-based embedding, we can determine whether a node
+came from a particular subgraph simply by examining its offset and sub-graph
+size. This is captured by `Embed.mem_range_iff`. -/
 
-/-- An embedding of graph `g` into graph `h`: an index translation that
-    preserves node payloads and edges. -/
+/-- A special-case embedding of `g` into `h` in which all edges and nodes
+    of `g` are present in `h` at a given offset `off`. -/
 structure Embed (g h : GGraph α) where
   f : g.Index → h.Index
+  off : ℕ
+  f_val : ∀ i, (f i).val = off + i.val
   nodes_eq : ∀ i, h.nodes (f i) = g.nodes i
   edges_mem : ∀ {e : g.Edge}, e ∈ g.edges → (f e.1, f e.2) ∈ h.edges
 
-/-- Embeddings compose. -/
+lemma Embed.f_inj {g h : GGraph α} (e : Embed g h) : Function.Injective e.f := by
+  intro i j hij
+  have := congrArg Fin.val hij
+  rw [e.f_val, e.f_val] at this
+  exact Fin.ext (by omega)
+
+/-- An embedding's range is the interval `[off, off + g.size)`. -/
+lemma Embed.mem_range_iff {g h : GGraph α} (e : Embed g h) (j : h.Index) :
+    (∃ i, e.f i = j) ↔ e.off ≤ j.val ∧ j.val < e.off + g.size := by
+  constructor
+  · rintro ⟨i, rfl⟩; have := i.isLt; rw [e.f_val]; omega
+  · rintro ⟨hlo, hhi⟩
+    refine ⟨⟨j.val - e.off, by omega⟩, Fin.ext ?_⟩
+    rw [e.f_val]
+    show e.off + (j.val - e.off) = j.val
+    omega
+
+/-- Build an embedding from an index map that is pointwise the shift.  The five
+    inclusions below are naturally written with `Fin.castAdd`/`Fin.natAdd` — the form
+    the `Fin.append` lemmas are stated in — so this lets them keep those proofs
+    verbatim.  The trailing argument is boilerplate at every call site and defaults
+    to discharging itself. -/
+private def Embed.ofIndexMap {g h : GGraph α} (off : ℕ) (k : g.Index → h.Index)
+    (hn : ∀ i, h.nodes (k i) = g.nodes i)
+    (hem : ∀ {e : g.Edge}, e ∈ g.edges → (k e.1, k e.2) ∈ h.edges)
+    (hk : ∀ i, (k i).val = off + i.val := by intro i; simp) :
+    Embed g h where
+  f := k
+  off := off
+  f_val := hk
+  nodes_eq := hn
+  edges_mem := hem
+
+/-- Embeddings compose (offsets add). -/
 def Embed.trans {g₁ g₂ g₃ : GGraph α} (e₁ : Embed g₁ g₂) (e₂ : Embed g₂ g₃) :
-    Embed g₁ g₃ where
-  f := e₂.f ∘ e₁.f
-  nodes_eq i := (e₂.nodes_eq (e₁.f i)).trans (e₁.nodes_eq i)
-  edges_mem he := e₂.edges_mem (e₁.edges_mem he)
+    Embed g₁ g₃ :=
+  ofIndexMap (e₂.off + e₁.off) (fun i => e₂.f (e₁.f i))
+    (fun i => (e₂.nodes_eq (e₁.f i)).trans (e₁.nodes_eq i))
+    (fun he => e₂.edges_mem (e₁.edges_mem he))
+    (hk := fun i => by rw [e₂.f_val, e₁.f_val]; omega)
 
 /-- The left operand's inclusion into a sequenced graph. -/
-def Embed.sequenceLeft (g₁ g₂ : GGraph α) : Embed g₁ (g₁ ⤳ g₂) where
-  f i := i.castAdd g₂.size
-  nodes_eq i := Fin.append_left g₁.nodes g₂.nodes i
-  edges_mem he := List.mem_append_left _ (List.mem_append_left _ (List.mem_map_of_mem _ he))
+def Embed.sequenceLeft (g₁ g₂ : GGraph α) : Embed g₁ (g₁ ⤳ g₂) :=
+  ofIndexMap 0 (fun i => i.castAdd g₂.size) (Fin.append_left g₁.nodes g₂.nodes)
+    (fun he => List.mem_append_left _ (List.mem_append_left _ (List.mem_map_of_mem _ he)))
 
 /-- The right operand's inclusion into a sequenced graph. -/
-def Embed.sequenceRight (g₁ g₂ : GGraph α) : Embed g₂ (g₁ ⤳ g₂) where
-  f i := i.natAdd g₁.size
-  nodes_eq i := Fin.append_right g₁.nodes g₂.nodes i
-  edges_mem he := List.mem_append_left _ (List.mem_append_right _ (List.mem_map_of_mem _ he))
+def Embed.sequenceRight (g₁ g₂ : GGraph α) : Embed g₂ (g₁ ⤳ g₂) :=
+  ofIndexMap g₁.size (fun i => i.natAdd g₁.size) (Fin.append_right g₁.nodes g₂.nodes)
+    (fun he => List.mem_append_left _ (List.mem_append_right _ (List.mem_map_of_mem _ he)))
 
 /-- The left operand's inclusion into an overlaid graph. -/
-def Embed.overlayLeft (g₁ g₂ : GGraph α) : Embed g₁ (g₁ ∙ g₂) where
-  f i := i.castAdd g₂.size
-  nodes_eq i := Fin.append_left g₁.nodes g₂.nodes i
-  edges_mem he := List.mem_append_left _ (List.mem_map_of_mem _ he)
+def Embed.overlayLeft (g₁ g₂ : GGraph α) : Embed g₁ (g₁ ∙ g₂) :=
+  ofIndexMap 0 (fun i => i.castAdd g₂.size) (Fin.append_left g₁.nodes g₂.nodes)
+    (fun he => List.mem_append_left _ (List.mem_map_of_mem _ he))
 
 /-- The right operand's inclusion into an overlaid graph. -/
-def Embed.overlayRight (g₁ g₂ : GGraph α) : Embed g₂ (g₁ ∙ g₂) where
-  f i := i.natAdd g₁.size
-  nodes_eq i := Fin.append_right g₁.nodes g₂.nodes i
-  edges_mem he := List.mem_append_right _ (List.mem_map_of_mem _ he)
+def Embed.overlayRight (g₁ g₂ : GGraph α) : Embed g₂ (g₁ ∙ g₂) :=
+  ofIndexMap g₁.size (fun i => i.natAdd g₁.size) (Fin.append_right g₁.nodes g₂.nodes)
+    (fun he => List.mem_append_right _ (List.mem_map_of_mem _ he))
 
 /-- The body's inclusion into a `loop` graph. -/
-def Embed.loop (g : GGraph (Option β)) : Embed g (loop g) where
-  f i := i.natAdd 2
-  nodes_eq i := Fin.append_right (fun _ : Fin 2 => none) g.nodes i
-  edges_mem he := List.mem_append_left _ (List.mem_append_left _
-    (List.mem_append_left _ (List.mem_map_of_mem _ he)))
+def Embed.loop (g : GGraph (Option β)) : Embed g (GGraph.loop g) :=
+  ofIndexMap 2 (fun i => i.natAdd 2) (Fin.append_right (fun _ : Fin 2 => none) g.nodes)
+    (fun he => List.mem_append_left _ (List.mem_append_left _
+      (List.mem_append_left _ (List.mem_map_of_mem _ he))))
+
+/-- A `singleton` subgraph has exactly one node; this is where it sits in the ambient graph. -/
+def Embed.singletonIndex {a : α} {h : GGraph α} (e : Embed (singleton a) h) : h.Index :=
+  e.f ⟨0, Nat.zero_lt_one⟩
+
+@[simp] lemma Embed.nodes_singletonIndex {a : α} {h : GGraph α}
+    (e : Embed (singleton a) h) : h.nodes e.singletonIndex = a :=
+  e.nodes_eq ⟨0, Nat.zero_lt_one⟩
 
 variable (g : GGraph α)
 

lean/Spa/Language/Program.lean

@@ -31,6 +31,11 @@ def cfg : Graph := Graph.wrap p.rootStmt.cfg
 /-- A state in the control flow `Spa.Graph` of this program. -/
 abbrev State : Type := p.cfg.Index
 
+/-- The root statement's CFG sits inside the program's CFG. -/
+def rootEmbed : GGraph.Embed p.rootStmt.cfg p.cfg :=
+  (GGraph.Embed.sequenceLeft p.rootStmt.cfg (Graph.singleton none)).trans
+    (GGraph.Embed.sequenceRight (Graph.singleton none) _)
+
 /-- Variables mentioned or defined in this program. -/
 def vars : List String := p.rootStmt.vars.sort (· ≤ ·)
 

lean/Spa/Language/Semantics.lean

@@ -33,7 +33,7 @@ inductive Env.Mem : String × Value → Env → Prop
 /-- Inference rules for evaluating an expression (`Spa.Expr`) in a given
     environment. Pretty standard big-step expression evaluation. -/
 inductive EvalExpr : Env → Expr → Value → Prop
-  | num (ρ : Env) (n : ℕ) : EvalExpr ρ (.num n) (.int n)
+  | num (ρ : Env) (z : ℤ) : EvalExpr ρ (.num z) (.int z)
   | var (ρ : Env) (x : String) (v : Value) :
       Env.Mem (x, v) ρ → EvalExpr ρ (.var x) v
   | add (ρ : Env) (e₁ e₂ : Expr) (z₁ z₂ : ℤ) :

lean/Spa/Language/Tagged/Basic.lean

@@ -1,18 +0,0 @@
-import Spa.Language.Base
-import Spa.Language.Tagged.Id
-import Spa.Language.Tagged.Derive
-
-derive_tagged Spa.Expr Spa.BasicStmt Spa.Stmt
-
-namespace Spa
-
-def tagStmt (s : Stmt) : Stmt.Tagged RawId := (s.tag 0).1
-
-def Stmt.Tagged.subtreeIds {τ : Type} (s : Stmt.Tagged τ) : List τ :=
-  s.foldTags (· :: ·) []
-
-def Stmt.Tagged.isInLoopBody {τ : Type} [DecidableEq τ]
-    (body : Stmt.Tagged τ) (id : τ) : Bool :=
-  decide (id ∈ body.subtreeIds)
-
-end Spa

lean/Spa/Language/Tagged/Derive.lean

@@ -1,509 +0,0 @@
-import Lean
-import Mathlib.Tactic.DeriveTraversable
-import Spa.Language.Base
-import Spa.Language.Tagged.Id
-
-/-!
-# The `derive_tagged` command
-
-`derive_tagged T₁ T₂ … Tₙ` takes a family of (possibly mutually recursive)
-inductive types and generates, for each `Tᵢ`:
-
-* a *tagged* mirror inductive `Tᵢ.Tagged (τ : Type)`, in which every constructor
-  carries a leading `tag : τ` field and every field whose type is a family
-  member is retyped to its `.Tagged τ` counterpart;
-* `Tᵢ.Tagged.erase : Tᵢ.Tagged τ → Tᵢ`, forgetting all tags;
-* `Tᵢ.tag : Tᵢ → ℕ → Tᵢ.Tagged RawId × ℕ`, assigning every node a unique
-  `RawId` (its postorder index) by a single unified traversal that threads a
-  counter; the whole family shares one counter, so identifiers are unique across
-  types.
-
-The generated declarations have exactly the shape of the hand-written reference;
-see `Spa/Language/Tagged/Basic.lean` (which invokes this command) and the proofs
-in `Spa/Language/Tagged/Properties.lean`.
-
-Scope: the generator handles non-indexed inductives whose constructor fields are
-either scalars or *direct* references to a family member (which covers the object
-language).  Nested occurrences such as `List Tᵢ` are not supported.
--/
-
-open Lean Elab Command Meta
-
-namespace Spa.DeriveTagged
-
-/-- One constructor field, classified as a recursive family reference or a scalar
-(whose type syntax we keep verbatim for the mirror inductive). -/
-structure FieldData where
-  isRec : Bool
-  recType : Name
-  typeStx : Term
-
-/-- A constructor: its original (full) name, short name, and fields. -/
-structure CtorData where
-  origName : Name
-  shortName : Name
-  fields : Array FieldData
-
-/-- A family member together with its constructors. -/
-structure TypeData where
-  name : Name
-  ctors : Array CtorData
-
-def taggedOf (n : Name) : Name := n ++ `Tagged
-def eraseOf (n : Name) : Name := n ++ `Tagged ++ `erase
-def rootTagOf (n : Name) : Name := n ++ `Tagged ++ `rootTag
-def tagOf (n : Name) : Name := n ++ `tag
-def foldTagsOf (n : Name) : Name := n ++ `Tagged ++ `foldTags
-def wfOf (n : Name) : Name := n ++ `Tagged ++ `WF
-def narrowOf (n : Name) : Name := n ++ `Tagged ++ `narrow
-def narrowEraseOf (n : Name) : Name := n ++ `Tagged ++ `narrow_erase
-def tagLeOf (n : Name) : Name := n ++ `tag_le
-def tagRootTagPostOf (n : Name) : Name := n ++ `tag_rootTag_post
-def tagWfOf (n : Name) : Name := n ++ `tag_wf
-
-/-- Project the `i`-th conjunct (1-based) out of `hyp`, which has type a
-right-nested `And` of `total` conjuncts, e.g. `hyp |>.2 |>.2 |>.1`. -/
-def projAnd {m : Type → Type} [Monad m] [MonadQuotation m]
-    (hyp : Term) (i total : Nat) : m Term := do
-  let mut t := hyp
-  for _ in [0:i-1] do
-    t ← `($t |>.2)
-  if i < total then
-    t ← `($t |>.1)
-  return t
-
-/-- Combine a non-empty array of propositions into a right-nested conjunction. -/
-def mkAndR {m : Type → Type} [Monad m] [MonadQuotation m]
-    (cs : Array Term) : m Term := do
-  let mut t := cs.back!
-  for c in cs.pop.reverse do
-    t ← `($c ∧ $t)
-  return t
-
-/-- For a constructor, return one entry per *recursive* field: its argument
-identifier, the family member it references, and the start-counter expression at
-which it is tagged (`n`, then `(a.tag n).2`, …) — the same threading `mkTag`
-uses. -/
-def recChildren (cd : CtorData) (argNames : Array Ident) (nStart : Term) :
-    CommandElabM (Array (Ident × Name × Term)) := do
-  let mut res : Array (Ident × Name × Term) := #[]
-  let mut cur := nStart
-  for (f, a) in cd.fields.zip argNames do
-    if f.isRec then
-      res := res.push (a, f.recType, cur)
-      cur ← `(($(mkIdent (tagOf f.recType)) $a $cur) |>.2)
-  return res
-
-/-- Inspect the family, classifying each constructor field. -/
-def gather (family : Array Name) (τ : Ident) : TermElabM (Array TypeData) := do
-  let famSet : NameSet := family.foldl (·.insert ·) {}
-  family.mapM fun tn => do
-    let iv ← getConstInfoInduct tn
-    let ctors ← iv.ctors.toArray.mapM fun cn => do
-      let cv ← getConstInfoCtor cn
-      let fields ← forallTelescopeReducing cv.type fun args _ => do
-        let fieldArgs := args.extract iv.numParams args.size
-        fieldArgs.mapM fun a => do
-          let ty ← inferType a
-          match ty.getAppFn.constName? with
-          | some hn =>
-            if famSet.contains hn then
-              return { isRec := true, recType := hn, typeStx := ← `($(mkIdent (taggedOf hn)) $τ) }
-            else
-              return { isRec := false, recType := default, typeStx := ← Lean.PrettyPrinter.delab ty }
-          | none =>
-            return { isRec := false, recType := default, typeStx := ← Lean.PrettyPrinter.delab ty }
-      return { origName := cn, shortName := cn.componentsRev.head!, fields }
-    return { name := tn, ctors }
-
-/-- The arrow type `τ → <fields…> → Self τ` of a tagged constructor. -/
-def ctorArrow (cd : CtorData) (self : Term) (τ : Ident) : TermElabM Term := do
-  let mut t := self
-  for f in cd.fields.reverse do
-    t ← `($(f.typeStx) → $t)
-  `($τ → $t)
-
-/-- The tagged mirror inductives, one per family member.  The family is a DAG
-(`Expr ← BasicStmt ← Stmt`), not genuinely mutual, so they are emitted as
-separate inductives in dependency order rather than a `mutual` block.
-
-`Functor`/`Traversable` instances are derived separately by `mkDeriveInstances`
-below rather than via an inline `deriving` clause. -/
-def mkInductives (tds : Array TypeData) (τ : Ident) :
-    CommandElabM (Array (TSyntax `command)) := do
-  tds.mapM fun td => do
-    let self ← `($(mkIdent (taggedOf td.name)) $τ)
-    let ctors ← td.ctors.mapM fun cd => do
-      let aty ← Command.liftTermElabM (ctorArrow cd self τ)
-      `(Lean.Parser.Command.ctor| | $(mkIdent cd.shortName):ident : $aty)
-    `(command| inductive $(mkIdent (taggedOf td.name)):ident ($τ : Type) where $ctors*)
-
-/-- A `deriving instance Functor, Traversable for Tᵢ.Tagged` command per family
-member.  Since every tagged type is a single-parameter, direct-recursive
-inductive in `τ`, Mathlib's deriving handler produces clean (`sorry`-free)
-instances, giving `map`, `traverse`, and the `Traversable.foldr`/`toList` folds
-for free.
-
-These are emitted as *separate* commands in dependency order (rather than an
-inline `deriving` clause on each inductive) for two reasons: deriving
-`Stmt.Tagged` needs the `Expr.Tagged`/`BasicStmt.Tagged` instances already in
-scope, and — because every member's type name ends in `.Tagged` — the handler's
-auto-generated instance name (`instFunctorTagged`, built from the type's last
-component) collides across the family unless each derive sees the environment
-the previous one updated; separate commands give it that, so the names
-disambiguate to `instFunctorTagged`, `instFunctorTagged_1`, ….
-
-The hand-written `foldTags` is retained alongside these: it is a
-structural-recursion fold that `simp`/`decide` reduce cleanly, unlike the
-abstract `Traversable.foldr` (defined via the `FreeMonoid`/`Const` applicative),
-which reduces under `decide`/`rfl` but not naive `simp` unfolding. -/
-def mkDeriveInstances (tds : Array TypeData) : CommandElabM (Array (TSyntax `command)) := do
-  tds.mapM fun td =>
-    `(command| deriving instance Functor, Traversable for $(mkIdent (taggedOf td.name)))
-
-/-- The `erase` functions, one per family member (separate defs in dependency
-order — each calls only already-defined lower members). -/
-def mkErase (tds : Array TypeData) : CommandElabM (Array (TSyntax `command)) := do
-  tds.mapM fun td => do
-    let mut pats : Array Term := #[]
-    let mut rhss : Array Term := #[]
-    for cd in td.ctors do
-      let argNames := (Array.range cd.fields.size).map (fun i => mkIdent (.mkSimple s!"a{i}"))
-      let pat ← `($(mkIdent (taggedOf td.name ++ cd.shortName)) _ $argNames*)
-      let eraseArgs ← (cd.fields.zip argNames).mapM fun (f, a) =>
-        if f.isRec then `($(mkIdent (eraseOf f.recType)) $a) else pure a
-      let rhs ← `($(mkIdent cd.origName) $eraseArgs*)
-      pats := pats.push pat
-      rhss := rhss.push rhs
-    `(command| def $(mkIdent (eraseOf td.name)) {τ : Type} :
-        $(mkIdent (taggedOf td.name)) τ → $(mkIdent td.name) :=
-        fun x => match x with $[| $pats => $rhss]*)
-
-/-- The `rootTag` accessors (one non-recursive `def` per type). -/
-def mkRootTag (tds : Array TypeData) : CommandElabM (Array (TSyntax `command)) := do
-  let tIdent := mkIdent `t
-  tds.mapM fun td => do
-    let mut pats : Array Term := #[]
-    let mut rhss : Array Term := #[]
-    for cd in td.ctors do
-      let hole ← `(_)
-      let wilds := Array.mkArray cd.fields.size hole
-      pats := pats.push (← `($(mkIdent (taggedOf td.name ++ cd.shortName)) $tIdent $wilds*))
-      rhss := rhss.push tIdent
-    `(command| def $(mkIdent (rootTagOf td.name)) {τ : Type} :
-        $(mkIdent (taggedOf td.name)) τ → τ :=
-        fun x => match x with $[| $pats => $rhss]*)
-
-/-- The postorder `tag` functions, one per family member (separate defs in
-dependency order). -/
-def mkTag (tds : Array TypeData) : CommandElabM (Array (TSyntax `command)) := do
-  let nId := mkIdent ``Spa.RawId
-  tds.mapM fun td => do
-    let mut pats : Array Term := #[]
-    let mut rhss : Array Term := #[]
-    for cd in td.ctors do
-      let argNames := (Array.range cd.fields.size).map (fun i => mkIdent (.mkSimple s!"a{i}"))
-      let pat ← `($(mkIdent cd.origName) $argNames*)
-      let mut cur : Term ← `(n)
-      let mut lets : Array (Ident × Term) := #[]
-      let mut taggedArgs : Array Term := #[]
-      let mut ri := 0
-      for (f, a) in cd.fields.zip argNames do
-        if f.isRec then
-          let rName := mkIdent (.mkSimple s!"r{ri}")
-          let rhsCall ← `($(mkIdent (tagOf f.recType)) $a $cur)
-          lets := lets.push (rName, rhsCall)
-          taggedArgs := taggedArgs.push (← `($rName |>.1))
-          cur ← `($rName |>.2)
-          ri := ri + 1
-        else
-          taggedArgs := taggedArgs.push a
-      let last := cur
-      let tagged ← `($(mkIdent (taggedOf td.name ++ cd.shortName))
-          (⟨$last⟩ : $nId) $taggedArgs*)
-      let mut body ← `(($tagged, $last + 1))
-      for (rName, rhs) in lets.reverse do
-        body ← `(let $rName := $rhs; $body)
-      pats := pats.push pat
-      rhss := rhss.push body
-    `(command| def $(mkIdent (tagOf td.name)) :
-        $(mkIdent td.name) → Nat → $(mkIdent (taggedOf td.name)) $nId × Nat :=
-        fun e n => match e with $[| $pats => $rhss]*)
-
-/-- The tag-fold functions: `foldTags f acc t` applies `f` to every tag in `t`,
-right-to-left, threading `acc`.  This is the `Foldable`/`foldr`-over-tags the
-hand-written collectors (e.g. `subtreeIds`) reduce to.  One separate def per
-family member (the family is a DAG, so no `mutual` block is needed). -/
-def mkFoldTags (tds : Array TypeData) : CommandElabM (Array (TSyntax `command)) := do
-  let τ := mkIdent `τ
-  let m := mkIdent `M
-  let fId := mkIdent `f
-  let accId := mkIdent `acc
-  let tagId := mkIdent `t
-  tds.mapM fun td => do
-    let mut pats : Array Term := #[]
-    let mut rhss : Array Term := #[]
-    for cd in td.ctors do
-      let argNames := (Array.range cd.fields.size).map (fun i => mkIdent (.mkSimple s!"a{i}"))
-      let pat ← `($(mkIdent (taggedOf td.name ++ cd.shortName)) $tagId $argNames*)
-      let mut body : Term := accId
-      for (fld, a) in (cd.fields.zip argNames).reverse do
-        if fld.isRec then
-          body ← `($(mkIdent (foldTagsOf fld.recType)) $fId $body $a)
-      body ← `($fId $tagId $body)
-      pats := pats.push pat
-      rhss := rhss.push body
-    `(command| def $(mkIdent (foldTagsOf td.name)) {$τ:ident : Type} {$m:ident : Type}
-        ($fId : $τ → $m → $m) ($accId : $m) :
-        $(mkIdent (taggedOf td.name)) $τ → $m :=
-        fun x => match x with $[| $pats => $rhss]*)
-
-/-- The well-formedness predicate `T.Tagged.WF : T.Tagged RawId → Prop`: every
-recursive child's root tag has a strictly smaller postorder index than the node's
-own tag, and each child is itself well-formed.  Leaf constructors are `True`. -/
-def mkWF (tds : Array TypeData) : CommandElabM (Array (TSyntax `command)) := do
-  let tId := mkIdent `t
-  let rawId := mkIdent ``Spa.RawId
-  tds.mapM fun td => do
-    let mut pats : Array Term := #[]
-    let mut rhss : Array Term := #[]
-    for cd in td.ctors do
-      let hasRec := cd.fields.any (·.isRec)
-      let mut patArgs : Array Term := #[]
-      let mut recArgs : Array Ident := #[]
-      let mut i := 0
-      for f in cd.fields do
-        if f.isRec then
-          let a := mkIdent (.mkSimple s!"a{i}")
-          patArgs := patArgs.push a
-          recArgs := recArgs.push a
-        else
-          patArgs := patArgs.push (← `(_))
-        i := i + 1
-      let tagBind : Term ← if hasRec then `($tId) else `(_)
-      let pat ← `($(mkIdent (taggedOf td.name ++ cd.shortName)) $tagBind $patArgs*)
-      let rhs ← if recArgs.isEmpty then `(True) else do
-        let bounds ← recArgs.mapM fun a => `($(a).rootTag.post < $(tId).post)
-        let wfs ← recArgs.mapM fun a => `($(a).WF)
-        mkAndR (bounds ++ wfs)
-      pats := pats.push pat
-      rhss := rhss.push rhs
-    `(command| def $(mkIdent (wfOf td.name)) :
-        $(mkIdent (taggedOf td.name)) $rawId → Prop :=
-        fun x => match x with $[| $pats => $rhss]*)
-
-/-- The `narrow` coercion `T.Tagged RawId → T.Tagged (Fin N)`, given a bound on
-the root tag and a well-formedness proof.  Each node's tag becomes the `Fin N`
-built from its postorder index, and recursion threads the bound through `lt_trans`
-and the (definitionally unfolded) `WF` conjunction. -/
-def mkNarrow (tds : Array TypeData) : CommandElabM (Array (TSyntax `command)) := do
-  let rawId := mkIdent ``Spa.RawId
-  let tId := mkIdent `t
-  let nId := mkIdent `N
-  let hId := mkIdent `h
-  let hwfId := mkIdent `hwf
-  let tgId := mkIdent `tg
-  tds.mapM fun td => do
-    let self ← `($(mkIdent (taggedOf td.name)) $rawId)
-    let mut patss : Array (Array Term) := #[]
-    let mut rhss : Array Term := #[]
-    for cd in td.ctors do
-      let argNames := (Array.range cd.fields.size).map fun i => mkIdent (.mkSimple s!"a{i}")
-      let ctorPat ← `($(mkIdent (taggedOf td.name ++ cd.shortName)) $tgId $argNames*)
-      let k := (cd.fields.filter (·.isRec)).size
-      let mut newArgs : Array Term := #[]
-      let mut ri := 0
-      for (f, a) in cd.fields.zip argNames do
-        if f.isRec then
-          let bound ← projAnd hwfId (ri + 1) (2 * k)
-          let wf ← projAnd hwfId (k + ri + 1) (2 * k)
-          newArgs := newArgs.push (← `($(a).narrow (lt_trans $bound $hId) $wf))
-          ri := ri + 1
-        else
-          newArgs := newArgs.push a
-      let built ← `($(mkIdent (taggedOf td.name ++ cd.shortName)) ⟨$(tgId).post, $hId⟩ $newArgs*)
-      let nPat ← `(_)
-      let hPat ← `($hId)
-      let hwfPat : Term ← if k == 0 then `(_) else `($hwfId)
-      patss := patss.push #[ctorPat, nPat, hPat, hwfPat]
-      rhss := rhss.push built
-    `(command| def $(mkIdent (narrowOf td.name)) : ($tId : $self) → {$nId : ℕ} →
-        $(tId).rootTag.post < $nId → $(tId).WF → $(mkIdent (taggedOf td.name)) (Fin $nId)
-        $[| $[$patss],* => $rhss]*)
-
-/-- `T.tag_rootTag_post`: the root tag of a freshly tagged node is exactly one
-below the threaded-out counter, i.e. the node itself is numbered last (postorder).
-A uniform `cases <;> simp` discharges every constructor. -/
-def mkTagRootTagPost (tds : Array TypeData) : CommandElabM (Array (TSyntax `command)) := do
-  let eId := mkIdent `e
-  let nId := mkIdent `n
-  tds.mapM fun td =>
-    `(command| theorem $(mkIdent (tagRootTagPostOf td.name))
-        ($eId : $(mkIdent td.name)) ($nId : ℕ) :
-        ($(eId).tag $nId).1.rootTag.post + 1 = ($(eId).tag $nId).2 := by
-        cases $eId:ident <;>
-          simp [$(mkIdent (tagOf td.name)):ident, $(mkIdent (rootTagOf td.name)):ident])
-
-/-- `T.tag_le`: tagging only ever advances the counter (`n ≤ (e.tag n).2`).
-Proved by induction; each arm threads the counter through its recursive children
-(using the relevant `tag_le`/induction hypothesis) and closes with `omega`. -/
-def mkTagLe (tds : Array TypeData) : CommandElabM (Array (TSyntax `command)) := do
-  let eId := mkIdent `e
-  let nId := mkIdent `n
-  tds.mapM fun td => do
-    let mut ctorLabels : Array Ident := #[]
-    let mut binderss : Array (Array Ident) := #[]
-    let mut tacs : Array (TSyntax ``Lean.Parser.Tactic.tacticSeq) := #[]
-    for cd in td.ctors do
-      let argNames := (Array.range cd.fields.size).map fun i => mkIdent (.mkSimple s!"a{i}")
-      let mut ihBinders : Array Ident := #[]
-      let mut haveTacs : Array (TSyntax `tactic) := #[]
-      let mut cur : Term ← `($nId)
-      let mut i := 0
-      for (f, a) in cd.fields.zip argNames do
-        if f.isRec then
-          let fact ← if f.recType == td.name then
-              `($(mkIdent (.mkSimple s!"ih{i}")) $cur)
-            else
-              `($(mkIdent (tagLeOf f.recType)) $a $cur)
-          if f.recType == td.name then
-            ihBinders := ihBinders.push (mkIdent (.mkSimple s!"ih{i}"))
-          haveTacs := haveTacs.push (← `(tactic| have := $fact))
-          cur ← `(($(mkIdent (tagOf f.recType)) $a $cur) |>.2)
-        i := i + 1
-      let simpTac ← `(tactic| simp only [$(mkIdent (tagOf td.name)):ident])
-      let omegaTac ← `(tactic| omega)
-      let allTacs := #[simpTac] ++ haveTacs ++ #[omegaTac]
-      ctorLabels := ctorLabels.push (mkIdent cd.shortName)
-      binderss := binderss.push (argNames ++ ihBinders)
-      tacs := tacs.push (← `(tacticSeq| $[$allTacs]*))
-    `(command| theorem $(mkIdent (tagLeOf td.name)) ($eId : $(mkIdent td.name)) ($nId : ℕ) :
-        $nId ≤ ($(eId).tag $nId).2 := by
-        induction $eId:ident generalizing $nId:ident with
-          $[| $ctorLabels:ident $binderss* => $tacs]*)
-
-/-- `T.tag_wf`: a freshly tagged term is well-formed.  Each recursive child's
-bound conjunct is closed by `omega` from that child's `tag_rootTag_post` plus the
-`tag_le` of every later child (which bounds the threaded-out counter), and each
-well-formedness conjunct is the child's induction hypothesis / `tag_wf`. -/
-def mkTagWf (tds : Array TypeData) : CommandElabM (Array (TSyntax `command)) := do
-  let eId := mkIdent `e
-  let nId := mkIdent `n
-  tds.mapM fun td => do
-    let mut ctorLabels : Array Ident := #[]
-    let mut binderss : Array (Array Ident) := #[]
-    let mut tacs : Array (TSyntax ``Lean.Parser.Tactic.tacticSeq) := #[]
-    for cd in td.ctors do
-      let argNames := (Array.range cd.fields.size).map fun i => mkIdent (.mkSimple s!"a{i}")
-      -- recursive children: (arg, recType, startCounter, sameType?, fieldIndex)
-      let mut recs : Array (Ident × Name × Term × Bool × Nat) := #[]
-      let mut cur : Term ← `($nId)
-      let mut i := 0
-      for (f, a) in cd.fields.zip argNames do
-        if f.isRec then
-          recs := recs.push (a, f.recType, cur, f.recType == td.name, i)
-          cur ← `(($(mkIdent (tagOf f.recType)) $a $cur) |>.2)
-        i := i + 1
-      let k := recs.size
-      let ihBinders := (recs.filter (·.2.2.2.1)).map fun r => mkIdent (.mkSimple s!"ih{r.2.2.2.2}")
-      let tac : TSyntax ``Lean.Parser.Tactic.tacticSeq ← if k == 0 then
-          `(tacticSeq| exact True.intro)
-        else do
-          let mut comps : Array Term := #[]
-          -- bound conjuncts
-          for idx in [0:k] do
-            let (a, rt, s, _, _) := recs[idx]!
-            let mut bHaves : Array (TSyntax `tactic) :=
-              #[← `(tactic| have := $(mkIdent (tagRootTagPostOf rt)) $a $s)]
-            for j in [idx+1:k] do
-              let (aj, rtj, sj, _, _) := recs[j]!
-              bHaves := bHaves.push (← `(tactic| have := $(mkIdent (tagLeOf rtj)) $aj $sj))
-            bHaves := bHaves.push (← `(tactic| omega))
-            comps := comps.push (← `(by $(← `(tacticSeq| $[$bHaves]*))))
-          -- well-formedness conjuncts
-          for idx in [0:k] do
-            let (a, rt, s, same, fi) := recs[idx]!
-            comps := comps.push <| ← if same then `($(mkIdent (.mkSimple s!"ih{fi}")) $s)
-              else `($(mkIdent (tagWfOf rt)) $a $s)
-          let simpTac ← `(tactic| simp only
-            [$(mkIdent (tagOf td.name)):ident, $(mkIdent (wfOf td.name)):ident])
-          let exactTac ← `(tactic| exact ⟨$comps,*⟩)
-          `(tacticSeq| $[$(#[simpTac, exactTac])]*)
-      ctorLabels := ctorLabels.push (mkIdent cd.shortName)
-      binderss := binderss.push (argNames ++ ihBinders)
-      tacs := tacs.push tac
-    `(command| theorem $(mkIdent (tagWfOf td.name)) ($eId : $(mkIdent td.name)) ($nId : ℕ) :
-        ($(eId).tag $nId).1.WF := by
-        induction $eId:ident generalizing $nId:ident with
-          $[| $ctorLabels:ident $binderss* => $tacs]*)
-
-/-- `T.Tagged.narrow_erase`: narrowing the tag type does not change the erased
-(untagged) term.  A per-constructor `simp` with the local `narrow`/`erase`
-equations, the lower members' `narrow_erase`, and the induction hypotheses. -/
-def mkNarrowErase (tds : Array TypeData) : CommandElabM (Array (TSyntax `command)) := do
-  let rawId := mkIdent ``Spa.RawId
-  let tId := mkIdent `t
-  let nId := mkIdent `N
-  let hId := mkIdent `h
-  let hwfId := mkIdent `hwf
-  let tgId := mkIdent `tg
-  tds.mapM fun td => do
-    let mut ctorLabels : Array Ident := #[]
-    let mut binderss : Array (Array Ident) := #[]
-    let mut tacs : Array (TSyntax ``Lean.Parser.Tactic.tacticSeq) := #[]
-    for cd in td.ctors do
-      let argNames := (Array.range cd.fields.size).map fun i => mkIdent (.mkSimple s!"a{i}")
-      let mut lemmas : Array Term :=
-        #[← `($(mkIdent (narrowOf td.name))), ← `($(mkIdent (eraseOf td.name)))]
-      let mut ihBinders : Array Ident := #[]
-      let mut seenLower : Array Name := #[]
-      let mut i := 0
-      for f in cd.fields do
-        if f.isRec then
-          if f.recType == td.name then
-            let ih := mkIdent (.mkSimple s!"ih{i}")
-            ihBinders := ihBinders.push ih
-            lemmas := lemmas.push (← `($ih))
-          else if !seenLower.contains f.recType then
-            seenLower := seenLower.push f.recType
-            lemmas := lemmas.push (← `($(mkIdent (narrowEraseOf f.recType))))
-        i := i + 1
-      let introTac ← `(tactic| intro $nId $hId $hwfId)
-      let simpTac ← `(tactic| simp [$[$lemmas:term],*])
-      ctorLabels := ctorLabels.push (mkIdent cd.shortName)
-      binderss := binderss.push (#[tgId] ++ argNames ++ ihBinders)
-      tacs := tacs.push (← `(tacticSeq| $[$(#[introTac, simpTac])]*))
-    `(command| theorem $(mkIdent (narrowEraseOf td.name)) :
-        ($tId : $(mkIdent (taggedOf td.name)) $rawId) → ∀ {$nId : ℕ}
-          ($hId : $(tId).rootTag.post < $nId) ($hwfId : $(tId).WF),
-          ($(tId).narrow $hId $hwfId).erase = $(tId).erase := by
-        intro $tId:ident
-        induction $tId:ident with
-          $[| $ctorLabels:ident $binderss* => $tacs]*)
-
-/-- `derive_tagged T₁ … Tₙ` — generate tagged mirrors, `erase`, and `tag` for the
-given family of inductives. -/
-syntax (name := deriveTaggedCmd) "derive_tagged " ident+ : command
-
-@[command_elab deriveTaggedCmd]
-def elabDeriveTagged : CommandElab := fun stx => do
-  match stx with
-  | `(derive_tagged $ids*) =>
-    let family ← ids.mapM fun i => Command.liftCoreM (realizeGlobalConstNoOverload i)
-    let τ := mkIdent `τ
-    let tds ← Command.liftTermElabM (gather family τ)
-    for d in (← mkInductives tds τ) do elabCommand d
-    for d in (← mkDeriveInstances tds) do elabCommand d
-    for d in (← mkRootTag tds) do elabCommand d
-    for d in (← mkErase tds) do elabCommand d
-    for d in (← mkTag tds) do elabCommand d
-    for d in (← mkFoldTags tds) do elabCommand d
-    for d in (← mkWF tds) do elabCommand d
-    for d in (← mkNarrow tds) do elabCommand d
-    for d in (← mkTagRootTagPost tds) do elabCommand d
-    for d in (← mkTagLe tds) do elabCommand d
-    for d in (← mkTagWf tds) do elabCommand d
-    for d in (← mkNarrowErase tds) do elabCommand d
-  | _ => throwUnsupportedSyntax
-
-end Spa.DeriveTagged

lean/Spa/Language/Tagged/Graphs.lean

@@ -1,104 +0,0 @@
-import Spa.Language
-import Spa.Language.Graphs
-import Spa.Language.Tagged.Basic
-import Spa.Language.Tagged.Properties
-
-namespace Spa
-
-open GGraph
-
-def Stmt.Tagged.cfg {τ : Type} : Stmt.Tagged τ → GGraph (Option (BasicStmt.Tagged τ))
-  | .basic _ bs => GGraph.singleton (some bs)
-  | .andThen _ s₁ s₂ => s₁.cfg ⤳  s₂.cfg
-  | .ifElse _ _ s₁ s₂ => s₁.cfg ∙ s₂.cfg
-  | .whileLoop _ _ s => GGraph.loop s.cfg
-
-theorem Stmt.Tagged.cfg_graph {τ : Type} : ∀ (t : Stmt.Tagged τ),
-    (Option.map BasicStmt.Tagged.erase) <$> t.cfg = t.erase.cfg
-  | .basic _ bs => by simp [Stmt.Tagged.cfg, Stmt.cfg, Stmt.Tagged.erase, BasicStmt.Tagged.erase]
-  | .andThen _ s₁ s₂ => by
-      simp [Stmt.Tagged.cfg, Stmt.cfg, Stmt.Tagged.erase, Stmt.Tagged.cfg_graph s₁, Stmt.Tagged.cfg_graph s₂]
-  | .ifElse _ _ s₁ s₂ => by
-      simp [Stmt.Tagged.cfg, Stmt.cfg, Stmt.Tagged.erase, Stmt.Tagged.cfg_graph s₁, Stmt.Tagged.cfg_graph s₂]
-  | .whileLoop _ _ s => by
-      simp [Stmt.Tagged.cfg, Stmt.cfg, Stmt.Tagged.erase, Stmt.Tagged.cfg_graph s]
-
-def GGraph.nodeLabel {τ : Type} (g : GGraph (Option (BasicStmt.Tagged τ))) (i : g.Index) :
-    Option τ :=
-  (g.nodes i).map BasicStmt.Tagged.rootTag
-
-def GGraph.stateOf {τ : Type} [DecidableEq τ] (g : GGraph (Option (BasicStmt.Tagged τ)))
-    (id : τ) : Option g.Index :=
-  g.indices.find? (fun i => decide (g.nodeLabel i = some id))
-
-theorem GGraph.stateOf_label {τ : Type} [DecidableEq τ]
-    {g : GGraph (Option (BasicStmt.Tagged τ))} {id : τ}
-    {i : g.Index} (h : g.stateOf id = some i) : g.nodeLabel i = some id := by
-  rw [GGraph.stateOf] at h
-  simpa using List.find?_some h
-
-namespace Program
-
-variable (p : Program)
-
-def tagged : Stmt.Tagged RawId := tagStmt p.rootStmt
-
-def size : ℕ := p.tagged.rootTag.post + 1
-
-theorem size_pos : 0 < p.size := Nat.succ_pos _
-
-abbrev NodeId : Type := Fin p.size
-
-theorem tagged_wf : p.tagged.WF := Stmt.tag_wf p.rootStmt 0
-
-def taggedFin : Stmt.Tagged p.NodeId :=
-  p.tagged.narrow (Nat.lt_succ_self _) p.tagged_wf
-
-def taggedCfg : GGraph (Option (BasicStmt.Tagged p.NodeId)) :=
-  GGraph.wrap p.taggedFin.cfg
-
-theorem taggedCfg_erase :
-    (Option.map BasicStmt.Tagged.erase) <$> p.taggedCfg = p.cfg := by
-  rw [taggedCfg, GGraph.map_wrap, Stmt.Tagged.cfg_graph, taggedFin,
-    Stmt.Tagged.narrow_erase, tagged, erase_tagStmt]
-  rfl
-
-theorem taggedCfg_size : p.taggedCfg.size = p.cfg.size := by
-  conv_rhs => rw [← p.taggedCfg_erase]
-  rfl
-
-def nodeIdOf (s : p.State) : Option p.NodeId :=
-  p.taggedCfg.nodeLabel (Fin.cast p.taggedCfg_size.symm s)
-
-def stateOfNodeId (id : p.NodeId) : Option p.State :=
-  (p.taggedCfg.stateOf id).map (Fin.cast p.taggedCfg_size)
-
-theorem cfg_nodes_eq (s : p.State) :
-    p.cfg.nodes s = Option.map BasicStmt.Tagged.erase
-      (p.taggedCfg.nodes (Fin.cast p.taggedCfg_size.symm s)) := by
-  have key : ∀ (g : Graph) (hsz : p.taggedCfg.size = g.size),
-      (Option.map BasicStmt.Tagged.erase) <$> p.taggedCfg = g →
-      ∀ i : Fin g.size,
-        g.nodes i = Option.map BasicStmt.Tagged.erase
-          (p.taggedCfg.nodes (Fin.cast hsz.symm i)) := by
-    intro g hsz hg i
-    subst hg
-    rfl
-  exact key p.cfg p.taggedCfg_size p.taggedCfg_erase s
-
-theorem nodeIdOf_isSome_of_code {s : p.State} {bs : BasicStmt}
-    (h : p.code s = some bs) : (p.nodeIdOf s).isSome = true := by
-  have hc : Option.map BasicStmt.Tagged.erase
-      (p.taggedCfg.nodes (Fin.cast p.taggedCfg_size.symm s)) = some bs := by
-    rw [← p.cfg_nodes_eq s]; exact h
-  unfold Program.nodeIdOf GGraph.nodeLabel
-  cases hcase : p.taggedCfg.nodes (Fin.cast p.taggedCfg_size.symm s) with
-  | none => rw [hcase] at hc; simp at hc
-  | some tbs => simp
-
-def nodeIdOfNonempty (s : p.State) {bs : BasicStmt} (h : p.code s = some bs) : p.NodeId :=
-  (p.nodeIdOf s).get (p.nodeIdOf_isSome_of_code h)
-
-end Program
-
-end Spa

lean/Spa/Language/Tagged/Id.lean

@@ -1,9 +0,0 @@
-import Mathlib.Data.Nat.Notation
-
-namespace Spa
-
-structure RawId where
-  post : ℕ
-  deriving DecidableEq, Repr
-
-end Spa

lean/Spa/Language/Tagged/Properties.lean

@@ -1,29 +0,0 @@
-import Spa.Language.Tagged.Basic
-
-namespace Spa
-
-@[simp] theorem Expr.erase_tag (e : Expr) (n : ℕ) : (e.tag n).1.erase = e := by
-  induction e generalizing n with
-  | add a b iha ihb => simp [Expr.tag, Expr.Tagged.erase, iha, ihb]
-  | sub a b iha ihb => simp [Expr.tag, Expr.Tagged.erase, iha, ihb]
-  | var x => simp [Expr.tag, Expr.Tagged.erase]
-  | num k => simp [Expr.tag, Expr.Tagged.erase]
-
-@[simp] theorem BasicStmt.erase_tag (bs : BasicStmt) (n : ℕ) :
-    (bs.tag n).1.erase = bs := by
-  cases bs with
-  | assign x e => simp [BasicStmt.tag, BasicStmt.Tagged.erase]
-  | noop => simp [BasicStmt.tag, BasicStmt.Tagged.erase]
-
-@[simp] theorem Stmt.erase_tag (s : Stmt) (n : ℕ) : (s.tag n).1.erase = s := by
-  induction s generalizing n with
-  | basic bs => simp [Stmt.tag, Stmt.Tagged.erase]
-  | andThen a b iha ihb => simp [Stmt.tag, Stmt.Tagged.erase, iha, ihb]
-  | ifElse e a b iha ihb => simp [Stmt.tag, Stmt.Tagged.erase, iha, ihb]
-  | whileLoop e s ih => simp [Stmt.tag, Stmt.Tagged.erase, ih]
-
-/-- Erasing a freshly tagged program recovers it. -/
-theorem erase_tagStmt (s : Stmt) : (tagStmt s).erase = s := by
-  simp [tagStmt]
-
-end Spa

lean/Spa/Language/Traces.lean

@@ -217,50 +217,30 @@ inductive EndToEndTrace (g : Graph) (ρ₁ ρ₂ : Env) : Type
       (idx₂ : g.Index) (idx₂_mem : idx₂ ∈ g.outputs)
       (trace : Trace g idx₁ idx₂ ρ₁ ρ₂) : EndToEndTrace g ρ₁ ρ₂
 
-inductive Reaches {prog : Program} : {s₁ s₂ : prog.State} → {ρ₁ ρ₂ : Env} →
-    Trace prog.cfg s₁ s₂ ρ₁ ρ₂ →
-    (s : prog.State) → (ρin ρout : Env) → Type
-  | single_here {s₁ : prog.State} {ρ₁ ρ₂ : Env}
-      (hnode : EvalBasicStmtOpt ρ₁ (prog.code s₁) ρ₂) :
-      Reaches (.single hnode) s₁ ρ₁ ρ₂
-  | edge_here {s₁ s₂ s₃ : prog.State} {ρ₁ ρ₂ ρ₃ : Env}
-      (hnode : EvalBasicStmtOpt ρ₁ (prog.code s₁) ρ₂)
-      (hedge : (s₁, s₂) ∈ prog.cfg.edges) (rest : Trace prog.cfg s₂ s₃ ρ₂ ρ₃) :
-      Reaches (.edge hnode hedge rest) s₁ ρ₁ ρ₂
-  | edge_there {s₁ s₂ s₃ : prog.State} {ρ₁ ρ₂ ρ₃ : Env}
-      (hnode : EvalBasicStmtOpt ρ₁ (prog.code s₁) ρ₂)
-      (hedge : (s₁, s₂) ∈ prog.cfg.edges) (rest : Trace prog.cfg s₂ s₃ ρ₂ ρ₃)
-      {s : prog.State} {ρin ρout : Env} :
-      Reaches rest s ρin ρout →
-      Reaches (.edge hnode hedge rest) s ρin ρout
+/-- Every trace splits into the prefix that arrives at its last node and that node's own step. -/
+def Trace.split {g : Graph} {i₁ i₂ : g.Index} {ρ₁ ρ₂ : Env} :
+    Trace g i₁ i₂ ρ₁ ρ₂ → Σ ρ, Traceₗ g i₁ i₂ ρ₁ ρ × EvalBasicStmtOpt ρ (g.nodes i₂) ρ₂
+  | .single hnode => ⟨_, .nil, hnode⟩
+  | .edge hnode hedge rest =>
+      let ⟨ρ, pre, step⟩ := rest.split
+      ⟨ρ, .cons hnode hedge pre, step⟩
 
-def Reaches.pre {prog : Program} {s₁ s₂ s: prog.State}
-    {ρ₁ ρ₂ ρin ρout : Env} {tr : Trace prog.cfg s₁ s₂ ρ₁ ρ₂} :
-    (r : Reaches tr s ρin ρout) → Traceₗ prog.cfg s₁ s ρ₁ ρin
-  | .single_here _ => .nil
-  | .edge_here _ _ _ => .nil
-  | .edge_there hnode hedge _ r => .cons hnode hedge r.pre
+@[simp] lemma Trace.split_append {g : Graph} {i₁ i₂ : g.Index} {ρ₁ ρ₂ : Env}
+    (tr : Trace g i₁ i₂ ρ₁ ρ₂) : tr.split.2.1 ++ tr.split.2.2 = tr := by
+  induction tr with
+  | single hnode => rfl
+  | edge hnode hedge rest ih =>
+      show Traceₗ.appendStep _ _ = _
+      simpa [Trace.split, Traceₗ.appendStep, Traceₗ.appendTrace] using ih
 
-def Reaches.post {prog : Program} {s₁ s₂ s: prog.State}
-    {ρ₁ ρ₂ ρin ρout : Env} {tr : Trace prog.cfg s₁ s₂ ρ₁ ρ₂} :
-    (r : Reaches tr s ρin ρout) → Trace prog.cfg s₁ s ρ₁ ρout
-  | .single_here hnode => .single hnode
-  | .edge_here hnode _ _ => .single hnode
-  | .edge_there hnode hedge _ r => .edge hnode hedge r.post
+structure Reaches {prog : Program} (s : prog.State) (ρin ρout : Env) : Type where
+  pre : Traceₗ prog.cfg prog.initialState s [] ρin
+  step : EvalBasicStmtOpt ρin (prog.code s) ρout
 
-def Reaches.first {prog : Program} {s₁ s₂ s: prog.State}
-    {ρ₁ ρ₂ ρin ρout : Env} {tr : Trace prog.cfg s₁ s₂ ρ₁ ρ₂} :
-    (r : Reaches tr s ρin ρout) → Σ ρ₁', Reaches tr s₁ ρ₁ ρ₁'
-  | .single_here hnode => ⟨_, .single_here hnode⟩
-  | .edge_here hnode hedge hrest => ⟨_, .edge_here hnode hedge hrest⟩
-  | .edge_there hnode hedge hrest tmp' => ⟨_, .edge_here hnode hedge hrest⟩
-
-def Reaches.step {prog : Program} {s₁ s₂ s: prog.State}
-    {ρ₁ ρ₂ ρin ρout : Env} {tr : Trace prog.cfg s₁ s₂ ρ₁ ρ₂} :
-    (r : Reaches tr s ρin ρout) → EvalBasicStmtOpt ρin (prog.code s) ρout
-  | .single_here hnode => hnode
-  | .edge_here hnode hedge hrest => hnode
-  | .edge_there hnode hedge hrest tmp' => tmp'.step
+/-- Forget the environment before the last evaluated state. -/
+def Reaches.post {prog : Program} {s : prog.State} {ρin ρout : Env}
+    (r : Reaches s ρin ρout) : Trace prog.cfg prog.initialState s [] ρout :=
+  r.pre ++ r.step
 
 
 end Spa

lean/Spa/Transformation/Licm.lean

@@ -1,26 +1,30 @@
 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`
-loop whose right-hand side depends only on definitions made *outside* the loop.
-These are the candidates a later LICM pass could hoist.
+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 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`);
-2. read the reaching definitions at the assignment's *entry*
-   (`joinForKey s result` — the join over predecessors, i.e. before the
-   assignment itself runs);
+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. map each definition site back to its `RawId` (`Program.nodeIdOf`) and check
-   it is **not** inside the loop body (structural `subtreeIds` membership).
+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
+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.
 -/
@@ -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). -/
-  loopId : prog.NodeId
-  /-- Every node id inside the loop body (the "is-child-of-loop" set). -/
-  bodyIds : List prog.NodeId
-  /-- The assignment `BasicStmt`'s tag — what labels its CFG node. -/
-  assignId : prog.NodeId
+  /-- 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.
-`enclosing` carries the current loop's tag and body id-set, or `none` outside any
-loop (in which case assignments are skipped — only in-loop assignments are
-candidates). -/
-def collectCandidates (prog : Program) (enc : Option (prog.NodeId × List prog.NodeId)) :
-    Stmt.Tagged prog.NodeId → List (Candidate prog)
-  | .basic _ bs =>
+`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 t _ e, some (loopId, bodyIds) =>
-          [{ loopId := loopId, bodyIds := bodyIds, assignId := t,
-             rhsVars := e.erase.vars.sort (· ≤ ·) }]
+      | .assign _ ex, some l =>
+          [{ encl := l, assignState := e.singletonIndex,
+             rhsVars := ex.vars.sort (· ≤ ·) }]
       | _, _ => []
-  | .andThen _ a b => collectCandidates prog enc a ++ collectCandidates prog enc b
-  | .ifElse _ _ a b => collectCandidates prog enc a ++ collectCandidates prog enc b
-  | .whileLoop loopT _ body =>
-      collectCandidates prog (some (loopT, body.subtreeIds)) body
+  | .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 ⊥
 
-/-- 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
-node id, so we compare against the loop-body ids directly (embedding the raw
-body ids into `p.NodeId`). -/
+its RHS variables lie outside the loop body? -/
 def isInvariant (prog : Program) (c : Candidate prog) : Bool :=
-  match prog.stateOfNodeId c.assignId with
-  | 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))
+  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 `(loopId, assignId)` pairs. -/
-def licmCandidates (prog : Program) : List (prog.NodeId × prog.NodeId) :=
-  (collectCandidates prog none prog.taggedFin).filterMap (fun c =>
-    if isInvariant prog c then some (c.loopId, c.assignId) else none)
+/-- 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 :=