-- 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

abbrev formula := (val → hProp) → hProp

def mkstruct (F : formula) :=
  fun (Q : val -> hProp) ↦ ∃ʰ Q', F Q' ∗ (Q' -∗ Q)

def structural (F : formula) :=
  forall Q, mkstruct F Q ==> F Q

def structural_pred (S : α -> formula) :=
  ∀ x, structural $ S x

def wpgen_val (v : val) : formula :=
  fun Q ↦ Q v

def wpgen_seq (F1 F2 : formula) : formula :=
  fun Q ↦ F1 (fun _ ↦ F2 Q)

def wpgen_for (v₁ v₂ : trm) (F1 : val -> formula) : formula :=
  mkstruct fun Q =>
    ∃ʰ n₁ n₂ : Int, ⌜v₁ = n₁⌝ ∗ ⌜v₂ = n₂⌝ ∗
      h∀ (S : Int -> formula),
        (let F i :=
          if i < n₂ then
            wpgen_seq (F1 (val_int i)) (S (i + 1))
          else wpgen_val val_unit
        ⌜structural_pred S /\ ∀ i, F i ===> S i⌝ -∗ S n₁ 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

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

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

content: notation:max "emp" => hempty

syntax " := " : bop

syntax "if " lang "then " lang "end " : lang

content: macro "h∀" xs:Lean.explicitBinders ", " b:term : term => Lean.expandExplicitBinders ``hforall xs b

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

syntax "let" ident " := " lang " in" ppDedent(ppLine lang) : lang

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

  ))

syntax " <= " : bop

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

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

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_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)

content: def hProp := heap -> Prop

content: abbrev heap := Heap.heap val

  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

  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

content: abbrev var := String

content: abbrev loc := Nat

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

content: class HWand (α : Type u) (β : Type v) (γ : outParam (Type w)) where
   
  hWand : α → β → γ

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

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 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: infixr:51 " ==> " => himpl

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

content: notation:max "emp" => hempty

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

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

content: infixr:55 " -∗ " => HWand.hWand

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

-- Repository lemmas used:
content: lemma himpl_trans H2 H1 H3 :
  (H1 ==> H2) → (H2 ==> H3) → (H1 ==> H3)

-- Local definitions used:
content: abbrev formula := (val → hProp) → hProp

content: def mkstruct (F : formula) :=
  fun (Q : val -> hProp) ↦ ∃ʰ Q', F Q' ∗ (Q' -∗ Q)

content: def structural (F : formula) :=
  forall Q, mkstruct F Q ==> F Q

content: def structural_pred (S : α -> formula) :=
  ∀ x, structural $ S x

content: def wpgen_val (v : val) : formula :=
  fun Q ↦ Q v

content: def wpgen_seq (F1 F2 : formula) : formula :=
  fun Q ↦ F1 (fun _ ↦ F2 Q)

content: def wpgen_for (v₁ v₂ : trm) (F1 : val -> formula) : formula :=
  mkstruct fun Q =>
    ∃ʰ n₁ n₂ : Int, ⌜v₁ = n₁⌝ ∗ ⌜v₂ = n₂⌝ ∗
      h∀ (S : Int -> formula),
        (let F i :=
          if i < n₂ then
            wpgen_seq (F1 (val_int i)) (S (i + 1))
          else wpgen_val val_unit
        ⌜structural_pred S /\ ∀ i, F i ===> S i⌝ -∗ S n₁ Q )

-- Local lemmas used:
content: lemma mkstruct_erase Q F :
  F Q ==> mkstruct F Q

-- Target theorem and proof:
lemma xfor_inv_lemma (I : Int -> hProp) (a b : Int)
  (F : val -> formula)
  (Q : val -> hProp) :
  structural_pred F ->
  a <= b ->
    (∃ H',
      H ==> I a ∗ H' ∧
      (∀ i, a <= i ∧ i < b -> I i ==> F i fun _ => I (i + 1)) ∧
      (fun _ => I b ∗ H') ===> Q) ->
    H ==> wpgen_for a b F Q
:= by
  move=> sF L ![H' Ma Mb Mc]
  unfold wpgen_for
  apply himpl_trans; rotate_left; apply mkstruct_erase
  unfold_let
  xsimp[a,b]=> //== ls
  srw OfNat.ofNat instOfNat instOfNatNat /== => hs
  -- shave-> ls hs: i + (OfNat.ofNat 1) = i + 1; sdone
  shave P: ∀ i, a <= i ∧ i <= b -> I i ==> S i fun _ => I b
  { move=> i [/[swap] iLb]
    apply (Int.le_induction_down _ _ _ iLb)
    { move=> ?
      xchange (hs b)=> /==
      sby xval }
    move=> i ? ih ?
    xchange hs
    scase_if=> // ?; rotate_left; omega
    xseq
    xchange Mb;
    srw OfNat.ofNat instOfNat instOfNatNat /==
    omega
    apply himpl_trans; rotate_left; apply sF
    unfold mkstruct; xsimp; apply ih; omega }
  xchange Ma; xchange P; omega
  apply himpl_trans; rotate_left; apply ls
  unfold mkstruct; xsimp; apply Mc