-- SPLean/Theories/WP1.lean

import Lean

import Mathlib.Data.Finmap

import Mathlib.Data.List.Indexes

import SPLean.Common.State

import SPLean.Common.Util

import SPLean.Theories.HProp

import SPLean.Theories.XSimp

import SPLean.Theories.XChange

import SPLean.Theories.SepLog

import SPLean.Theories.WPUtil

open trm val prim

namespace Theories

def wp (t : trm) (Q : val → hProp) : hProp :=
  fun s ↦ eval s t Q

section tactics

open Lean Elab Tactic

section xapp

end xapp

end tactics

open AList

section funs_fixs_eval_like

variable (xs : List var) (vs : List val) (t : trm) (v0 : trm)
  (heqt : t = trm_apps v0 ts)
  (hconv : trms_to_vals ts = vs)
  (hform : var_funs xs vs.length)   -- NOTE: can be relaxed to `vs.length ≤ xs.length`

variable (f : var) (hf : f ∉ xs)

end funs_fixs_eval_like

end Theories

open Theories

open Lean.Elab.Tactic in

-- Repository definitions used:
content: notation:max "emp" => hempty

syntax "fun" ident+ " => " lang : lang

syntax "⟨" term "⟩" : lang

syntax "⟨" term ":" term "⟩" : lang

content: macro "xsimp" : tactic =>
  `(tactic| (
    xsimp_start
    repeat xsimp_step
    try rev_pure
    try hide_mvars
    try hsimp
    rotate_left

  ))

content: macro "xif" : tactic => do
  `(tactic|
  (xseq_xlet_if_needed; xstruct_if_needed; apply xif_lemma))

content: macro "xif" : tactic => `(tactic| (xwp; xif))

syntax " := " : bop

content: macro "xstruct" : tactic => do
  `(tactic| apply xstruct_lemma)

content: macro "xval" : tactic => do
  `(tactic| (xstruct_if_needed; apply xval_lemma))

content: macro "xval" : tactic => `(tactic| (xwp; xval))

content: macro "xapp" : tactic =>
  `(tactic|
    (xapp_nosubst;
     try xapp_try_subst
     first
       | done
       | all_goals try srw wp_equiv
         all_goals try subst_vars))

content: notation:max "⌜" P "⌝" => hpure P

content: macro "xwp" : tactic =>
  `(tactic|
    (intros
     first | apply xwp_lemma_fix; rfl
           | apply xwp_lemma_fun; rfl))

content: macro "xwp" : tactic =>
  `(tactic|
    (intros
     try srw trm_apps1
     srw ?trm_apps2
     first
           
           | (apply xwp_lemma_funs; rfl; rfl; rfl; sdone)=> //
           | apply wp_of_wpgen
     all_goals try simp [wpgen, subst, isubst, subst, trm_apps, AList.lookup, List.dlookup]))

content: macro "xseq" : tactic => do
  `(tactic| (xstruct_if_needed; apply xseq_lemma))

content: macro "xchange" l:term : tactic =>
  `(tactic| (xchange_core $l; xsimp))

content: macro "∃ʰ" xs:Lean.explicitBinders ", " b:term : term => Lean.expandExplicitBinders ``hexists xs b

content: macro "xpull" : tactic =>
  `(tactic| (
    xpull_start
    repeat' xsimp_step
    try rev_pure
    hsimp
  ))

content: macro_rules
  | `([lang| ()])                       => `(trm_val (val_unit))
  | `([lang| $n:num])                   => `(trm_val (val_int $n))
  | `([lang| $t1 $t2])                  => `(trm_app [lang| $t1] [lang| $t2])
  | `([lang| if $t1 then $t2 else $t3]) => `(trm_if [lang| $t1] [lang| $t2] [lang| $t3])
  | `([lang| if $t1 then $t2 end])      => `(trm_if [lang| $t1] [lang| $t2] (trm_val val_unit))
  | `([lang| let $x := $t1:lang in $t2:lang])     =>
    `(trm_let $(%x) [lang| $t1] [lang| $t2])
  | `([lang| ref $x := $t1:lang in $t2:lang])     =>
    `(trm_ref $(%x) [lang| $t1] [lang| $t2])
  | `([lang| alloc $t1:lang as $x in $t2:lang])   =>
    `(trm_alloc $(%x) [lang| $t1] [lang| $t2])
  | `([lang| $t1 ; $t2])                => `(trm_seq [lang| $t1] [lang| $t2])
  | `([lang| fun_ $xs* => $t])             => do
    let xs <- xs.mapM fun x => `(term| $(%x))
    `(trm_funs [ $xs,* ] [lang| $t])
  | `([lang| fun $xs* => $t])             => do
    let xs <- xs.mapM fun x => `(term| $(%x))
    `(val_funs [ $xs,* ] [lang| $t])
  | `([lang| fix_ $f $xs* => $t])    => do
      let xs <- xs.mapM fun x => `(term| $(%x))
      `(trm_fixs $(%f) [ $xs,* ] [lang| $t])
  | `([lang| fix $f $xs* => $t])    => do
      let xs <- xs.mapM fun x => `(term| $(%x))
      `(val_fixs $(%f) [ $xs,* ] [lang| $t])
  
  | `([lang| free $t])                  => `(trm_val (val_prim val_free) [lang| $t])
  | `([lang| not $t])                   => `(trm_val (val_prim val_not) [lang| $t])
  
  | `([lang| !$t])                      => `(trm_val val_get [lang| $t])
  | `([lang| $t1 := $t2])               => `(trm_val val_set [lang| $t1] [lang| $t2])
  | `([lang| $t1 + $t2])                => `(trm_val val_add [lang| $t1] [lang| $t2])
  | `([lang| $t1 * $t2])                => `(trm_val val_mul [lang| $t1] [lang| $t2])
  | `([lang| $t1 - $t2])                => `(trm_val val_sub [lang| $t1] [lang| $t2])
  | `([lang| $t1 / $t2])                => `(trm_val val_div [lang| $t1] [lang| $t2])
  | `([lang| $t1 < $t2])                => `(trm_val val_lt [lang| $t1] [lang| $t2])
  | `([lang| $t1 > $t2])                => `(trm_val val_gt [lang| $t1] [lang| $t2])
  | `([lang| $t1 <= $t2])               => `(trm_val val_le [lang| $t1] [lang| $t2])
  | `([lang| $t1 >= $t2])               => `(trm_val val_ge [lang| $t1] [lang| $t2])
  | `([lang| -$t])                      => `(trm_val val_opp [lang| $t])
  | `([lang| $t1 = $t2])                => `(trm_val val_eq [lang| $t1] [lang| $t2])
  | `([lang| $t1 != $t2])               => `(trm_val val_neq [lang| $t1] [lang| $t2])
  | `([lang| $t1 mod $t2])              => `(trm_val val_mod [lang| $t1] [lang| $t2])
  | `([lang| $t1 ++ $t2])               => `(trm_val val_ptr_add [lang| $t1] [lang| $t2])
  | `([lang| ($t)]) => `([lang| $t])
  | `([lang| ⟨$t : $tp⟩]) => `(trm_val (($t : $tp)))
  | `([lang| for $x in [$n1 : $n2] { $t } ]) =>
    `(trm_for $(%x) [lang| $n1] [lang| $n2] [lang| $t])
  | `([lang| while $c:lang { $t:lang } ]) =>
     `(trm_while [lang| $c] [lang| $t] )

content: macro_rules
  | `([lang| len $p])                => `(trm_val val_array_length [lang| $p])
  | `([lang| $arr[$i] ])              => `(trm_val val_array_get [lang| $arr] [lang| $i])
  
  | `([lang| $arr[$i] := $v])        => `(trm_app val_array_set [lang| $arr] [lang| $i] [lang| $v])
  | `([lang| mkarr $n:lang $v:lang]) => `(trm_val val_array_make [lang| $n] [lang| $v])

content: macro_rules
  | `({ $P }[$t:lang]{$v, $Q}) => `(triple [lang| $t] $P (fun $v => $Q))
  | `({ $P }[$t:lang]{$Q})     => `(triple [lang| $t] $P (fun _ => $Q))
  | `(WP[$t:lang]{$v, $Q})     => `(wp [lang| $t] (fun $v => $Q))
  | `(WP[$t:lang]{$Q})         => `(wp [lang| $t] (fun _ => $Q))

content: macro_rules
  | `(tactic| xstep $(t)? ) => `(tactic| (xwp; xapp $(t)?))

content: macro_rules | `($x ∗ $y)   => `(binop% HStar.hStar $x $y)

  inductive trm : Type where
    | trm_val   : val -> trm
    | trm_var   : var -> trm
    | trm_fun   : var -> trm -> trm
    | trm_fix   : var -> var -> trm -> trm
    | trm_app   : trm -> trm -> trm
    | trm_seq   : trm -> trm -> trm
    | trm_let   : var -> trm -> trm -> trm
    | trm_if    : trm -> trm -> trm -> trm
    | trm_for   : var -> trm -> trm -> trm -> trm
    | trm_while : trm -> trm -> trm
    | trm_ref   : var → trm → trm → trm
    | trm_alloc : var → trm → trm → trm

  inductive val : Type where
    | val_unit   : val
    | val_bool   : Bool → val
    | val_int    : Int → val
    | val_real   : ℝ → val
    | val_loc    : loc → val
    | val_prim   : prim → val
    | val_fun    : var -> trm -> val
    | val_fix    : var -> var -> trm -> val
    | val_uninit : val
    | val_error : val

content: inductive prim where
  
  | val_get : prim
  | val_set : prim
  
  | val_neg : prim
  | val_opp : prim
  | val_eq : prim
  | val_add : prim
  | val_neq : prim
  | val_sub : prim
  | val_mul : prim
  | val_div : prim
  | val_mod : prim
  
  | val_le : prim
  | val_lt : prim
  | val_ge : prim
  | val_gt : prim
  | val_ptr_add : prim

content: abbrev var := String

content: abbrev loc := Nat

content: def hProp := heap -> Prop

content: abbrev heap := Heap.heap val

content: abbrev Heap.heap (val : Type) := Finmap (λ _ : loc ↦ val)

content: inductive eval : state → trm → (val → state → Prop) -> Prop where
  | eval_val : forall s v Q,
      Q v s ->
      eval s (trm_val v) Q
  | eval_fun : forall s x t1 Q,
      Q (val_fun x t1) s ->
      eval s (trm_fun x t1) Q
  | eval_fix : forall s f x t1 Q,
      Q (val_fix f x t1) s ->
      eval s (trm_fix f x t1) Q
  | eval_app_arg1 : forall s1 t1 t2 Q1 Q,
      ¬ trm_is_val t1 ->
      eval s1 t1 Q1 ->
      (forall v1 s2, Q1 v1 s2 -> eval s2 (trm_app v1 t2) Q) ->
      eval s1 (trm_app t1 t2) Q
  | eval_app_arg2 : forall s1 (v1 : val) t2 Q1 Q,
      ¬ trm_is_val t2 ->
      eval s1 t2 Q1 ->
      (forall v2 s2, Q1 v2 s2 -> eval s2 (trm_app v1 v2) Q) ->
      eval s1 (trm_app v1 t2) Q
  | eval_app_fun : forall s1 v1 (v2 :val) x t1 Q,
      v1 = val_fun x t1 ->
      eval s1 (subst x v2 t1) Q ->
      eval s1 (trm_app v1 v2) Q
  | eval_app_fix : forall s (v1 v2 : val) f x t1 Q,
      v1 = val_fix f x t1 ->
      eval s (subst x v2 (subst f v1 t1)) Q ->
      eval s (trm_app v1 v2) Q
  | eval_seq : forall Q1 s t1 t2 Q,
      eval s t1 Q1 ->
      (forall v1 s2, Q1 v1 s2 -> eval s2 t2 Q) ->
      eval s (trm_seq t1 t2) Q
  | eval_let : forall Q1 s x t1 t2 Q,
      eval s t1 Q1 ->
      (forall v1 s2, Q1 v1 s2 -> eval s2 (subst x v1 t2) Q) ->
      eval s (trm_let x t1 t2) Q
  | eval_if : forall s (b : Bool) t1 t2 Q,
      eval s (if b then t1 else t2) Q ->
      eval s (trm_if (val_bool b) t1 t2) Q
  | eval_unop : forall op s v1 P Q,
      evalunop op v1 P ->
      purepostin s P Q ->
      eval s (trm_app op v1) Q
  | eval_binop : forall op s (v1 v2 : val) P Q,
      evalbinop op v1 v2 P ->
      purepostin s P Q ->
      eval s (trm_app (trm_app op v1) v2) Q
  | eval_ref : forall s x t1 t2 (Q Q₁ : val → state → Prop),
    eval s t1 Q₁ →
    (∀ v1 s1, Q₁ v1 s1 → ∀ p ∉ s1,
     eval (s1.insert p v1) (subst x p t2) fun v s ↦ Q v (s.erase p)) →
    eval s (trm_ref x t1 t2) Q
  | eval_get : forall s p Q,
      p ∈ s ->
      Q (read_state p s) s ->
      eval s (trm_app val_get (val_loc p)) Q
  | eval_set : forall s p v Q,
      v = trm_val v' ->
      p ∈ s ->
      Q val_unit (Finmap.insert p v' s) ->
      eval s (trm_app (trm_app val_set (val_loc p)) v) Q
  | eval_alloc_arg : forall s Q₁ Q,
    ¬ trm_is_val t1 →
    eval s t1 Q₁ →
    (∀ v' s', Q₁ v' s' → eval s' (trm_alloc x v' t2) Q) →
    eval s (trm_alloc x t1 t2) Q
  | eval_alloc : forall (sa : state) (n : ℤ) Q,
    n ≥ 0 →
    (∀ (p : loc) (sb : state),
      sb = conseq (make_list n.natAbs val_uninit) p →
      p ≠ null →
      Finmap.Disjoint sa sb →
      eval (sb ∪ sa) (subst x p t2) fun v s ↦ Q v (s \ sb)) →
    eval sa (trm_alloc x n t2) Q
  
  
  
  
  
  
  
  
  | eval_for (n₁ n₂ : Int) (Q : val -> state -> Prop) :
    eval s (if (n₁ < n₂) then
               (trm_seq (subst x n₁ t₁) (trm_for x (val_int (n₁ + 1)) n₂ t₁))
            else val_unit) Q ->
    eval s (trm_for x n₁ n₂ t₁) Q
  | eval_while (t₁ t₂ : trm) (Q : val -> state -> Prop) :
    eval s t₁ Q₁ ->
    (∀ s v, Q₁ v s -> eval s (trm_if v (trm_seq t₂ (trm_while t₁ t₂)) val_unit) Q) ->
    eval s (trm_while t₁ t₂) Q

content: class HStar (α : Type u) (β : Type v) (γ : outParam (Type w)) where
   
  hStar : α → β → γ

content: def hpure (P : Prop) : hProp :=
  hexists (fun (_ : P) => emp)

content: def hexists {A} (J : A → hProp) : hProp :=
  fun h => exists x, J x h

content: def hempty : hProp :=
  fun h => (h = ∅)

content: abbrev himpl (H1 H2 : hProp) : Prop :=
  forall h, H1 h -> H2 h

content: def qimpl {A} (Q1 Q2 : A → hProp) : Prop :=
  forall (v:A), Q1 v ==> Q2 v

content: def purepostin (s : state) (P : val → Prop) (Q : val → state → Prop) : Prop :=
  
  forall v, P v → Q v s

content: abbrev state := Finmap (λ _ : loc ↦ val)

content: abbrev tohProp (h : heap -> Prop) : hProp := h

content: abbrev ofhProp (h : val -> hProp) : val -> heap -> Prop := h

content: infixr:51 " ==> " => himpl

content: infixr:51 " ===> " => qimpl

content: notation:max "emp" => hempty

content: infixr:55 " ∗ " => HStar.hStar

content: notation:max "⌜" P "⌝" => hpure P

content: notation "funloc" p "↦" H =>
  fun (r : val) ↦ hexists (fun p ↦ ⌜r = val_loc p⌝ ∗ H)

-- Repository lemmas used:
content: lemma eval_conseq s t Q1 Q2 :
  eval s t Q1 →
  Q1 ===> Q2 →
  eval s t Q2

content: lemma qstarE α (Q1 : α → hProp)  (H : hProp):
  Q1 ∗ H = fun x => Q1 x ∗ H

content: lemma eval_frame (h1 h2 : state) t (Q : val -> hProp) :
  eval h1 t (ofhProp Q) →
  Finmap.Disjoint h1 h2 →
  eval (h1 ∪ h2) t (Q ∗ (tohProp (fun h ↦ h = h2)))

content: lemma hstar_comm H1 H2 :
  H1 ∗ H2 = H2 ∗ H1

content: lemma hprop_op_comm (op : hProp → hProp → hProp) :
  (forall H1 H2, op H1 H2 ==> op H2 H1) →
  (forall H1 H2, op H1 H2 = op H2 H1)

content: lemma himpl_antisym H1 H2:
  (H1 ==> H2) → (H2 ==> H1) → (H1 = H2)

content: lemma hstar_inv (H1 H2 : hProp) h:
  (H1 ∗ H2) h →
  exists h1 h2, H1 h1 ∧ H2 h2 ∧ Finmap.Disjoint h1 h2 ∧ h = h1 ∪ h2

content: lemma himpl_trans H2 H1 H3 :
  (H1 ==> H2) → (H2 ==> H3) → (H1 ==> H3)

-- Local definitions used:
content: def wp (t : trm) (Q : val → hProp) : hProp :=
  fun s ↦ eval s t Q

-- Local lemmas used:
content: lemma wp_conseq t Q1 Q2 :
  Q1 ===> Q2 →
  wp t Q1 ==> wp t Q2

content: lemma wp_frame t H Q :
  (wp t Q) ∗ H ==> wp t (Q ∗ H)

-- Target theorem and proof:
lemma xwhile_inv_basic_lemma (I : Bool -> α -> hProp) R
  -- (F1 F2 : formula)
  :
  WellFounded R ->
  -- structural F1 ->
  -- structural F2 ->
  (H ==> H' ∗ ∃ʰ b a, I b a) ->
  (∀ b X, I b X ==> wp F1 (fun bv => I b X ∗ ⌜bv = b⌝)) ->
  (∀ X, I true X ==> wp F2 (fun _ => ∃ʰ b X', ⌜R X' X⌝ ∗ I b X')) ->
  H ==> wp (trm_while F1 F2) (fun _ => H' ∗ ∃ʰ a, I false a)
:= by
  move=> wf hini hf1 hf2
  xchange hini=> b sR
  move: b
  apply WellFounded.induction wf sR=> X ih []
  -- apply eval.eval_while
  -- unfold wpgen_while ; unfold_let ; xstruct ; xsimp=> [] sR hstep; rename_i wfR
  -- frame H' out, using `structural`?
  { xchange hf1
    apply himpl_trans;  rotate_left
    { srw hstar_comm; apply wp_frame }
    xsimp
    apply himpl_trans; rotate_left
    { move=> ? H
      apply eval.eval_while; apply H
      move=> > // }
    apply wp_conseq=> ? /=; xpull
    xwp; xif=> // _; xwp; xval; xsimp }
  -- move: sR
  -- apply himpl_trans;  rotate_left
  -- { srw hstar_comm; apply wp_frame }
  -- xsimp
  apply himpl_trans; rotate_left
  { move=> ? H
    apply eval.eval_while; apply H
    move=> > // }
  xchange hf1
  apply himpl_trans; apply wp_frame

  apply wp_conseq=> ? /==; srw qstarE /=; xpull
  xwp; xif=> // _
  xwp; xseq; xapp hf2=> // ?? /ih;
  srw [2]hstar_comm //; sapply