From iris.heap_lang Require Export lifting metatheory.
From iris.heap_lang Require Import notation proofmode.
From iris_examples.logrel_heaplang Require Import ltyping ltyping_safety.
From iris.algebra Require Import gmap auth agree.
From iris_examples.logrel_heaplang.lib Require Import invariants arrays.

Set Default Proof Using "Type".

(* vec = ref array *)
Section Vectors.

  (* API *)
  Definition use : val :=
    λ: <> "a" "f" ,
    let: "b" := copy "a" (!"a") in
    let: "v" := ref "b" in "f" #() "v".

  Definition idx : val :=
    λ: <> "v" "i", if: ((#-1) < "i") && ("i" < !(!"v")) then SOME "i" else NONE.

  Definition get : val :=
    λ: <> "v" "i", unsafe_get (!"v") "i".

  Definition set : val :=
    λ: <> "v" "i" "x", (unsafe_set (!"v") "i") "x".

  Definition app : val :=
    λ: <> , rec: "f" "v" "x" :=
      let: "a" := !"v" in
      let: "b" := copy "a" (!"a" + #1) in
      "b" + !"a" + #1 <- "x";;
      if: CAS "v" "a" "b" then #() else "f" "v" "x".

  (* History *)
  Definition Hist := gmap loc nat.
  Definition HistUR  := gmapUR loc (agreeR (leibnizO nat)).
  Definition hist : Hist → HistUR := fmap (λ v: nat, to_agree (v : leibnizO nat)).
  Implicit Types (h : Hist) (a: loc).

  Context `{!heapG Σ,
            !inG Σ (authR mnatUR), (* maximum length *)
            !inG Σ (authR HistUR) (* history *)
           }.
  Variable (N: namespace).

  (* Maximum Length *)
  Global Instance auth_mnat_persistent (n: mnat) γ: Persistent (own γ (◯ n)).
  Proof.
    eapply own_core_persistent, auth_frag_core_id, mnat_core_id.
  Qed.

  Lemma auth_mnat_alloc (n: mnat) :
    (|==> ∃ γ, own γ (● n) ∗ own γ (◯ n))%I.
  Proof.
    iMod (own_alloc (● n ⋅ ◯ n)) as (γ) "[??]". eapply auth_both_valid_2; done.
    eauto with iFrame.
  Qed.

  Lemma auth_mnat_agree γ (m n : mnat):
    own γ (● m) -∗ own γ (◯ n) -∗ ⌜n ≤ m⌝.
  Proof.
    iIntros "Hγ● Hγ◯".
    iDestruct (own_valid_2 with "Hγ● Hγ◯") as "%".
    eapply auth_both_valid in H as [H _].
    iPureIntro. eapply mnat_included in H. lia.
  Qed.

  Lemma auth_mnat_fork γ (n: mnat):
    own γ (● n)  ==∗  own γ (● n) ∗ own γ (◯ n).
  Proof.
    iIntros "H".
    iMod (own_update _ _ (● n ⋅ ◯ n) with "H") as "[$$]"; eauto.
    apply auth_update_alloc. now eapply mnat_local_update.
  Qed.

  Lemma auth_mnat_previous γ (n m: mnat):
    n ≤ m → own γ (◯ m) -∗ own γ (◯ n).
  Proof.
    intros H. eapply own_mono. eapply (auth_frag_mono n m), mnat_included. lia.
  Qed.

  Lemma auth_mnat_update γ (n p: mnat):
    n ≤ p → own γ (● n) ==∗ own γ (● p).
  Proof.
    iIntros (H) "Hγ●".
    iMod (auth_mnat_fork with "Hγ●") as "[Hγ● Hγ◯]"; eauto.
    iMod (own_update_2 _ _ _ (● p ⋅ ◯ p) with "Hγ● Hγ◯") as "[$ _]"; eauto.
    eapply auth_update, mnat_local_update; lia.
  Qed.

  Section hist.
    (** Conversion to hists and back *)
    Lemma hist_valid h : ✓ hist h.
    Proof. intros l. rewrite lookup_fmap. by case (h !! l). Qed.

    Lemma lookup_hist_None h a : h !! a = None → hist h !! a = None.
    Proof. by rewrite /hist lookup_fmap=> ->. Qed.

    Lemma hist_singleton_included h a m :
      {[a := to_agree m]} ≼ hist h → h !! a = Some m.
    Proof.
      rewrite singleton_included=>  -[a' []].
      rewrite /hist lookup_fmap fmap_Some_equiv => -[m' [-> H]]. rewrite H.
      move=> /Some_included_total /to_agree_included /leibniz_equiv_iff -> //.
    Qed.

    Lemma hist_insert h a m :
      hist (<[a:=m]> h) = <[a:= to_agree (m: leibnizO nat)]> (hist h).
    Proof. by rewrite /hist fmap_insert. Qed.

    Lemma hist_delete a h :
      hist (delete a h) = delete a (hist h).
    Proof. by rewrite /hist fmap_delete. Qed.
  End hist.

  Definition hist_all γ h : iProp Σ := (own γ (● (hist h)))%I.
  Definition hist_entry γ a m : iProp Σ :=
    (own γ (◯ {[ a := to_agree (m : leibnizO nat)]}))%I.

  Global Instance hist_persistent a m γ: Persistent (hist_entry γ a m).
  Proof. eapply _. Qed.

  Lemma hist_alloc: (|==> ∃ γ, hist_all γ ∅)%I.
  Proof.
    iMod (own_alloc (● (hist ∅))) as (γ) "H".
    eapply auth_auth_valid, hist_valid.
    eauto with iFrame.
  Qed.

  Lemma hist_update γ h a m:
    (h !! a = None) →
    hist_all γ h ==∗ hist_all γ (<[a := m]> h) ∗ hist_entry γ a m.
  Proof.
    intros H. iIntros "Hγ●".
    iMod (own_update _ _ _ with "Hγ●") as "H"; eauto.
    2: iModIntro; iDestruct "H" as "[H1 H2]"; iFrame "H1"; eauto.
    rewrite hist_insert. eapply auth_update_alloc.
    eapply alloc_singleton_local_update.
    now eapply lookup_hist_None. done.
  Qed.

  Lemma entry_agree γ a m n: hist_entry γ a m -∗ hist_entry γ a n -∗ ⌜m = n⌝.
  Proof.
    apply bi.wand_intro_r.
    rewrite -own_op -auth_frag_op own_valid uPred.discrete_valid.
    f_equiv. rewrite auth_frag_valid /=. rewrite op_singleton singleton_valid.
      by intros ?%agree_op_invL'.
  Qed.

  Lemma hist_agree γ h a m:
    hist_all γ h -∗ hist_entry γ a m -∗ ⌜h !! a = Some m⌝.
  Proof.
    iIntros "Hh Ha". rewrite /hist_all /hist_entry.
    iDestruct (own_valid_2 with "Hh Ha") as %H.
    eapply auth_both_valid in H as [H _].
    now eapply hist_singleton_included in H.
  Qed.

  Definition Vec_inv γ1 γ2 (l: loc) : iProp Σ :=
    (∃  (h: Hist) (a: loc) (m: mnat),
        l ↦ #a ∗
        own γ1 (● m) ∗
        hist_all γ2 h ∗
        hist_entry γ2 a m ∗
        [∗ map] a ↦ m ∈ h, array (Z.of_nat m) a)%I.

  Definition Vec γ1 γ2 l := inv N (Vec_inv γ1 γ2 l).
  Definition Idx (γ1: gname) (i: mnat) := own γ1 (◯ (S i: mnat)).

  Lemma Vec_project_array γ1 γ2 a m l:
    hist_entry γ2 a m -∗ Vec γ1 γ2 l -∗ inv N (array m a).
  Proof.
    iIntros "#E #V". iApply inv_acc; eauto.
    iNext. iModIntro. iDestruct 1 as (h b n) "(Hl & Hγ1● & Hγ2● & Hγ2◯ & Hst)".
    iAssert (⌜h !! a = Some m⌝)%I with "[Hγ2●]" as "#H".
    { iApply (hist_agree with "[Hγ2●]"); eauto. }
    iAssert ((array m a ∗ (array m a -∗ [∗ map] a ↦ m ∈ h, array (Z.of_nat m) a)))%I
      with "[Hst]" as "(Ha & Hst)".
    { iDestruct "H" as %H1.
      now iApply (big_sepM_lookup_acc with "Hst").  }
    iFrame. iIntros "Ha".
    iExists h, b, n. iFrame. now iApply "Hst".
  Qed.

  (* Function Verification *)
  Definition ι γ1 γ2 : val := PairV (#(Z.pos γ1)) (#(Z.pos γ2)).
  Definition β γ1 γ2 : lty Σ := Lty (λ v, ⌜v = ι γ1 γ2⌝)%I.
  Definition Index (S: lty Σ) : lty Σ :=
    Lty (λ v, (□ ∃ γ1 γ2 (i: mnat), (∀ v, S v ∗-∗ β γ1 γ2 v) ∗ ⌜v = #i⌝ ∗ Idx γ1 i)%I).
  Definition Vector (S: lty Σ) : lty Σ :=
    Lty (λ v, (□ ∃ γ1 γ2 (l: loc), (∀ v, S v ∗-∗ β γ1 γ2 v) ∗  ⌜v = #l⌝ ∗ Vec γ1 γ2 l)%I).

  Lemma typing_use: ∅ ⊨ use : ∀ T,  Array N → (∀ S, Vector S → T) → T.
  Proof. unfold use.
    iIntros (s) "!# #T /=".
    iApply wp_value.
    iIntros (S) "!#". wp_pures.
    iIntros (v1) "!# #Array". wp_pures.
    iIntros (v2) "!# #Cont". wp_pures.
    iDestruct "Array" as (a m) "[Array ->]".
    wp_bind (! #a)%E. iPoseProof (array_length_inv with "Array") as "L".
    iInv "L" as "Ha" "Hclose". wp_load. iMod ("Hclose" with "Ha") as "_".
    iModIntro. wp_apply (copy_spec); eauto with lia.
    iIntros (b) "B". wp_pures.
    wp_bind (ref _)%E. wp_apply wp_fupd. wp_alloc l as "Hl".
    iMod (auth_mnat_alloc m) as (γ1) "[Hγ1● _]".
    iMod (hist_alloc) as (γ2) "Hγ2●".
    iMod (hist_update γ2 ∅ b m with "[Hγ2●]") as "[Hγ2● #Hγ2◯]"; eauto.
    iMod (inv_alloc N _ (Vec_inv γ1 γ2 l) with "[Hγ1● Hγ2● Hl B]") as "#I".
    - iNext. iExists {[ b := m ]}, b, m. iFrame. iSplit; eauto.
      now rewrite big_sepM_singleton.
    - iModIntro. wp_pures. iSpecialize ("Cont" $! (β γ1 γ2)).
      wp_bind (v2 _). wp_apply (wp_wand with "Cont").
      iIntros (v) "#H". wp_apply ("H"). iModIntro.
      iExists γ1, γ2, l. repeat iSplit; eauto.
  Qed.

  Lemma typing_idx: ∅ ⊨ idx : (∀ S, Vector S → lty_int → Option (Index S)).
  Proof. unfold idx.
    iIntros (s) "!# T /=". iApply wp_value.
    iIntros (S) "!#". wp_pures.
    iIntros (v1) "!# #Vec". wp_pures.
    iIntros (v2) "!# #Int". wp_pures.
    iDestruct "Vec" as (γ1 γ2 l) "(EQ & -> & V)".
    iDestruct "Int" as (n) "->". wp_pures.
    bool_decide H; last (wp_pures; now iLeft).
    wp_pures. wp_bind (! #l)%E. iInv "V" as (h a m) "(H1 & H2 & H3 & #Ha & H4)" "Hclose".
    wp_load. iMod (auth_mnat_fork with "H2") as "[H2 #Hi]".
    iMod("Hclose" with "[H1 H2 H3 H4]") as "_".
    iNext; iExists h, a, m; iFrame; eauto.
    iModIntro. iPoseProof (Vec_project_array with "Ha V") as "I".
    wp_bind (! _)%E. iPoseProof (array_length_inv with "I") as "L".
    iInv "L" as "H" "Hclose". wp_load. iMod ("Hclose" with "H") as "_".
    iModIntro. wp_pures.
    bool_decide H'; last (wp_pures; now iLeft).
    wp_pures. iRight. iExists #n. iSplit; eauto.
    destruct (Z_to_nat n) as [k ->]; first lia.
    iExists γ1, γ2, k.
    iModIntro. repeat iSplit; eauto.
    iApply auth_mnat_previous; eauto. lia.
  Qed.

  Lemma typing_get: ∅ ⊨ get : (∀ S, Vector S → Index S → lty_int).
  Proof. unfold get.
    iIntros (s) "!# T /=". iApply wp_value.
    iIntros (S) "!#". wp_pures.
    iIntros (v1) "!# #Vec". wp_pures.
    iIntros (v2) "!# #Idx". wp_pures.
    iDestruct "Vec" as (γ1 γ2 l) "(EQ & -> & V)".
    iDestruct "Idx" as (γ1' γ2' i) "(EQ' & -> & I)".
    iAssert (β γ1 γ2 (ι γ1' γ2'))%I as "%".
    { iApply "EQ". iApply "EQ'". now iPureIntro. }
    injection H as -> ->.
    wp_bind (! _)%E. iInv "V" as (h a m) "(H1 & H2 & H3 & #Ha & H4)" "Hclose".
    wp_load. iDestruct (auth_mnat_agree with "H2 I") as %H.
    iMod("Hclose" with "[H1 H2 H3 H4]") as "_".
    iNext; iExists h, a, m; iFrame; eauto.
    iModIntro. wp_apply (unsafe_get_spec); eauto. iSplit.
    iApply Vec_project_array; eauto. iPureIntro; lia.
  Qed.

  Lemma typing_set: ∅ ⊨ set : (∀ S, Vector S → Index S → lty_int → lty_unit).
  Proof. unfold set.
    iIntros (s) "!# T /=". iApply wp_value.
    iIntros (S) "!#". wp_pures.
    iIntros (v1) "!# #Vec". wp_pures.
    iIntros (v2) "!# #Idx". wp_pures.
    iIntros (v3) "!# #Int". wp_pures.
    iDestruct "Vec" as (γ1 γ2 l) "(EQ & -> & V)".
    iDestruct "Idx" as (γ1' γ2' i) "(EQ' & -> & I)".
    iDestruct "Int" as (n) "->".
    iAssert (β γ1 γ2 (ι γ1' γ2'))%I as "%".
    { iApply "EQ". iApply "EQ'". now iPureIntro. }
    injection H as -> ->.
    wp_bind (! _)%E. iInv "V" as (h a m) "(H1 & H2 & H3 & #Ha & H4)" "Hclose".
    wp_load. iDestruct (auth_mnat_agree with "H2 I") as %H.
    iMod("Hclose" with "[H1 H2 H3 H4]") as "_".
    iNext; iExists h, a, m; iFrame; eauto.
    iModIntro. wp_apply (unsafe_set_spec); eauto. iSplit.
    iApply Vec_project_array; eauto. iPureIntro; lia.
  Qed.

  Lemma typing_app: ∅ ⊨ app : (∀ S, Vector S → lty_int → lty_unit).
  Proof.
    iIntros (s) "!# #T /=".
     unfold app.
    iApply wp_value.
    iIntros (S) "!#". wp_pures.
    iIntros (v1) "!# #Vec". wp_pures.
    iIntros (v2) "!# #Int". wp_pures. fold app.
    iLöb as "IH" forall (v1 v2) "Vec Int".
    iDestruct "Vec" as (γ1 γ2 l) "(EQ & -> & Vec)".
    iDestruct "Int" as (n) "->".
    wp_bind (! #l)%E.
    iInv "Vec" as (h a m) "(Hl & Hγ1● & Hγ2 & #Ha & Hst)" "Hclose". wp_load.
    iMod ("Hclose" with "[Hl Hγ1● Hγ2 Hst]") as "_". iNext.  iExists h, a, m. iFrame; eauto.
    clear h. iModIntro. iPoseProof (Vec_project_array with "Ha Vec") as "#HA".
    wp_pures. wp_bind (! #a)%E. iPoseProof (array_length_inv with "HA") as "L".
    iInv "L" as "H" "Hclose". wp_load. iMod ("Hclose" with "H") as "_".
    iModIntro. wp_pures.
    wp_apply (copy_spec with "[]"); first eauto with lia.
    iIntros (b) "A". wp_pures.
    wp_bind (! #a)%E. iInv "L" as "Hb" "Hclose". wp_load.
    iMod("Hclose" with "Hb") as "_". iModIntro. wp_pures.
    iPoseProof (array_entry_acc m with "[] A") as "[E A]"; first (iPureIntro; lia).
    iDestruct "E" as (v) "E". wp_store. iSpecialize ("A" with "[E]"); first by iExists _.
    wp_bind (CmpXchg _ _ _).
    iInv "Vec" as (h a' m') "(Hl & Hγ1● & Hγ2 & #Ha' & Hst)" "Hclose".
    wp_cmpxchg as H|H.
    - injection H as ?; subst.
      iPoseProof (entry_agree with "Ha Ha'") as "%"; subst m'.
      iMod (auth_mnat_update _ _ (m + 1)%nat with "Hγ1●") as "Hγ1●"; try lia.
      iAssert (⌜h !! b = None⌝)%I with "[A Hst]" as "%".
      { destruct (h !! b) eqn: H; eauto.
        iPoseProof (big_sepM_insert_acc with "Hst") as "[Hst _]"; eauto.
        iPoseProof (array_length_acc with "A") as "[H1 _]".
        iPoseProof (array_length_acc with "Hst") as "[H2 _]".
        iPoseProof (mapsto_valid_2 with "H1 H2") as "%". exfalso.
        destruct (frac.frac_valid' 2%Qp) as [H1 _]. eapply H1 in H0.
        now eapply Qp_not_plus_q_ge_1 in H0. }
      iMod (hist_update _ _ _ (m + 1)%nat with "Hγ2") as "[Hγ2 #Hb_ent]"; eauto.
      iMod ("Hclose" with "[Hl Hγ1● Hγ2 A Hst]") as "_".
      iNext. iExists (<[b:=(m + 1)%nat]> h), b, ((m + 1)%nat). iFrame.
      iSplit; eauto. rewrite big_sepM_insert; eauto; iFrame.
      replace ((m + 1)%nat : Z) with (m + 1) by lia.
      iFrame. iModIntro. wp_pures. done.
    - iMod ("Hclose" with "[Hl Hγ1● Hγ2 A Hst]") as "_".
      iNext. iExists h, a', m'. iFrame; eauto.
      iModIntro. wp_pures.
      iApply "IH". iModIntro. iExists γ1, γ2, l. repeat iSplit; eauto.
      iModIntro. iExists n. done.
  Qed.
End Vectors.