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

/-!

# Program Traces

This module defines program traces tied to Control Flow Graphs, or CFGs
(see `Spa.GGraph` and `Spa.Graph`). These traces boil town 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.

While the regular `Trace` is just _any_ path through the graph, an
`EndToEndTrace` is a path from the entry node to the exit node, 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 partial trace through a graph `g`, starting right before
    the execution of the basic block at the first index, and
    ending right after the execution of the basic block at the last index. -/
inductive Trace (g : Graph) : g.Index → g.Index → Env → Env → Type
  | single {ρ₁ ρ₂ : Env} {idx : g.Index} :
      EvalBasicStmtOpt ρ₁ (g.nodes idx) ρ₂ → Trace g idx idx ρ₁ ρ₂
  | edge {ρ₁ ρ₂ ρ₃ : Env} {idx₁ idx₂ idx₃ : g.Index} :
      EvalBasicStmtOpt ρ₁ (g.nodes idx₁) ρ₂ → (idx₁, idx₂) ∈ g.edges →
      Trace g idx₂ idx₃ ρ₂ ρ₃ → Trace g idx₁ idx₃ ρ₁ ρ₃

/-!

## Open Traces

A normal `Trace` starts right before one state, and ends right after another.
This is convenient for inductively proving correctness / sufficience, but
awkward because 1) no empty traces exist and 2) concatenation requires an extra
edge.

However, when attempting an "empty" trace, two types are equally possible:
traces that end _right before_ executing a state (`Traceₗ`) and
traces that begin _right after_ executing a state (`Traceᵣ`). They
are symmetric and can be concatenated with full traces on the left
and right, respectively. -/

/-- Left-open trace, representing execution that ends right before `idx₂`. -/
inductive Traceₗ (g : Graph) : g.Index → g.Index → Env → Env → Type where
  | nil {idx : g.Index} {ρ : Env} : Traceₗ g idx idx ρ ρ
  | cons {idx₁ idx₂ idx₃ : g.Index} {ρ₁ ρ₂ ρ₃ : Env} :
      EvalBasicStmtOpt ρ₁ (g.nodes idx₁) ρ₂ →
      (idx₁, idx₂) ∈ g.edges →
      Traceₗ g idx₂ idx₃ ρ₂ ρ₃ → Traceₗ g idx₁ idx₃ ρ₁ ρ₃

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

/-- Right-open trace, representing execution that starts right after `idx₁`. -/
inductive Traceᵣ (g : Graph) : g.Index → g.Index → Env → Env → Type where
  | nil {idx : g.Index} {ρ : Env} : Traceᵣ g idx idx ρ ρ
  | cons {idx₁ idx₂ idx₃ : g.Index} {ρ₁ ρ₂ ρ₃ : Env} :
      Traceᵣ g idx₁ idx₂ ρ₁ ρ₂ →
      (idx₂, idx₃) ∈ g.edges →
      EvalBasicStmtOpt ρ₂ (g.nodes idx₃) ρ₃ → Traceᵣ g idx₁ idx₃ ρ₁ ρ₃

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

/-- Sequence two traces together. Since the endpoint of the first trace
    is _after_ its last basic block's execution, and the beginning of
    the next trace is _before_ its first basic block's execution,
    there must be an edge to connect the two. -/
def Trace.concat {g : Graph} {idx₁ idx₂ idx₃ idx₄ : g.Index}
    {ρ₁ ρ₂ ρ₃ : Env} (tr₁ : Trace g idx₁ idx₂ ρ₁ ρ₂)
    (he : (idx₂, idx₃) ∈ g.edges) (tr₂ : Trace g idx₃ idx₄ ρ₂ ρ₃) :
    Trace g idx₁ idx₄ ρ₁ ρ₃ :=
  match tr₁ with
  | single hbs => edge hbs he tr₂
  | edge hbs he' tr₁' => edge hbs he' (tr₁'.concat he tr₂)

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

def Trace.addEdge {g : Graph} {idx₁ idx₂ idx₃ : g.Index} {ρ₁ ρ₂ : Env} :
    Trace g idx₁ idx₂ ρ₁ ρ₂ →
    (idx₂, idx₃) ∈ g.edges →
    Traceₗ g idx₁ idx₃ ρ₁ ρ₂
  | .single hnode, hedge => .cons hnode hedge .nil
  | .edge hnode hedge' rest, hedge => .cons hnode hedge' (rest.addEdge hedge)

@[aesop simp]
def Traceₗ.append {g : Graph} {idx₁ idx₂ idx₃ : g.Index} {ρ₁ ρ₂ ρ₃ : Env} :
    Traceₗ g idx₁ idx₂ ρ₁ ρ₂ → Traceₗ g idx₂ idx₃ ρ₂ ρ₃ →
    Traceₗ g idx₁ idx₃ ρ₁ ρ₃
  | .nil, rhs => rhs
  | .cons hnode hedge rest, rhs => .cons hnode hedge (rest.append rhs)

@[simp] def traceₗ_append_nil {g : Graph} {idx₁ idx₂ : g.Index} {ρ₁ ρ₂ : Env}
    {trₗ : Traceₗ g idx₁ idx₂ ρ₁ ρ₂} : trₗ.append Traceₗ.nil = trₗ := by
  induction trₗ <;> aesop

def Traceₗ.appendTrace {g : Graph} {idx₁ idx₂ idx₃ : g.Index} {ρ₁ ρ₂ ρ₃ : Env} :
    Traceₗ g idx₁ idx₂ ρ₁ ρ₂ → Trace g idx₂ idx₃ ρ₂ ρ₃ →
    Trace g idx₁ idx₃ ρ₁ ρ₃
  | .nil, rhs => rhs
  | .cons hnode hedge rest, rhs => .edge hnode hedge (rest.appendTrace rhs)

def Traceₗ.appendStep {g : Graph} {idx₁ idx₂ : g.Index} {ρ₁ ρ₂ ρ₃ : Env} :
    Traceₗ g idx₁ idx₂ ρ₁ ρ₂ → EvalBasicStmtOpt ρ₂ (g.nodes idx₂) ρ₃ →
    Trace g idx₁ idx₂ ρ₁ ρ₃ := fun trₗ hbs => trₗ.appendTrace (Trace.single hbs)

def Trace.appendRight {g : Graph} {idx₁ idx₂ idx₃ : g.Index} {ρ₁ ρ₂ ρ₃ : Env} :
    Trace g idx₁ idx₂ ρ₁ ρ₂ → Traceᵣ g idx₂ idx₃ ρ₂ ρ₃ →
    Trace g idx₁ idx₃ ρ₁ ρ₃
  | lhs, .nil => lhs
  | lhs, .cons rest hedge hnode => Trace.concat (lhs.appendRight rest) hedge (.single hnode)

instance instHAppendTraceLTraceL {g : Graph} {idx₁ idx₂ idx₃ : g.Index} {ρ₁ ρ₂ ρ₃ : Env} :
    HAppend (Traceₗ g idx₁ idx₂ ρ₁ ρ₂) (Traceₗ g idx₂ idx₃ ρ₂ ρ₃) (Traceₗ g idx₁ idx₃ ρ₁ ρ₃) where
  hAppend := Traceₗ.append

instance instHAppendTraceLTrace {g : Graph} {idx₁ idx₂ idx₃ : g.Index} {ρ₁ ρ₂ ρ₃ : Env} :
    HAppend (Traceₗ g idx₁ idx₂ ρ₁ ρ₂) (Trace g idx₂ idx₃ ρ₂ ρ₃) (Trace g idx₁ idx₃ ρ₁ ρ₃) where
  hAppend := Traceₗ.appendTrace

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

instance instHAppendTraceTraceR {g : Graph} {idx₁ idx₂ idx₃ : g.Index} {ρ₁ ρ₂ ρ₃ : Env} :
    HAppend (Trace g idx₁ idx₂ ρ₁ ρ₂) (Traceᵣ g idx₂ idx₃ ρ₂ ρ₃) (Trace g idx₁ idx₃ ρ₁ ρ₃) where
  hAppend := Trace.appendRight

@[simp] lemma Traceₗ.append_addEdge {g : Graph}
    {idx₁ idx₂ idx₃ idx₄ : g.Index} {ρ₁ ρ₂ ρ₃ ρ₄ : Env}
    (trₗ : Traceₗ g idx₁ idx₂ ρ₁ ρ₂)
    (hnode : EvalBasicStmtOpt ρ₂ (g.nodes idx₂) ρ₃)
    (hedge : (idx₂, idx₃) ∈ g.edges)
    (rest : Traceₗ g idx₃ idx₄ ρ₃ ρ₄) :
    trₗ.append (Traceₗ.cons hnode hedge rest) =
      (Trace.addEdge (trₗ.appendStep hnode) hedge).append rest := by
  induction trₗ <;> simp [Traceₗ.append, Traceₗ.appendStep, Traceₗ.appendTrace, Trace.addEdge, *]

@[simp] lemma Traceₗ.appendTrace_addEdge {g : Graph}
    {idx₁ idx₂ idx₃ idx₄ : g.Index} {ρ₁ ρ₂ ρ₃ ρ₄ : Env}
    (trₗ : Traceₗ g idx₁ idx₂ ρ₁ ρ₂)
    (hnode : EvalBasicStmtOpt ρ₂ (g.nodes idx₂) ρ₃)
    (hedge : (idx₂, idx₃) ∈ g.edges)
    (rest : Trace g idx₃ idx₄ ρ₃ ρ₄) :
    trₗ.appendTrace (Trace.edge hnode hedge rest) =
      (Trace.addEdge (trₗ.appendStep hnode) hedge).appendTrace rest := by
  induction trₗ <;> simp [Traceₗ.appendTrace, Traceₗ.appendStep, Trace.addEdge, *]

/-- A beginning-to-end trace corresponding to the CFG `g`. -/
inductive EndToEndTrace (g : Graph) (ρ₁ ρ₂ : Env) : Type
  | intro (idx₁ : g.Index) (idx₁_mem : idx₁ ∈ g.inputs)
      (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

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

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

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


end Spa