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

namespace Spa

structure Program where
  rootStmt : Stmt
  /-- A memoized copy of the control-flow graph. This field is an
      implementation detail to avoid re-computing `Graph.wrap` and `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

@[implemented_by cfgImpl]
def cfg : Graph := Graph.wrap p.rootStmt.cfg

abbrev State : Type := p.cfg.Index

def vars : List String := p.rootStmt.vars.sort (· ≤ ·)

lemma vars_nodup : p.vars.Nodup := Finset.sort_nodup _ _

def states : List p.State := p.cfg.indices

lemma states_complete (s : p.State) : s ∈ p.states := p.cfg.mem_indices s

lemma states_nodup : p.states.Nodup := p.cfg.nodup_indices

def code (st : p.State) : Option BasicStmt := p.cfg.nodes st

def incoming (s : p.State) : List p.State := p.cfg.predecessors s

def initialState : p.State := Graph.wrapInput p.rootStmt.cfg

def finalState : p.State := Graph.wrapOutput p.rootStmt.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

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