passert3.lean

import data.rbmap

namespace experiment

def prog (σ : Type) (α : Type) :=
state_t σ option α

instance (σ : Type) : monad (prog σ) :=
by delta prog; apply_instance

variables {σ α β : Type}

def passert (p : prog σ α) (s₀ : σ) (post : σ → α → Prop) : Prop :=
match p s₀ with
| (some (a, s)) := post s a
| none          := true
end

lemma passert_of_post (p : prog σ α) (s₀ : σ) {post : σ → α → Prop} :
      (∀ s a, post s a) → passert p s₀ post :=
begin
  intro h,
  simp [passert],
  generalize h₁ : p s₀ = o,
  cases o with r,
  { simp [passert] },
  { cases r with a s, simp [passert], apply h}
end

/-
- post₂ is the post condition we want to prove for p.
- post₁ is a post condition we have already proved. Example: it is the spec for p.
-/
lemma passert_mp {p : prog σ α} {s₀ : σ} {post₁ post₂ : σ → α → Prop} :
     passert p s₀ post₁ → (∀ s a, post₁ s a → post₂ s a) → passert p s₀ post₂ :=
begin
  simp [passert],
  generalize h₁ : p s₀ = o,
  cases o with r,
  { simp [passert] },
  { cases r with a s,
    simp [passert], intros h₁ h₂, exact h₂ _ _ h₁ }
end

/- Similar to passert_mp, but allows us to apply/accumulate several specs for p.
   After we apply all specs we want, we can eliminate passert using passert_of_post -/
lemma passert_mp' {p : prog σ α} {s₀ : σ} {post₁ post₂ : σ → α → Prop} :
      passert p s₀ post₁ → passert p s₀ (λ s a, post₁ s a → post₂ s a) → passert p s₀ post₂ :=
begin
  simp [passert],
  generalize h₁ : p s₀ = o,
  cases o with r,
  { simp [passert] },
  { cases r with a s, simp [passert], intros h₁ h₂, exact h₂ h₁ }
end

@[simp]
lemma return_spec_eq (post : σ → α → Prop) (s : σ) (a : α) : passert (return a) s post = post s a :=
rfl

@[simp]
lemma bind_spec_eq (p₁ : prog σ α) (p₂ : α → prog σ β) (post : σ → β → Prop) (s₀ : σ) :
                   passert (p₁ >>= p₂) s₀ post = passert p₁ s₀ (λ s₁ a, passert (p₂ a) s₁ post) :=
begin
  dsimp [passert, experiment.monad, state_t_bind, option.bind, monad.to_has_bind],
  generalize h₁ : p₁ s₀ = o₁,
  cases o₁ with r₁,
  { simp [passert, option.bind] },
  { cases r₁ with a s₁, simp [passert],
    generalize h₂ : p₂ a s₁ = o₂,
    cases o₂ with r₂,
    { simp [passert, option.bind, state_t_bind, *] },
    { cases r₂ with b s₂, simp [passert, *, option.bind, state_t_bind] }
   }
end

def fail : prog σ α :=
λ s, none

@[simp]
lemma fail_spec (s₀ : σ) (post : σ → α → Prop) : passert fail s₀ post :=
trivial

def write (s : σ) : prog σ unit :=
state_t.write s

def read : prog σ σ :=
state_t.read

def when (c : prog σ bool) (p : prog σ unit) : prog σ unit :=
c >>= (λ b, if b = tt then p else return ())

def assert (c : σ → bool) : prog σ unit :=
read >>= (λ s, if c s = tt then return () else fail)

@[simp]
lemma write_spec_eq (s₀ s : σ) (post : σ → unit → Prop) : passert (write s) s₀ post = post s () :=
rfl

@[simp]
lemma read_spec_eq (s₀ : σ) (post : σ → σ → Prop) : passert read s₀ post = post s₀ s₀ :=
rfl

@[simp]
lemma when_spec_eq (c : prog σ bool) (p : prog σ unit) (post : σ → unit → Prop) (s₀ : σ) :
      passert (when c p) s₀ post = passert c s₀ (λ s₁ b, (b = tt → passert p s₁ post) ∧ (b = ff → post s₁ ())) :=
begin
  simp [when],
  congr, apply funext, intro s₁, apply funext, intro b,
  cases b; simp
end

@[simp]
lemma assert_spec_eq (c : σ → bool) (s₀ : σ) (post : σ → unit → Prop) : passert (assert c) s₀ post = (c s₀ = tt → post s₀ ()) :=
begin
  simp [assert],
  cases c s₀; simp,
end

def is_zero : prog nat bool :=
read >>= (λ x, return (to_bool (x = 0)))

def inc : prog nat unit :=
read >>= (λ x, write (x + 1))

lemma is_zero_spec (s₀ : nat) : passert is_zero s₀ (λ s b, s = 0 ↔ b = tt) :=
by simp [is_zero]

/- For simple procedures such as isZero, we can create very precise summaries.
   We can always do it when in the spec s is a function of s₀. -/
lemma is_zero_spec_eq (s₀ : nat) (post : nat → bool → Prop) : passert is_zero s₀ post = post s₀ (to_bool (s₀ = 0)) :=
by simp [is_zero]

lemma inc_spec (s₀ : nat) : passert inc s₀ (λ s _, s = s₀ + 1) :=
by simp [inc]

lemma inc_spec_eq (s₀ : nat) (post : nat → unit → Prop) : passert inc s₀ post = post (s₀ + 1) () :=
by simp [inc]

def prog1 : prog nat unit :=
when is_zero inc

/- Proce prog1 simple spec using *_spec_eq lemmas -/
lemma prog1_spec1 (s₀ : nat) : passert prog1 s₀ (λ s _, s > 0) :=
begin
  cases s₀; simp [prog1, nat.succ_ne_zero, inc_spec_eq, is_zero_spec_eq],
  /- Goals are
    ⊢ 1 > 0
    a : ℕ
    ⊢ nat.succ a > 0
  -/
  tactic.comp_val,
  apply nat.zero_lt_succ
end

/- Prove spec again using unfolding -/
lemma prog1_spec2 (s₀ : nat) : passert prog1 s₀ (λ s _, s > 0) :=
begin
  simp [prog1, is_zero, inc],
  cases s₀; simp [nat.succ_ne_zero],
  tactic.comp_val,
  apply nat.zero_lt_succ
end

/- Prove spec again using is_zero and inc specs, and passert_mp.
   Remark: this process looks verbose, but I think it can be automated
   by "rewriting" with "->" lemmas and having support for applying specs
   for exisiting functions + passert_mp.
-/
lemma prog1_spec3 (s₀ : nat) : passert prog1 s₀ (λ s _, s > 0) :=
begin
  simp [prog1],
  apply passert_mp (is_zero_spec s₀), simp, clear s₀,
  intros s₁ b hz, -- hz is the isZero spec,
  cases s₁,
  { simp [*] at *,
    apply passert_mp (inc_spec 0), simp, intros, simp [*], tactic.comp_val },
  { simp [*, nat.succ_ne_zero] at *,
    apply nat.zero_lt_succ }
end

/- Prove spec with precondition -/
lemma prog1_spec4 (s₀ : nat) : s₀ > 0 → passert prog1 s₀ (λ s _, s = s₀) :=
begin
  cases s₀; intro h; simp [prog1, is_zero, inc, nat.succ_ne_zero],
  exact absurd h (lt_irrefl _)
end

def mem := rbmap nat nat

def read_mem (m : mem) (id : nat) : nat :=
match m.find id with
| some v := v
| none   := default nat
end

def write_mem (m : mem) (id : nat) (v : nat) : mem :=
m.insert id v

@[simp]
lemma read_write_mem (m : mem) (id : nat) (v : nat) : read_mem (write_mem m id v) id = v :=
by simp [read_mem, write_mem, rbmap.find_insert]

@[simp]
lemma read_write_mem_of_ne (m : mem) (id₁ id₂ : nat) (v : nat) : ¬ id₁ = id₂ → read_mem (write_mem m id₁ v) id₂ = read_mem m id₂ :=
by intros; simp [read_mem, write_mem, rbmap.find_insert_of_ne, *]

@[simp]
lemma read_write_mem_of_ne' (m : mem) (id₁ id₂ : nat) (v : nat) : ¬ id₂ = id₁ → read_mem (write_mem m id₁ v) id₂ = read_mem m id₂ :=
by intros h; simp [read_mem, write_mem, rbmap.find_insert_of_ne, ne.symm h]

def read_var (id : nat) : prog mem nat :=
do m ← read, return $ read_mem m id

def write_var (id : nat) (val : nat) : prog mem unit :=
do m ← read, write (write_mem m id val)

infix `::=`:50 := write_var

/- TODO: should we create a quotient for maps and use `=`? -/
def equiv_mem (m₁ m₂ : mem) : Prop :=
∀ id, read_mem m₁ id = read_mem m₂ id

/- TODO: should we define type class has_equiv? -/
infix `≈` := equiv_mem

attribute [simp] equiv_mem

@[simp]
lemma write_var_spec_eq (m₀ : mem) (id : nat) (val : nat) (post : mem → unit → Prop) : passert (id ::= val) m₀ post = post (write_mem m₀ id val) () :=
rfl

@[simp]
lemma read_var_sec_eq (m₀ : mem) (id : nat) (post : mem → nat → Prop) : passert (read_var id) m₀ post = post m₀ (read_mem m₀ id) :=
rfl

/- TODO: move to core -/
attribute [simp] has_bind.and_then

def swap (id₁ id₂ : nat) : prog mem unit :=
do v₁ ← read_var id₁,
   v₂ ← read_var id₂,
   write_var id₁ v₂,
   write_var id₂ v₁

lemma swap_swap (m₀ : mem) (id₁ id₂ : nat) : passert (swap id₁ id₂ >> swap id₁ id₂) m₀ (λ m _, m ≈ m₀) :=
begin
  simp [swap], intros id,
  by_cases id = id₁; by_cases id = id₂; by_cases id₁ = id₂; simp [*] { discharger := tactic.assumption }
end

def equiv_prog (p₁ p₂ : prog mem α) : Prop :=
∀ m₀ : mem,
  match p₁ m₀ : _ → Prop with
  | some (a₁, m₁) :=
    match p₂ m₀ with
    | some (a₂, m₂) := m₁ ≈ m₂ ∧ a₁ = a₂
    | none          := false
    end
  | none         :=
    match p₂ m₀ with
    | some (a₂, m₂) := false
    | none          := true
    end
  end

end experiment