agda-spa/lean/Spa/Language/Graphs.lean

import Spa.Language.Base
import Mathlib.Data.Fin.Tuple.Basic
import Mathlib.Data.List.ProdSigma
import Mathlib.Data.List.FinRange

/-!

# Algebraic Control Flow Graphs

This file defines control flow graphs and operations to naturally compose them,
making it possible to inductively covnert a program in the object language
(see `Spa.Stmt` in `Spa/Language/Base.lean`) into its corresponding graph.

Graphs are, in general, parameterized by their "payload" (the per-node data); see `GGraph`.
This is useful because other operations, such as finding the CFG node corresponding
to an AST node, are performed by embellishing a graph's basic blocks with their AST
identifiers.

The operations are deliberately a little bit sloppy here, creating empty / statement-less
CFG nodes. Additionally, the current CFG construction algorithm doesn't group
consecutive statements in a single notional basic block into one node.
This makes graph construction much easier to define, and might save us the
trouble of (when trying to find the CFG node for an AST node) doing
indexing into a list.

-/

/-- Logically, when combining `Fin`s from two distinct pools,
    the combination is disjoint. -/
lemma Fin.castAdd_ne_natAdd {n m : ℕ} (i : Fin n) (j : Fin m) :
    Fin.castAdd m i ≠ Fin.natAdd n j := by
  intro h
  have := congrArg Fin.val h
  simp only [Fin.coe_castAdd, Fin.coe_natAdd] at this
  omega

/-- Bump the upper bound of a list of `Fin`s without changing their value. -/
def List.finCastAdd {n : ℕ} (l : List (Fin n)) (m : ℕ) : List (Fin (n + m)) :=
  l.map (Fin.castAdd m)

/-- Bump the upper bound of a list of `Fin`s by adding the amount to their value. -/
def List.finNatAdd {m : ℕ} (l : List (Fin m)) (n : ℕ) : List (Fin (n + m)) :=
  l.map (Fin.natAdd n)

/-- Bump the upper bound of a list of `Fin` pairs without changing their value. -/
def List.finCastAddProd {n : ℕ} (l : List (Fin n × Fin n)) (m : ℕ) :
    List (Fin (n + m) × Fin (n + m)) :=
  l.map (fun e => (e.1.castAdd m, e.2.castAdd m))

/-- Bump the upper bound of a list of `Fin` pairs by adding the amount to their value. -/
def List.finNatAddProd {m : ℕ} (l : List (Fin m × Fin m)) (n : ℕ) :
    List (Fin (n + m) × Fin (n + m)) :=
  l.map (fun e => (e.1.natAdd n, e.2.natAdd n))

namespace Spa

/-- Graph with general (`α`-labeled) nodes. By using a tuple `Fin size → α`
    and writing `edges` over the `Fin size`, guarantees all edges are between real nodes.

    To make graph composition via operations not force a
    [`alga`](https://hackage.haskell.org/package/algebraic-graphs)-style "connect"-based
    algebra, explicitly defines `inputs` and `outputs`, which are the only nodes that
    get connected when graphs are sequenced. This makes the graph construction
    operations more naturally fit with how CFGs are created from `Stmt`s. -/
structure GGraph (α : Type) where
  size : ℕ
  nodes : Fin size → α
  edges : List (Fin size × Fin size)
  inputs : List (Fin size)
  outputs : List (Fin size)

namespace GGraph

variable {α β : Type}

/-- An index (node) in the CFG. -/
abbrev Index (g : GGraph α) : Type := Fin g.size

/-- An edge in the CFG. -/
abbrev Edge (g : GGraph α) : Type := g.Index × g.Index

instance : Functor GGraph where
  map {α β : Type} (f : α → β) (g : GGraph α) : GGraph β :=
    { size := g.size,
      nodes := f ∘ g.nodes
      edges := g.edges,
      inputs := g.inputs,
      outputs := g.outputs }

@[simp] lemma map_size (f : α → β) (g : GGraph α) : (f <$> g).size = g.size := rfl
@[simp] lemma map_edges (f : α → β) (g : GGraph α) : (f <$> g).edges = g.edges := rfl
@[simp] lemma map_inputs (f : α → β) (g : GGraph α) : (f <$> g).inputs = g.inputs := rfl
@[simp] lemma map_outputs (f : α → β) (g : GGraph α) : (f <$> g).outputs = g.outputs := rfl

/-- Overlay two graphs: create a new graph whose nodes and edges come from two
    sub-graphs, without inserting any additional edges. Also combines the
    input and output node sets. -/
def overlay (g₁ g₂ : GGraph α) : GGraph α where
  size := g₁.size + g₂.size
  nodes := Fin.append g₁.nodes g₂.nodes
  edges := g₁.edges.finCastAddProd g₂.size ++ g₂.edges.finNatAddProd g₁.size
  inputs := g₁.inputs.finCastAdd g₂.size ++ g₂.inputs.finNatAdd g₁.size
  outputs := g₁.outputs.finCastAdd g₂.size ++ g₂.outputs.finNatAdd g₁.size

@[inherit_doc] scoped infixr:70 " ∙ " => GGraph.overlay

/-- Sequence two CFGs: create a combined graph whose nodes and edges come
    from two subgraphs, __and__ make all the outputs of the left graph have edges to
    all the inputs of the right graph. By the semantics of CFGs, this
    encodes the fact that code first traverses the basic blocks in theleft
    graph, and does the same for the right graph. -/
def sequence (g₁ g₂ : GGraph α) : GGraph α where
  size := g₁.size + g₂.size
  nodes := Fin.append g₁.nodes g₂.nodes
  edges := g₁.edges.finCastAddProd g₂.size ++ g₂.edges.finNatAddProd g₁.size ++
    (g₁.outputs.finCastAdd g₂.size).product (g₂.inputs.finNatAdd g₁.size)
  inputs := g₁.inputs.finCastAdd g₂.size
  outputs := g₂.outputs.finNatAdd g₁.size

@[inherit_doc] scoped infixr:70 " ⤳ " => GGraph.sequence

/-- When a graph `g` is wrapped in a `loop`, the index / node corresponding
    to the input of the new loop. -/
def loopIn (g : GGraph α) : Fin (2 + g.size) := (0 : Fin 2).castAdd g.size

/-- When a graph `g` is wrapped in a `loop`, the index / node corresponding
    to the output of the new loop. -/
def loopOut (g : GGraph α) : Fin (2 + g.size) := (1 : Fin 2).castAdd g.size

/-- Creates a zero-or-more loop loop in the CFG: connects all the output
    nodes of the CFG back to the graph's beginning, and also introduces a path
    to a new ending node (see `loopOut`) which bypasses the entire graph.

    Notably, both the new input (`loopIn`) and new output (`loopOut`)
    nodes are necessary for correctness: adding a path from inputs to a
    hypothetical no-op end node encodes something like "just the first statement is executed".
    Similarly, just adding a path from a a hypothetical no-op beginning node
    to the outputs encodes "just the last statement is executed".

    This is technically sloppy (see module comment), but it's simple.
-/
def loop (g : GGraph (Option β)) : GGraph (Option β) where
  size := 2 + g.size
  nodes := Fin.append (fun _ : Fin 2 => none) g.nodes
  edges := g.edges.finNatAddProd 2 ++
    ((g.loopIn, ·) <$> g.inputs.finNatAdd 2)  ++
    ((·, g.loopOut) <$> g.outputs.finNatAdd 2) ++
    [(g.loopOut, g.loopIn), (g.loopIn, g.loopOut)]
  inputs := [g.loopIn]
  outputs := [g.loopOut]

@[simp] lemma loop_inputs (g : GGraph (Option β)) : (loop g).inputs = [g.loopIn] := rfl

@[simp] lemma loop_outputs (g : GGraph (Option β)) : (loop g).outputs = [g.loopOut] := rfl

/-- Creates a single-node graph whose node contains the given value. -/
def singleton (a : α) : GGraph α where
  size := 1
  nodes := fun _ => a
  edges := []
  inputs := [0]
  outputs := [0]

/-- Creates a new graph with a single input and single output node. Useful to ensure there's
    a single point of entry and single point of exit. -/
def wrap (g : GGraph (Option β)) : GGraph (Option β) :=
  singleton none ⤳ g ⤳ singleton none

/-- The input / entry node generated by `GGraph.wrap`. -/
def wrapInput (g : GGraph (Option β)) : (wrap g).Index :=
  (0 : Fin 1).castAdd ((g ⤳ singleton none).size)

/-- The output / exit node generated by `GGraph.wrap`. -/
def wrapOutput (g : GGraph (Option β)) : (wrap g).Index :=
  Fin.natAdd 1 ((Fin.natAdd g.size (0 : Fin 1)))

/-- The `wrapInput` is, indeed, the graph's only input after `wrap`. -/
lemma wrap_inputs (g : GGraph (Option β)) :
    (wrap g).inputs = [g.wrapInput] := rfl

/-- The `wrapInput` is, indeed, the graph's only output after `wrap`. -/
lemma wrap_outputs (g : GGraph (Option β)) :
    (wrap g).outputs = [g.wrapOutput] := rfl

@[simp] lemma map_singleton (f : α → β) (a : α) :
    f <$> singleton a = singleton (f a) := rfl

@[simp] lemma map_overlay (f : α → β) (g₁ g₂ : GGraph α) :
    f<$> (g₁ ∙ g₂) = f <$> g₁ ∙ f <$> g₂ := by
  rcases g₁ with ⟨n₁, nd₁, e₁, i₁, o₁⟩; rcases g₂ with ⟨n₂, nd₂, e₂, i₂, o₂⟩
  simp only [Functor.map, GGraph.overlay]
  congr 1
  funext i
  refine Fin.addCases ?_ ?_ i <;> intro j <;> simp [Fin.append_left, Fin.append_right]

@[simp] lemma map_sequence (f : α → β) (g₁ g₂ : GGraph α) :
    f <$> (g₁ ⤳ g₂) = (f <$> g₁) ⤳  (f <$> g₂) := by
  rcases g₁ with ⟨n₁, nd₁, e₁, i₁, o₁⟩; rcases g₂ with ⟨n₂, nd₂, e₂, i₂, o₂⟩
  simp only [Functor.map, GGraph.sequence]
  congr 1
  funext i
  refine Fin.addCases ?_ ?_ i <;> intro j <;> simp [Fin.append_left, Fin.append_right]

@[simp] lemma map_loop (h : β → γ) (g : GGraph (Option β)) :
    (Option.map h) <$> (loop g) = loop (Option.map h <$> g) := by
  rcases g with ⟨n, nd, e, i, o⟩
  simp only [Functor.map, GGraph.loop]
  congr 1
  funext i
  refine Fin.addCases ?_ ?_ i <;> intro j <;> simp [Fin.append_left, Fin.append_right]

@[simp] lemma map_wrap (h : β → γ) (g : GGraph (Option β)) :
    (Option.map h) <$> wrap g = wrap (Option.map h <$> g) := by
  simp [GGraph.wrap, GGraph.map_sequence, GGraph.map_singleton]

/-! ### 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.

`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.

Concretely, every embedding here is a *constant index shift*: `∙` and `⤳` lay
their operands out in consecutive blocks via `Fin.append`, and `loop` prepends its
two synthetic nodes.  Storing the offset rather than an arbitrary function makes
the range of an embedding an interval by construction (`Embed.mem_range_iff`),
which is the "is this node inside that subgraph?" test. -/

/-- Translate an index of `g` into `h` by a constant offset. -/
def shift {g h : GGraph α} (off : ℕ) (hle : g.size + off ≤ h.size) (i : g.Index) :
    h.Index := ⟨off + i.val, by omega⟩

@[simp] lemma shift_val {g h : GGraph α} {off : ℕ} (hle : g.size + off ≤ h.size)
    (i : g.Index) : (shift (h := h) off hle i).val = off + i.val := rfl

/-- An embedding of `g` into `h`: a constant index shift preserving node payloads
    and edges. -/
structure Embed (g h : GGraph α) where
  off : ℕ
  size_le : g.size + off ≤ h.size
  nodes_eq : ∀ i, h.nodes (shift off size_le i) = g.nodes i
  edges_mem : ∀ {e : g.Edge}, e ∈ g.edges →
    (shift off size_le e.1, shift off size_le e.2) ∈ h.edges

/-- The index translation of an embedding. -/
abbrev Embed.f {g h : GGraph α} (e : Embed g h) (i : g.Index) : h.Index :=
  shift e.off e.size_le i

lemma Embed.f_inj {g h : GGraph α} (e : Embed g h) : Function.Injective e.f :=
  fun _ _ hij => Fin.ext (by simpa using congrArg Fin.val hij)

/-- 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; simp only [Embed.f, shift_val]; omega
  · rintro ⟨hlo, hhi⟩
    exact ⟨⟨j.val - e.off, by omega⟩, Fin.ext (by simp only [Embed.f, shift_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 while `Embed` stores only the offset.  The trailing two arguments are
    boilerplate at every call site and default to discharging themselves. -/
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)
    (hle : g.size + off ≤ h.size := by
      first | omega | (simp [GGraph.sequence, GGraph.overlay, GGraph.loop] <;> omega))
    (hk : ∀ i, (k i).val = off + i.val := by intro i; simp) :
    Embed g h where
  off := off
  size_le := hle
  nodes_eq i := by rw [show shift off hle i = k i from Fin.ext (by simp [hk])]; exact hn i
  edges_mem hmem := by
    have hs : ∀ i, shift (h := h) off hle i = k i := fun i => Fin.ext (by simp [hk])
    rw [hs, hs]; exact hem hmem

/-- Embeddings compose; offsets add. -/
def Embed.trans {g₁ g₂ g₃ : GGraph α} (e₁ : Embed g₁ g₂) (e₂ : Embed g₂ g₃) :
    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))
    (hle := by have := e₁.size_le; have := e₂.size_le; omega)
    (hk := fun i => by simp; omega)

/-! The five inclusions `Stmt.cfg` uses, one per composition operator.  Each `_f`
lemma recovers the `Fin.castAdd`/`Fin.natAdd` form for callers stating indices
that way (`Spa/Language/Properties.lean`). -/

/-- The left operand's inclusion into a sequenced graph. -/
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₂) :=
  .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₂) :=
  .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₂) :=
  .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 (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))))

/-! `shift` puts the offset on the left (`off + i.val`), matching `Fin.natAdd`, so
every right inclusion is *definitionally* the form `Spa/Language/Properties.lean`
states its trace indices in.  Only the left inclusions need a bridge, since their
offset is `0` and `0 + i.val` does not reduce.

Deliberately not `@[simp]`: they rewrite `Embed.f` back into `Fin.castAdd` form,
discarding the offset that `Embed.mem_range_iff` — and the subgraph-containment
tests built on it — reason with. -/

lemma Embed.sequenceLeft_f (g₁ g₂ : GGraph α) (i : g₁.Index) :
    (Embed.sequenceLeft g₁ g₂).f i = i.castAdd g₂.size := Fin.ext (Nat.zero_add _)
lemma Embed.overlayLeft_f (g₁ g₂ : GGraph α) (i : g₁.Index) :
    (Embed.overlayLeft g₁ g₂).f i = i.castAdd g₂.size := Fin.ext (Nat.zero_add _)

variable (g : GGraph α)

/-- All the nodes in the graph. -/
def indices : List g.Index := List.finRange g.size

/-- All of the graph's indices are listed in `indices`. -/
lemma mem_indices (idx : g.Index) : idx ∈ g.indices :=
  List.mem_finRange idx

/-- `indices` does not have duplicates. -/
lemma nodup_indices : g.indices.Nodup :=
  List.nodup_finRange g.size

/-- Predecessors of a particular node in the graph. --/
def predecessors (idx : g.Index) : List g.Index :=
  g.indices.filter (fun idx' => (idx', idx) ∈ g.edges)

/-- When sequencing (proven here with `Graph.singleton` on the left), no edges
    exist from the right-hand graph back to the left. -/
private lemma not_mem_edges_castAdd_sequence {g₂ : GGraph (Option β)} (i : Fin 1)
    (idx : (singleton none ⤳ g₂).Index) :
    ((idx, i.castAdd g₂.size) : (singleton none ⤳ g₂).Edge)
      ∉ (singleton none ⤳ g₂).edges := by
  intro h
  rcases List.mem_append.mp h with h' | h'
  · rcases List.mem_append.mp h' with h'' | h''
    · -- lifted edges of `singleton []`: there are none
      simp [singleton, List.finCastAddProd] at h''
    · -- lifted edges of g₂: targets are natAdd
      obtain ⟨e, _, heq⟩ := List.mem_map.mp h''
      exact Fin.castAdd_ne_natAdd i e.2 (congrArg Prod.snd heq).symm
  · -- product edges: targets are natAdd'd inputs of g₂
    obtain ⟨-, hb⟩ := List.mem_product.mp h'
    obtain ⟨j, -, heq⟩ := List.mem_map.mp hb
    exact Fin.castAdd_ne_natAdd i j heq.symm

/-- The input node of a graph after `Graph.wrap` has no predecessors. -/
lemma wrap_predecessors_eq_nil (g : GGraph (Option β)) (idx : (wrap g).Index)
    (h : idx ∈ (wrap g).inputs) :
    (wrap g).predecessors idx = [] := by
  rw [wrap_inputs, List.mem_singleton] at h
  subst h
  rw [GGraph.predecessors, List.filter_eq_nil_iff]
  intro idx' _
  simpa using not_mem_edges_castAdd_sequence (g₂ := g ⤳ singleton none) 0 idx'

/-- There's there's an edge between two nodes `idx₁` and `idx₂`,
    then `idx₁` is the predecessor of `idx₂`. -/
lemma mem_predecessors_of_edge {idx₁ idx₂ : g.Index}
    (h : (idx₁, idx₂) ∈ g.edges) : idx₁ ∈ g.predecessors idx₂ :=
  List.mem_filter.mpr ⟨g.mem_indices idx₁, by simpa using h⟩

/-- A node is a predecessor of another node only if there's an
    edge between them. -/
lemma edge_of_mem_predecessors {idx₁ idx₂ : g.Index}
    (h : idx₁ ∈ g.predecessors idx₂) : (idx₁, idx₂) ∈ g.edges := by
  simpa using (List.mem_filter.mp h).2

end GGraph

/-- "Normal" graphs, for the purposes of the analyses in this
    framework, have basic statements in their nodes, and nothing else. -/
abbrev Graph : Type := GGraph (Option BasicStmt)

namespace Graph

export GGraph (overlay sequence loop singleton wrap loop_inputs loop_outputs wrapInput wrapOutput wrap_inputs wrap_outputs)

@[inherit_doc] scoped infixr:70 " ∙ " => GGraph.overlay
@[inherit_doc] scoped infixr:70 " ⤳ " => GGraph.sequence

end Graph

open Graph in
def Stmt.cfg : Stmt → Graph
  -- A basic statement goes into a single basic block
  | .basic bs => singleton (some bs)
  -- Sequencing of statements corresponds naturally to CFG sequencing
  | .andThen s₁ s₂ => s₁.cfg ⤳  s₂.cfg
  -- An if can execute either one branch or the other; overlap them.
  -- Subsequent sequencing (etc.) will end up creating the forks and joins.
  | .ifElse _ s₁ s₂ => s₁.cfg ∙ s₂.cfg
  -- The `loop` construct was developed specifically for zero-or-more loops like this.
  | .whileLoop _ s => loop s.cfg

end Spa