import Spa.Language.Graphs
import Spa.Language.Program
import Spa.Language.Semantics

/-!

# Program Traces

This module defines program traces tied to Control Flow Graphs, or CFGs
(see `Spa.GGraph` and `Spa.Graph`). These traces boil down to sequences of
basic-block executions (really, `Spa.BasicStmt` executions), each of which must
have an actual basic block in the graph _and_ be connected to the previous
basic block by an edge. In this way, traces encode executions admitted
by the CFG.

`Path` interleaves execution and edge steps, with endpoints recording whether
we are before or after a node. `Trace`, `Traceₗ`, and `Traceᵣ` are endpoint
specializations of this one type. An `EndToEndTrace` runs from a graph input
to a graph output, denoting full program execution.

Properties about graphs and language semantics (especially,
the fact that the graph contains the proper basic block and edges
to represent any program execution according to the
language's big-step semantics `EvalStmt`) is found
in `Spa/Language/Properties.lean`.

-/

namespace Spa

/-- A node together with the phase of its execution. -/
inductive Position (α : Type) where
  | before : α → Position α
  | after : α → Position α
  deriving DecidableEq

abbrev Configuration (g : Graph) := Position g.Index × Env

/-- Executing a node changes the environment; following an edge preserves it. -/
inductive Step (g : Graph) : Configuration g → Configuration g → Type where
  | execute {i : g.Index} {ρ ρ' : Env}
      (h : EvalBasicStmtOpt ρ (g.nodes i) ρ') :
      Step g (.before i, ρ) (.after i, ρ')
  | edge {i j : g.Index} {ρ : Env} (h : (i, j) ∈ g.edges) :
      Step g (.after i, ρ) (.before j, ρ)

/-- A concrete CFG path, including executions of statement-less nodes. -/
inductive Path (g : Graph) : Configuration g → Configuration g → Type where
  | nil {a} : Path g a a
  | cons {a b c} : Step g a b → Path g b c → Path g a c

namespace Path

variable {g : Graph} {a b c d : Configuration g}

@[match_pattern] def single (s : Step g a b) : Path g a b := .cons s .nil

def append {a b c : Configuration g} : Path g a b → Path g b c → Path g a c
  | .nil, q => q
  | .cons s p, q => .cons s (p.append q)

instance : HAppend (Path g a b) (Path g b c) (Path g a c) := ⟨append⟩

@[simp] lemma nil_append (p : Path g a b) : Path.nil.append p = p := rfl

@[simp] lemma append_nil (p : Path g a b) : p.append Path.nil = p := by
  induction p <;> aesop (add simp append)

lemma append_assoc (p : Path g a b) (q : Path g b c) (r : Path g c d) :
    (p.append q).append r = p.append (q.append r) := by
  induction p <;> aesop (add simp append)

end Path

def GGraph.Embed.mapConfiguration {g h : Graph} (e : GGraph.Embed g h) :
    Configuration g → Configuration h
  | (.before i, ρ) => (.before (e.f i), ρ)
  | (.after i, ρ) => (.after (e.f i), ρ)

lemma GGraph.Embed.mapConfiguration_trans {g h k : Graph}
    (e : GGraph.Embed g h) (f : GGraph.Embed h k) (a : Configuration g) :
    f.mapConfiguration (e.mapConfiguration a) = (e.trans f).mapConfiguration a := by
  rcases a with ⟨_ | _, ρ⟩ <;> rfl

noncomputable def Step.embed {g h : Graph} (e : GGraph.Embed g h)
    {a b : Configuration g} : Step g a b → Step h (e.mapConfiguration a) (e.mapConfiguration b)
  | .execute h => .execute (_root_.cast (congrArg (EvalBasicStmtOpt _ · _) (e.nodes_eq _).symm) h)
  | .edge h => .edge (e.edges_mem h)

noncomputable def Path.embed {g h : Graph} (e : GGraph.Embed g h)
    {a b : Configuration g} : Path g a b → Path h (e.mapConfiguration a) (e.mapConfiguration b)
  | .nil => .nil
  | .cons s p => .cons (s.embed e) (p.embed e)

lemma Path.embed_append {g h : Graph} (e : GGraph.Embed g h)
    {a b c : Configuration g} (p : Path g a b) (q : Path g b c) :
    (p.append q).embed e = (p.embed e).append (q.embed e) := by
  induction p <;> aesop (add simp [append, embed])

/-- Transport endpoints without changing the path. -/
def Path.cast {g : Graph} {a b a' b' : Configuration g}
    (ha : a = a') (hb : b = b') (p : Path g a b) : Path g a' b' := ha ▸ hb ▸ p

lemma Path.embed_trans {g h k : Graph} (e : GGraph.Embed g h) (f : GGraph.Embed h k)
    {a b : Configuration g} (p : Path g a b) :
    ((p.embed e).embed f).cast (e.mapConfiguration_trans f a)
      (e.mapConfiguration_trans f b) = p.embed (e.trans f) := by
  induction p with
  | @nil a => rcases a with ⟨_ | _, ρ⟩ <;> rfl
  | @cons a b c s p ih =>
    rcases c with ⟨_ | _, ρ⟩ <;> cases s <;>
      aesop (add simp [embed, Step.embed, cast, GGraph.Embed.mapConfiguration, cast_cast])

/-- A trace includes the executions of both endpoint nodes. -/
abbrev Trace (g : Graph) (i j : g.Index) (ρ ρ' : Env) :=
  Path g (.before i, ρ) (.after j, ρ')

/-- A prefix ending before execution of its final node. -/
abbrev Traceₗ (g : Graph) (i j : g.Index) (ρ ρ' : Env) :=
  Path g (.before i, ρ) (.before j, ρ')

/-- A suffix starting after execution of its initial node. -/
abbrev Traceᵣ (g : Graph) (i j : g.Index) (ρ ρ' : Env) :=
  Path g (.after i, ρ) (.after j, ρ')

/-- Compatibility patterns for an execution and an execution-edge pair. -/
@[match_pattern] abbrev Trace.single {g : Graph} {ρ₁ ρ₂ : Env} {idx : g.Index}
    (h : EvalBasicStmtOpt ρ₁ (g.nodes idx) ρ₂) : Trace g idx idx ρ₁ ρ₂ :=
  .cons (.execute h) .nil

@[match_pattern] abbrev Trace.edge {g : Graph} {ρ₁ ρ₂ ρ₃ : Env}
    {idx₁ idx₂ idx₃ : g.Index} (h : EvalBasicStmtOpt ρ₁ (g.nodes idx₁) ρ₂)
    (he : (idx₁, idx₂) ∈ g.edges) (p : Trace g idx₂ idx₃ ρ₂ ρ₃) :
    Trace g idx₁ idx₃ ρ₁ ρ₃ := Path.cons (.execute h) (.cons (.edge he) p)

@[match_pattern] abbrev Traceₗ.nil {g : Graph} {idx : g.Index} {ρ : Env} :
    Traceₗ g idx idx ρ ρ := Path.nil

@[match_pattern] abbrev Traceₗ.cons {g : Graph} {ρ₁ ρ₂ ρ₃ : Env}
    {idx₁ idx₂ idx₃ : g.Index} (h : EvalBasicStmtOpt ρ₁ (g.nodes idx₁) ρ₂)
    (he : (idx₁, idx₂) ∈ g.edges) (p : Traceₗ g idx₂ idx₃ ρ₂ ρ₃) :
    Traceₗ g idx₁ idx₃ ρ₁ ρ₃ := Path.cons (.execute h) (.cons (.edge he) p)

@[match_pattern] abbrev Traceᵣ.nil {g : Graph} {idx : g.Index} {ρ : Env} : Traceᵣ g idx idx ρ ρ := Path.nil

abbrev Traceᵣ.cons {g : Graph} {ρ₁ ρ₂ ρ₃ : Env} {idx₁ idx₂ idx₃ : g.Index}
    (p : Traceᵣ g idx₁ idx₂ ρ₁ ρ₂) (he : (idx₂, idx₃) ∈ g.edges)
    (h : EvalBasicStmtOpt ρ₂ (g.nodes idx₃) ρ₃) : Traceᵣ g idx₁ idx₃ ρ₁ ρ₃ :=
  p.append (.cons (.edge he) (.single (.execute h)))

abbrev Traceₗ.single (g : Graph) (idx : g.Index) (ρ : Env) : Traceₗ g idx idx ρ ρ := .nil
abbrev Traceᵣ.single (g : Graph) (idx : g.Index) (ρ : Env) : Traceᵣ g idx idx ρ ρ := .nil

abbrev Trace.concat {g : Graph} {idx₁ idx₂ idx₃ idx₄ : g.Index} {ρ₁ ρ₂ ρ₃ : Env}
    (p : Trace g idx₁ idx₂ ρ₁ ρ₂) (he : (idx₂, idx₃) ∈ g.edges)
    (q : Trace g idx₃ idx₄ ρ₂ ρ₃) : Trace g idx₁ idx₄ ρ₁ ρ₃ :=
  (p.append (.single (.edge he))).append q

scoped notation:65 tr₁:66 " ++< " he " >++ " tr₂:65 => Trace.concat tr₁ he tr₂

abbrev Trace.addEdge {g : Graph} {idx₁ idx₂ idx₃ : g.Index} {ρ₁ ρ₂ : Env}
    (p : Trace g idx₁ idx₂ ρ₁ ρ₂) (he : (idx₂, idx₃) ∈ g.edges) :
    Traceₗ g idx₁ idx₃ ρ₁ ρ₂ := p.append (.single (.edge he))

abbrev Traceₗ.append {g : Graph} {i j k : g.Index} {ρ₁ ρ₂ ρ₃ : Env}
    (p : Traceₗ g i j ρ₁ ρ₂) (q : Traceₗ g j k ρ₂ ρ₃) : Traceₗ g i k ρ₁ ρ₃ :=
  Path.append p q

abbrev Traceₗ.appendTrace {g : Graph} {i j k : g.Index} {ρ₁ ρ₂ ρ₃ : Env}
    (p : Traceₗ g i j ρ₁ ρ₂) (q : Trace g j k ρ₂ ρ₃) : Trace g i k ρ₁ ρ₃ :=
  Path.append p q

abbrev Trace.appendRight {g : Graph} {i j k : g.Index} {ρ₁ ρ₂ ρ₃ : Env}
    (p : Trace g i j ρ₁ ρ₂) (q : Traceᵣ g j k ρ₂ ρ₃) : Trace g i k ρ₁ ρ₃ :=
  Path.append p q

noncomputable abbrev Trace.embed {g h : Graph} (e : GGraph.Embed g h)
    {i j : g.Index} {ρ₁ ρ₂ : Env} (p : Trace g i j ρ₁ ρ₂) :
    Trace h (e.f i) (e.f j) ρ₁ ρ₂ := Path.embed e p

abbrev Traceₗ.appendStep {g : Graph} {idx₁ idx₂ : g.Index} {ρ₁ ρ₂ ρ₃ : Env}
    (p : Traceₗ g idx₁ idx₂ ρ₁ ρ₂) (h : EvalBasicStmtOpt ρ₂ (g.nodes idx₂) ρ₃) :
    Trace g idx₁ idx₂ ρ₁ ρ₃ := Path.append p (.single (.execute h))

instance {g : Graph} {idx₁ idx₂ : g.Index} {ρ₁ ρ₂ ρ₃ : Env} :
    HAppend (Traceₗ g idx₁ idx₂ ρ₁ ρ₂) (EvalBasicStmtOpt ρ₂ (g.nodes idx₂) ρ₃)
      (Trace g idx₁ idx₂ ρ₁ ρ₃) := ⟨Traceₗ.appendStep⟩

/-- The node executed by an optional-statement step; empty nodes are omitted. -/
def EvalBasicStmtOpt.steps {α : Type*} (idx : α) {ρ₁ ρ₂ : Env} {obs : Option BasicStmt} :
    EvalBasicStmtOpt ρ₁ obs ρ₂ → List α
  | .none => []
  | .some _ => [idx]

def Step.steps {g : Graph} {a b : Configuration g} : Step g a b → List g.Index
  | .execute (i := i) h => h.steps i
  | .edge _ => []

/-- Executed nodes in chronological order; edges and empty nodes contribute nothing.
The instruction at each node is given by `g.nodes`, rather than copied into the history. -/
def Path.steps {g : Graph} {a b : Configuration g} : Path g a b → List g.Index
  | .nil => []
  | .cons s p => s.steps ++ p.steps

abbrev Trace.steps {g : Graph} {i j : g.Index} {ρ₁ ρ₂ : Env}
    (p : Trace g i j ρ₁ ρ₂) : List g.Index := Path.steps p
abbrev Traceₗ.steps {g : Graph} {i j : g.Index} {ρ₁ ρ₂ : Env}
    (p : Traceₗ g i j ρ₁ ρ₂) : List g.Index := Path.steps p
abbrev Traceᵣ.steps {g : Graph} {i j : g.Index} {ρ₁ ρ₂ : Env}
    (p : Traceᵣ g i j ρ₁ ρ₂) : List g.Index := Path.steps p

@[simp] lemma Path.steps_append {g : Graph} {a b c : Configuration g}
    (p : Path g a b) (q : Path g b c) :
    (p.append q).steps = p.steps ++ q.steps := by
  induction p <;> aesop (add simp [append, steps, List.append_assoc])

@[simp] lemma Traceₗ.steps_appendStep {g : Graph} {idx₁ idx₂ : g.Index}
    {ρ₁ ρ₂ ρ₃ : Env} (tr : Traceₗ g idx₁ idx₂ ρ₁ ρ₂)
    (hbs : EvalBasicStmtOpt ρ₂ (g.nodes idx₂) ρ₃) :
    (tr ++ hbs).steps = tr.steps ++ hbs.steps idx₂ := by
  change Path.steps (Path.append tr (Path.single (.execute hbs))) = _
  aesop (add simp [Trace.steps, Traceₗ.steps, Path.single, Path.steps, Step.steps])

@[simp] lemma Trace.steps_addEdge {g : Graph} {idx₁ idx₂ idx₃ : g.Index}
    {ρ₁ ρ₂ : Env} (tr : Trace g idx₁ idx₂ ρ₁ ρ₂) (he : (idx₂, idx₃) ∈ g.edges) :
    (tr.addEdge he).steps = tr.steps := by
  change Path.steps (Path.append tr (Path.single (.edge he))) = _
  aesop (add simp [Trace.steps, Traceₗ.steps, Path.single, Path.steps, Step.steps])

/-- A beginning-to-end trace corresponding to the CFG `g`. -/
structure EndToEndTrace (g : Graph) (ρ₁ ρ₂ : Env) : Type where
  intro ::
  entry : g.Index
  entry_mem : entry ∈ g.inputs
  exit : g.Index
  exit_mem : exit ∈ g.outputs
  trace : Trace g entry exit ρ₁ ρ₂

/-- Every trace splits into the prefix arriving at its last node and that node's execution. -/
def Trace.split {g : Graph} {i₁ i₂ : g.Index} {ρ₁ ρ₂ : Env} :
    Trace g i₁ i₂ ρ₁ ρ₂ → Σ ρ, Traceₗ g i₁ i₂ ρ₁ ρ × EvalBasicStmtOpt ρ (g.nodes i₂) ρ₂
  | Trace.single h => ⟨_, .nil, h⟩
  | Trace.edge h he rest =>
      let ⟨ρ, pre, step⟩ := rest.split
      ⟨ρ, Traceₗ.cons h he pre, step⟩

@[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
  match tr with
  | Trace.single h => rw [Trace.split.eq_1]; rfl
  | Trace.edge h he rest =>
      have ih := Trace.split_append rest
      rw [Trace.split.eq_2]
      aesop (add simp [HAppend.hAppend, Traceₗ.appendStep, Path.append])

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

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