-- Veil/DSL/Action/Theory.lean

import Veil.DSL.Base

namespace Veil

section Veil

open Classical

section Types

inductive Mode where
  | internal : Mode
  | external : Mode
deriving BEq

variable (m : Mode) (σ ρ : Type)

@[inline] abbrev SProp := σ -> Prop

@[inline] abbrev RProp (ρ : Type u) := ρ → SProp σ

abbrev Wp (m : Mode) (σ ρ : Type) := σ -> RProp σ ρ -> Prop

abbrev Wlp (m : Mode) (σ ρ : Type) := σ -> RProp σ ρ -> Prop

return value `r : ρ` and a post-state `s' : σ` -/

abbrev BigStep := σ -> ρ -> σ -> Prop

end Types

section Theory

variable {σ ρ : Type}

@[actSimp] def Wp.bind (wp : Wp m σ ρ) (wp_cont : ρ -> Wp m σ ρ') : Wp m σ ρ' :=
  fun s post => wp s (fun r s' => wp_cont r s' post)

@[actSimp] def Wp.get : Wp m σ σ := fun s post => post s s

@[actSimp] def Wp.set (s' : σ) : Wp m σ Unit := fun _ post => post () s'

@[actSimp] def Wp.modifyGet (act : σ -> ρ × σ) : Wp m σ ρ := fun s post => let (ret, s') := act s ; post ret s'

class IsSubStateOf (σ : semiOutParam Type) (σ' : Type) where
   
  setIn : σ -> σ' -> σ'
   
  getFrom : σ' -> σ

  setIn_getFrom_idempotent : ∀ σ', setIn (getFrom σ') σ' = σ'
  getFrom_setIn_idempotent : ∀ σ σ', getFrom (setIn σ σ') = σ

@[actSimp] def Wp.lift {σ σ'} [IsSubStateOf σ σ'] (act : Wp m σ ρ) : Wp m σ' ρ :=
  fun s' post => act (getFrom s') (fun r s => post r (setIn s s'))

abbrev Wp.triple {σ ρ} (req : SProp σ) (act : Wp m σ ρ) (ens : RProp σ ρ) : Prop :=
  ∀ s, req s -> act s ens

class LawfulAction {σ ρ : Type} (act : Wp m σ ρ) where
  inter {τ : Type} [Inhabited τ] (post : τ -> RProp σ ρ) :
    ∀ s : σ, (∀ t : τ, act s (post t)) -> act s (∀ t, post t · ·)

  impl (post post' : RProp σ ρ) : ∀ s,
    (∀ r s, post r s -> post' r s) -> act s post -> act s post'

@[actSimp]
abbrev Wp.toWlp {σ ρ : Type} {m : Mode} (wp : Wp m σ ρ) : Wlp m σ ρ :=
  -- `wlp(P, φ, s) = ¬ wp(P, ¬φ, s)`
  fun (s : σ) (post : RProp σ ρ) => ¬ wp s (fun r s' => ¬ post r s')

@[actSimp]
def Wp.toBigStep {σ} (wp : Wp m σ ρ) : BigStep σ ρ :=
  fun s r' s' =>
    wp.toWlp s (fun r₀ s₀ => r' = r₀ ∧ s' = s₀)

@[actSimp]
def Wp.toTwoState {σ} (wp : Wp m σ ρ) : TwoState σ :=
  fun s s' =>
    wp.toWlp s (fun _ s₀ => (s' = s₀))

@[actSimp]
def BigStep.toWp {σ} (act : BigStep σ ρ) : Wp .internal σ ρ :=
  fun s post => ∀ r s', act s r s' -> post r s'

@[actSimp]
def _root_.Function.toWp (m : Mode) (r : TwoState σ) : Wp m σ Unit :=
  fun s post => ∀ s', r s s' -> post () s'

abbrev refines {σ ρ} (act : Wp m σ ρ) (act' : Wp m σ ρ) : Prop :=
  ∀ s post, act s post -> act' s post

instance : LE (Wp m σ ρ) where
  le := refines

abbrev Wp.triple {σ ρ} (req : SProp σ) (act : Wp m σ ρ) (ens : RProp σ ρ) : Prop :=
  ∀ s, req s -> act s ens

abbrev Wp.alwaysSuccessfullyTerminates {σ } (req : SProp σ) (act : Wp m σ ρ)  : Prop :=
  ∀ s, req s -> act s (fun _ _ => True)

abbrev TwoState.triple {σ } (req : SProp σ) (act : TwoState σ) (ens : SProp σ) : Prop :=
  ∀ s s', req s -> act s s' -> ens s'

abbrev BigStep.triple {σ } (req : SProp σ) (act : BigStep σ ρ) (ens : RProp σ ρ) : Prop :=
  ∀ s r' s', req s -> act s r' s' -> ens r' s'
section TwoStateSoundness

-- Repository definitions used:
content: syntax (name:= assumption) "assumptionDef" : attr

-- Repository lemmas used:
content: theorem List.getElem_append_left{α : Type u_1} {l₁ l₂ : List α} {i : Nat} (hn : i < l₁.length) :\n(l₁ ++ l₂)[i] = l₁[i]

-- Local definitions used:
content: inductive Mode where
  | internal : Mode
  | external : Mode
deriving BEq

content: @[inline] abbrev SProp := σ -> Prop

content: @[inline] abbrev RProp (ρ : Type u) := ρ → SProp σ

content: abbrev Wp (m : Mode) (σ ρ : Type) := σ -> RProp σ ρ -> Prop

content: @[actSimp] def Wp.bind (wp : Wp m σ ρ) (wp_cont : ρ -> Wp m σ ρ') : Wp m σ ρ' :=
  fun s post => wp s (fun r s' => wp_cont r s' post)

content: @[actSimp] def Wp.get : Wp m σ σ := fun s post => post s s

content: @[actSimp] def Wp.set (s' : σ) : Wp m σ Unit := fun _ post => post () s'

content: @[actSimp] def Wp.modifyGet (act : σ -> ρ × σ) : Wp m σ ρ := fun s post => let (ret, s') := act s ; post ret s'

content: class IsSubStateOf (σ : semiOutParam Type) (σ' : Type) where
   
  setIn : σ -> σ' -> σ'
   
  getFrom : σ' -> σ

  setIn_getFrom_idempotent : ∀ σ', setIn (getFrom σ') σ' = σ'
  getFrom_setIn_idempotent : ∀ σ σ', getFrom (setIn σ σ') = σ

content: @[actSimp] def Wp.lift {σ σ'} [IsSubStateOf σ σ'] (act : Wp m σ ρ) : Wp m σ' ρ :=
  fun s' post => act (getFrom s') (fun r s => post r (setIn s s'))

content: abbrev Wp.triple {σ ρ} (req : SProp σ) (act : Wp m σ ρ) (ens : RProp σ ρ) : Prop :=
  ∀ s, req s -> act s ens

content: class LawfulAction {σ ρ : Type} (act : Wp m σ ρ) where
  inter {τ : Type} [Inhabited τ] (post : τ -> RProp σ ρ) :
    ∀ s : σ, (∀ t : τ, act s (post t)) -> act s (∀ t, post t · ·)

  impl (post post' : RProp σ ρ) : ∀ s,
    (∀ r s, post r s -> post' r s) -> act s post -> act s post'

-- Local lemmas used:
content: theorem TwoState_sound' [LawfulAction act] (req : SProp σ) (ens : RProp σ ρ) :
  act.triple req ens → act.toTwoState.triple req (∃ r, ens r ·)

content: theorem exists_over_PUnit (p : PUnit → Prop) : (∃ (u : PUnit), p u) = p ()

-- Target theorem and proof:
theorem TwoState_sound'_ret_unit [LawfulAction act] (req : SProp σ) (ens : RProp σ PUnit) :
  act.triple req ens → act.toTwoState.triple req (ens () ·)
:= by
  have heq : (ens () ·) = (∃ r, ens r ·) := by ext ; rw [exists_over_PUnit]
  rw [heq] ; apply TwoState_sound'