import Spa.Language.Base
import Spa.Language.Semantics
import Spa.Language.Graphs
import Mathlib.Data.Finset.Sort
import Mathlib.Data.String.Basic

namespace Spa

/-- A self-contained program to be evaluated, analyzed, and transformed. -/
structure Program where
  /-- The statement at the top level of the program. Since `Spa.Stmt` contains
      sequencing via `Spa.Stmt.andThen`, this can encode any number of
      statements. -/
  rootStmt : Stmt
  /-- A memoized copy of the control-flow graph. This field is an
      implementation detail to avoid re-computing `Spa.GGraph.wrap` and `Spa.Stmt.cfg`
      every time the program's control flow graph is needed -/
  cfgCache : Thunk Graph := Thunk.mk fun _ => Graph.wrap rootStmt.cfg

namespace Program

variable (p : Program)

-- Runtime implementation of `cfg`: read the memoized graph.
private def cfgImpl : Graph := p.cfgCache.get

/-- The control flow graph corresponding to this graph. -/
@[implemented_by cfgImpl]
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

/-- Variables mentioned or defined in this program. -/
def vars : List String := p.rootStmt.vars.sort (· ≤ ·)

/-- `vars` has no duplicates. -/
lemma vars_nodup : p.vars.Nodup := Finset.sort_nodup _ _

/-- All the states in the program's control flow `Spa.Graph`. -/
def states : List p.State := p.cfg.indices

/-- All states in the CFG are contained in `states`. -/
lemma states_complete (s : p.State) : s ∈ p.states := p.cfg.mem_indices s

/-- `states` has no duplicates. -/
lemma states_nodup : p.states.Nodup := p.cfg.nodup_indices

/-- Given a node of the program's CFG, return the code at that node.
    At this time, for convenience of proofs, the CFGs have at most
    one basic statement, and multi-statement basic blocks are encoded
    as chains of blocks. Thus, this returns at most one `Spa.BasicStmt`. -/
def code (st : p.State) : Option BasicStmt := p.cfg.nodes st

/-- Get the predecessors of a particular CFG node / program state. -/
def incoming (s : p.State) : List p.State := p.cfg.predecessors s

/-- The entry point of the program's CFG. -/
def initialState : p.State := Graph.wrapInput p.rootStmt.cfg

/-- The exit point of the program's CFG. -/
def finalState : p.State := Graph.wrapOutput p.rootStmt.cfg

/-- `incoming` is a faithful representation of edges in the CFG. -/
lemma mem_incoming_of_edge {s₁ s₂ : p.State}
    (h : (s₁, s₂) ∈ p.cfg.edges) : s₁ ∈ p.incoming s₂ :=
  p.cfg.mem_predecessors_of_edge h

/-- The `initialState` has no incoming edges (it's the program start). -/
lemma incoming_initialState_eq_nil : p.incoming p.initialState = [] :=
  GGraph.wrap_predecessors_eq_nil p.rootStmt.cfg p.initialState
    (by rw [Graph.wrap_inputs]; exact List.mem_singleton_self _)

end Program

end Spa