import Spa.Analysis.Forward.Evaluation

namespace Spa

variable {L : Type} [Lattice L] {prog : Program} [E : ExprEvaluator L prog]

def updateVariablesFromExpression (k : String) (e : Expr)
    (vs : VariableValues L prog) : VariableValues L prog :=
  FiniteMap.generalizedUpdate id (fun _ vs => E.eval e vs) [k] vs

theorem updateVariablesFromExpression_mono (k : String) (e : Expr) :
    Monotone (updateVariablesFromExpression (L := L) (prog := prog) k e) :=
  FiniteMap.generalizedUpdate_monotone monotone_id (fun _ => E.eval_mono e)

def evalB (_ : prog.State) (bs : BasicStmt)
    (vs : VariableValues L prog) : VariableValues L prog :=
  match bs with
  | .assign k e => updateVariablesFromExpression k e vs
  | .noop => vs

theorem evalB_mono (s : prog.State) (bs : BasicStmt) :
    Monotone (evalB (L := L) (prog := prog) s bs) := by
  cases bs with
  | assign k e => exact updateVariablesFromExpression_mono k e
  | noop => exact monotone_id

instance ExprEvaluator.toStmtEvaluator : StmtEvaluator L prog :=
  ⟨evalB, evalB_mono⟩

instance ExprEvaluator.toStmtEvaluator_valid [LatticeInterpretation L]
    [ValidExprEvaluator L prog] : ValidStmtEvaluator L prog := by
  constructor
  intro s vs ρ₁ ρ₂ bs hbs hvs
  cases hbs with
  | noop => exact hvs
  | assign k e v hev =>
    intro k' l hk'l v' hv'
    cases hv' with
    | here =>
      have hk'l₀ : (k, l) ∈ FiniteMap.generalizedUpdate (ks := prog.vars) id
          (fun _ vs => E.eval e vs) [k] vs := hk'l
      have hl := FiniteMap.generalizedUpdate_mem_eq (f := id)
        (g := fun _ vs => E.eval e vs) (List.mem_singleton_self k) hk'l₀
      rw [hl]
      exact ValidExprEvaluator.valid hev hvs
    | there _ _ _ _ _ hne hmem' =>
      have hk'l₀ : (k', l) ∈ FiniteMap.generalizedUpdate (ks := prog.vars) id
          (fun _ vs => E.eval e vs) [k] vs := hk'l
      have hk'l' : (k', l) ∈ (id vs : VariableValues L prog) :=
        FiniteMap.generalizedUpdate_not_mem_backward
          (fun hmem => hne (List.mem_singleton.mp hmem)) hk'l₀
      exact hvs _ _ hk'l' _ hmem'

end Spa