(** Cancellable locks *)
From iris.heap_lang Require Export proofmode notation.
From iris.heap_lang Require Import spin_lock.
From iris.base_logic.lib Require Import cancelable_invariants auth saved_prop.
From iris.algebra Require Import auth agree excl frac gmap.
From iris.bi.lib Require Import fractional.

Inductive lockstate :=
  | Locked : frac → lockstate
  | Unlocked.
Canonical Structure lockstateC := leibnizC lockstate.

Instance lockstate_inhabited : Inhabited lockstate := populate Unlocked.

Record flock_name := {
  flock_lock_name : gname;
  flock_cinv_name : gname;
  flock_state_name : gname;
  flock_props_name : gname;
  flock_props_active_name : gname
}.

(* positive so that we have a `Fresh` instance for `gset positive` *)
Definition prop_id := positive.

Class flockG Σ :=
  FlockG {
    flock_stateG :> inG Σ (authR (optionUR (exclR lockstateC)));
    flock_cinvG  :> cinvG Σ;
    flock_lockG  :> lockG Σ;
    flock_props :> authG Σ (optionUR (exclR (gmapC prop_id (iProp Σ))));
    flock_props_activeG :> authG Σ (gmapUR prop_id (agreeR (iProp Σ)))
  }.

Section cinv_lemmas.
  Context `{invG Σ, cinvG Σ}.

  Lemma cinv_alloc_strong (G : gset gname) E N (Φ : gname → iProp Σ) :
    ▷ (∀ γ, ⌜γ ∉ G⌝ → Φ γ) ={E}=∗ ∃ γ, ⌜γ ∉ G⌝ ∧ cinv N γ (Φ γ) ∗ cinv_own γ 1.
  Proof.
    iIntros "HP".
    iMod (own_alloc_strong 1%Qp G) as (γ) "[% H1]"; first done.
    iMod (inv_alloc N _ (Φ γ ∨ own γ 1%Qp)%I with "[HP]") as "#Hinv".
    - iNext. iLeft. by iApply "HP".
    - iExists γ. iModIntro. rewrite /cinv /cinv_own. iFrame "H1".
      iSplit; eauto. iExists _. iFrame "Hinv".
      iIntros "!# !>". iSplit; by iIntros "?".
  Qed.

  Lemma cinv_cancel_vs (P Q : iProp Σ) E N γ :
    ↑N ⊆ E → cinv N γ P
    -∗ (cinv_own γ 1 ∗ ▷ P ={E∖↑N}=∗ cinv_own γ 1 ∗ ▷ Q)
    -∗ cinv_own γ 1 ={E}=∗ ▷ Q.
  Proof.
    iIntros (?) "#Hinv Hvs Hγ". rewrite /cinv.
    iDestruct "Hinv" as (P') "[#HP' Hinv]".
    iInv N as "[HP|>Hγ']" "Hclose"; last first.
    - iDestruct (cinv_own_1_l with "Hγ Hγ'") as %[].
    - iMod ("Hvs" with "[$Hγ HP]") as "[Hγ HQ]".
      { by iApply "HP'". }
      iFrame. iApply "Hclose". iNext. by iRight.
  Qed.

End cinv_lemmas.

Section flock.
  Context `{heapG Σ, flockG Σ}.

  Variable N : namespace.
  Definition flockN := N.@"flock".

  Definition to_props_map (f : gmap prop_id (iProp Σ))
    : gmapUR prop_id (agreeR (iProp Σ)) := to_agree <$> f.
  
  Lemma to_props_map_insert f i P :
    to_props_map (<[i:=P]>f) = <[i:=to_agree P]>(to_props_map f).
  Proof. by rewrite /to_props_map fmap_insert. Qed.

  Lemma to_props_map_dom f :
    dom (gset prop_id) (to_props_map f) = dom (gset prop_id) f.
  Proof. by rewrite /to_props_map dom_fmap_L. Qed.

  Definition all_props (f : gmap prop_id (iProp Σ)) : iProp Σ :=
    ([∗ map] i ↦ R ∈ f, R)%I.

  Lemma all_props_union f g :
    all_props f ∗ all_props g ⊢ all_props (f ∪ g).
  Proof.
    revert g.
    simple refine (map_ind (fun f => ∀ g, all_props f ∗ all_props g -∗ all_props (f ∪ g)) _ _ f); simpl; rewrite /all_props.
    - intros g. by rewrite bi.big_sepM_empty !left_id.
    - intros i P f' Hi IH.
      intros g. rewrite bi.big_sepM_insert; last done.
      iIntros "[[HP Hf] Hg]".
      rewrite -insert_union_l.
      remember (g !! i) as Z. destruct Z as [R|].
      + assert (g = <[i:=R]>(delete i g)) as Hfoo.
        { rewrite insert_delete insert_id; eauto. }
        rewrite {1}Hfoo.
        rewrite bi.big_sepM_insert; last first.
        { apply lookup_delete. }
        iDestruct "Hg" as "[HR Hg]".
        iApply (bi.big_sepM_insert_override_2 _ _ i R P with "[-HP] [HP]"); simpl.
        { apply lookup_union_Some_raw. right. eauto. }
        * iApply (IH g).
          rewrite Hfoo bi.big_sepM_insert; last by apply lookup_delete.
          rewrite delete_insert; last by apply lookup_delete. iFrame.
        * eauto.
      + rewrite bi.big_sepM_insert; last first.
        { apply lookup_union_None. eauto. }
        iFrame. iApply (IH g). iFrame.
  Qed.

  Global Instance all_props_proper : Proper ((≡) ==> (≡))
    (all_props : gmapC prop_id (iProp Σ) → iProp Σ).
  Proof.
    intros f. rewrite /all_props. (* rewrite /all_props /big_opM. *)
    simple refine (map_ind (fun f => ∀ g, f ≡ g → _ ≡ _) _ _ f); simpl.
    - intros g. rewrite bi.big_sepM_empty.
      simple refine (map_ind (fun g => ∅ ≡ g → _ ≡ _) _ _ g); simpl.
      + by rewrite bi.big_sepM_empty.
      + intros j R g' Hj IH Hg'.
        exfalso. specialize (Hg' j). revert Hg'.
        rewrite lookup_insert lookup_empty. inversion 1.
    - intros i P f' Hi IH g Hf'.
      rewrite bi.big_sepM_insert; last done.
      specialize (IH (delete i g)).
      assert (f' ≡ delete i g) as Hf'g.
      { by rewrite -Hf' delete_insert. }
      specialize (IH Hf'g). rewrite IH.
      assert (∃ R, g !! i = Some R ∧ P ≡ R) as [R [Hg HR]].
      { specialize (Hf' i).
        destruct (g !! i) as [R|]; last first.
        - exfalso. revert Hf'. rewrite lookup_insert. inversion 1.
        - exists R. split; auto.
          revert Hf'. rewrite lookup_insert. by inversion 1. }
      rewrite (bi.big_sepM_delete _ g i R); eauto.
      by apply bi.sep_proper.
  Qed.

  Definition flock_inv (γ : flock_name) : iProp Σ :=
    (∃ (s : lockstate) (fp fa : gmap prop_id (iProp Σ)),
       (** fa -- active propositions, fp -- inactive propositions *)
       own (flock_state_name γ) (● (Excl' s)) ∗
       own (flock_props_name γ) (● to_props_map (fp ∪ fa)) ∗
       own (flock_props_active_name γ) (◯ Excl' fa) ∗
       match s with
       | Locked q =>
         cinv_own (flock_cinv_name γ) (q/2) ∗
         locked (flock_lock_name γ) ∗
         all_props fp
       | Unlocked => own (flock_props_active_name γ) (● Excl' fa) ∗
                     all_props fa ∗
                     ⌜fp = ∅⌝
       end)%I.

  Definition is_flock (γ : flock_name) (lk : val) : iProp Σ :=
    (cinv (flockN .@ "inv") (flock_cinv_name γ) (flock_inv γ) ∗
     is_lock (flockN .@ "lock") (flock_lock_name γ) lk
         (own (flock_state_name γ) (◯ (Excl' Unlocked))))%I.
  
  Definition flock_res (γ : flock_name) (s : prop_id) (R : iProp Σ) : iProp Σ :=
    (own (flock_props_name γ) (◯ {[ s := to_agree R ]}))%I.

  Global Instance is_flock_persistent γ lk : Persistent (is_flock γ lk).
  Proof. apply _. Qed.

  Global Instance flock_res_persistent γ s R : Persistent (flock_res γ s R).
  Proof. apply _. Qed.

  Definition unflocked (γ : flock_name) (q : frac) : iProp Σ :=
    cinv_own (flock_cinv_name γ) q.

  Definition flocked
    (γ : flock_name) (q : frac) (f : gmap prop_id (iProp Σ)) : iProp Σ :=
    (own (flock_state_name γ) (◯ (Excl' (Locked q)))
     ∗ cinv_own (flock_cinv_name γ) (q/2)
     ∗ own (flock_props_active_name γ) (● Excl' f)
     ∗ ▷ all_props f)%I.

  Lemma unflocked_op (γ : flock_name) (π1 π2 : frac) :
    unflocked γ (π1+π2) ⊣⊢ unflocked γ π1 ∗ unflocked γ π2.
  Proof. by rewrite /unflocked fractional. Qed.

  Global Instance unflocked_fractional γ : Fractional (unflocked γ).
  Proof. intros p q. apply unflocked_op. Qed.
  Global Instance unflocked_as_fractional γ π :
    AsFractional (unflocked γ π) (unflocked γ) π.
  Proof. split. done. apply _. Qed.

  Lemma flock_res_alloc_unflocked γ lk π R :
    is_flock γ lk -∗ unflocked γ π -∗ ▷ R ={⊤}=∗ ∃ s, flock_res γ s R ∗ unflocked γ π.
  Proof.
    iIntros "Hl Hunfl HR". rewrite /is_flock. iDestruct "Hl" as "(#Hcinv & #Hlk)".
    iMod (cinv_open with "Hcinv Hunfl") as "(Hstate & Hunfl & Hcl)"; first done.
    rewrite {2}/flock_inv.
    iDestruct "Hstate" as ([q|] fp fa) "(>Hstate & >Hauth & >Hfactive & Hrest)".
    - iDestruct "Hrest" as "(>Hcown & >Hlocked2 & Hfp)".
      pose (s := (fresh (dom (gset prop_id) (fp ∪ fa)))).
      iMod (own_update  with "Hauth") as "Hauth".
      { apply (auth_update_alloc _ (to_props_map (<[s := R]> fp ∪ fa))
                                 {[ s := to_agree R ]}).
        rewrite -insert_union_l.
        rewrite to_props_map_insert.
        apply alloc_local_update; last done.
        apply (not_elem_of_dom (to_props_map (fp ∪ fa)) s (D:=gset prop_id)).
        rewrite to_props_map_dom.
        apply is_fresh. }
      iDestruct "Hauth" as "[Hauth Hres]".
      iExists s. iFrame "Hres Hunfl".
      iApply ("Hcl" with "[-]").
      iNext. iExists _,_,_. iFrame. iFrame "Hcown Hlocked2".
      rewrite /all_props bi.big_sepM_insert. iFrame.
      apply (not_elem_of_dom _ s (D:=gset prop_id)).
      assert (s ∉ (dom (gset prop_id) (fp ∪ fa))) as Hs.
      { apply is_fresh. }
      revert Hs. rewrite dom_union_L not_elem_of_union. set_solver.
    - iDestruct "Hrest" as "(>Hactive & Hfa & >%)".
      simplify_eq/=. rewrite left_id.
      pose (s := (fresh (dom (gset prop_id) fa))).
      iMod (own_update  with "Hauth") as "Hauth".
      { apply (auth_update_alloc _ (to_props_map (<[s := R]> fa))
                                 {[ s := to_agree R ]}).
        rewrite to_props_map_insert.
        apply alloc_local_update; last done.
        apply (not_elem_of_dom (to_props_map fa) s (D:=gset prop_id)).
        rewrite to_props_map_dom.
        apply is_fresh. }
      iDestruct "Hauth" as "[Hauth Hres]".
      iExists s. iFrame "Hres Hunfl".
      iMod (own_update_2 _ _ _ ((● Excl' (<[s:=R]>fa)) ⋅ (◯ Excl' (<[s:=R]>fa)))
              with "Hactive Hfactive") as "[Hactive Hfactive]".
      { by apply auth_update, option_local_update, exclusive_local_update. }
      iApply ("Hcl" with "[-]").
      iNext. iExists _,∅,_. rewrite left_id. iFrame. iFrame "Hactive".
      iSplit; auto.
      rewrite /all_props bi.big_sepM_insert. iFrame.
      apply (not_elem_of_dom _ s (D:=gset prop_id)).
      apply is_fresh.
  Qed.

  Lemma newlock_cancel_spec :
    {{{ True }}} newlock #() {{{ lk γ, RET lk; is_flock γ lk ∗ unflocked γ 1 }}}.
  Proof.
    iIntros (Φ) "_ HΦ". rewrite -wp_fupd.
    iMod (own_alloc (● (Excl' Unlocked) ⋅ ◯ (Excl' Unlocked)))
      as (γ_state) "[Hauth Hfrag]"; first done.
    iMod (own_alloc (● to_props_map ∅)) as (γ_props) "Hprops".
    { apply auth_valid_discrete. simpl.
      split; last done. apply ucmra_unit_least. }
    iMod (own_alloc ((● Excl' ∅) ⋅ (◯ Excl' ∅))) as (γ_ac_props) "[Hacprops1 Hacprops2]".
    { apply auth_valid_discrete. simpl.
      split; last done. by rewrite left_id. }
    iApply (newlock_spec (flockN.@"lock") (own γ_state (◯ (Excl' Unlocked))) with "Hfrag").
    iNext. iIntros (lk γ_lock) "#Hlock".
    iMod (cinv_alloc_strong ∅ _ (flockN.@"inv")
       (fun γ => flock_inv (Build_flock_name γ_lock γ γ_state γ_props γ_ac_props)) with "[-HΦ]") as (γ_cinv) "(_ & #Hcinv & Hown)".
    { iNext. rewrite /flock_inv /=. iIntros (γ _).
      iExists Unlocked, ∅, ∅. rewrite left_id. iFrame.
      by rewrite /all_props !bi.big_sepM_empty. }
    pose (γ := (Build_flock_name γ_lock γ_cinv γ_state γ_props γ_ac_props)).
    iModIntro. iApply ("HΦ" $! lk γ with "[-]"). iFrame "Hown".
    rewrite /is_flock. by iFrame "Hlock".
  Qed.

  Lemma acquire_cancel_spec γ π lk s R :
    {{{ is_flock γ lk ∗ unflocked γ π  ∗ flock_res γ s R }}}
      acquire lk
    {{{ f, RET #(); ▷ R ∗ (▷ R -∗ flocked γ π f) }}}.
  Proof.
    iIntros (Φ) "(Hl & Hunfl & #HRres) HΦ". rewrite -wp_fupd.
    rewrite /is_flock. iDestruct "Hl" as "(#Hcinv & #Hlk)".
    iApply (acquire_spec with "Hlk"). iNext.
    iIntros "[Hlocked Hunlk]".
    iMod (cinv_open with "Hcinv Hunfl") as "(Hstate & Hunfl & Hcl)"; first done.
    rewrite {2}/flock_inv.
    iDestruct "Hstate" as ([q|] fp fa) "(>Hstate & >Hauth & >Hfactive & Hrest)".
    - iDestruct "Hrest" as "(>Hcown & >Hlocked2 & Hfp)".
      iExFalso. iApply (locked_exclusive with "Hlocked Hlocked2").
    - iDestruct "Hrest" as "(>Hactive & Hfa & >%)".
      simplify_eq/=. rewrite left_id.
      iMod (own_update_2 with "Hstate Hunlk") as "Hstate".
      { apply (auth_update _ _ (Excl' (Locked π)) (Excl' (Locked π))).
        apply option_local_update.
        by apply exclusive_local_update. }
      iDestruct "Hstate" as "[Hstate Hflkd]".
      iDestruct "Hunfl" as "[Hunfl Hcown]".
      iAssert (∃ R', ⌜R ≡ R' ∧ fa !! s = Some R'⌝)%I
        with "[HRres Hauth]" as %Hfoo.
      { iDestruct (own_valid_2 with "Hauth HRres")
          as %[Hfoo _]%auth_valid_discrete.
        iPureIntro. revert Hfoo.
        rewrite /= left_id singleton_included=> -[R' [Hf HR]].
        revert Hf HR.
        rewrite /to_props_map lookup_fmap fmap_Some_equiv=>-[R'' [/= Hf ->]].
        intros [?%to_agree_inj | ?%to_agree_included]%Some_included;
          simplify_eq/=; exists R''; eauto. }
      destruct Hfoo as [R' [HReq Hf]].
      rewrite /all_props {1}(bi.big_sepM_lookup_acc _ fa s R'); last done.
      iDestruct "Hfa" as "[HR' Hfa]".
      iMod ("Hcl" with "[Hstate Hcown Hlocked Hauth Hfactive]").
      { iNext. iExists (Locked π),∅,fa. rewrite left_id. iFrame.
        by rewrite /all_props bi.big_sepM_empty. }
      iModIntro. 
      iApply ("HΦ" $! fa). rewrite -HReq. iFrame.
      by iApply bi.later_wand.
  Qed.

  Lemma release_cancel_spec γ π lk f :
    {{{ is_flock γ lk ∗ flocked γ π f }}}
      release lk
    {{{ RET #(); unflocked γ π }}}.
  Proof.
    iIntros (Φ) "(Hl & (Hstate' & Hflkd & Hactive & Hf)) HΦ". rewrite -fupd_wp.
    rewrite /is_flock. iDestruct "Hl" as "(#Hcinv & #Hlk)".
    iMod (cinv_open with "Hcinv Hflkd") as "(Hstate & Hunfl & Hcl)"; first done.
    rewrite {2}/flock_inv.
    iDestruct "Hstate" as ([q|] fp fa) "(>Hstate & >Hauth & >Hfactive & Hrest)"; last first.
    - iDestruct "Hrest" as "(>Hactive2 & Hfa & >%)".
      iExFalso. iDestruct (own_valid_2 with "Hstate Hstate'") as %Hfoo.
      iPureIntro. revert Hfoo.
      rewrite auth_valid_discrete /= left_id.
      intros [Hexcl _].
      apply Some_included_exclusive in Hexcl; try (done || apply _).
      simplify_eq.
    - iDestruct "Hrest" as "(>Hcown & >Hlocked & Hfp)".
      iAssert (⌜q = π⌝)%I with "[Hstate Hstate']" as  %->.
      { iDestruct (own_valid_2 with "Hstate Hstate'") as %Hfoo.
        iPureIntro. revert Hfoo.
        rewrite auth_valid_discrete /= left_id.
        intros [Hexcl _].
        apply Some_included_exclusive in Hexcl; try (done || apply _).
        by simplify_eq. }
      iMod (own_update_2 with "Hstate Hstate'") as "Hstate".
      { apply (auth_update _ _ (Excl' Unlocked) (Excl' Unlocked)).
        apply option_local_update.
        by apply exclusive_local_update. }
      iDestruct "Hstate" as "[Hstate Hstate']".
      iAssert (⌜(f : gmapC prop_id (iProp Σ)) ≡ fa⌝%I) with "[Hactive Hfactive]" as %Hfoo.
      { iDestruct (own_valid_2 with "Hactive Hfactive") as %Hfoo.
        iPureIntro. apply auth_valid_discrete_2 in Hfoo.
        destruct Hfoo as [Hfoo _].
        apply Some_included_exclusive in Hfoo; [|apply _|done].
        by apply Excl_inj. }
      rewrite {1}Hfoo.
      iMod (own_update_2 _ _ _ (● Excl' (fp ∪ fa) ⋅ ◯ Excl' (fp ∪ fa))
              with "Hactive Hfactive") as "[Hactive Hfactive]".
      { by apply auth_update, option_local_update, exclusive_local_update. }
      iMod ("Hcl" with "[Hstate Hauth Hfp Hf Hactive Hfactive]").
      { iNext. iExists Unlocked,∅,(fp∪fa). rewrite left_id. iFrame.
        rewrite Hfoo. iSplit; eauto.
        iApply all_props_union; iFrame. }
      iApply (release_spec with "[$Hlk $Hlocked $Hstate']").
      iModIntro. iNext. iIntros "_". iApply "HΦ".
      rewrite /unflocked. by iSplitL "Hcown".
  Qed.

  Lemma cancel_lock γ lk s R :
    is_flock γ lk -∗ flock_res γ s R -∗ unflocked γ 1 ={⊤}=∗ ▷ R.
  Proof.
    rewrite /is_flock /unflocked.
    iDestruct 1 as "(#Hcinv & #Hislock)". iIntros "#Hres Hcown".
    rewrite /flock_res.
    iMod (cinv_cancel_vs _ R with "Hcinv [] Hcown") as "$"; eauto.
    iIntros "[Hcown Hstate]".
    iDestruct "Hstate" as ([q|] fp fa) "(Hstate & >Hauth & Hfactive & Hrest)".
    - iDestruct "Hrest" as "(>Hcown2 & >Hlocked & Hfp)".
      iDestruct (cinv_own_1_l with "Hcown Hcown2") as %[].
    - iDestruct "Hrest" as "(>Hactive & Hfa & >%)".
      simplify_eq/=. iFrame. rewrite left_id.
      iAssert (∃ R', ⌜R ≡ R' ∧ fa !! s = Some R'⌝)%I
        with "[Hres Hauth]" as %Hfoo.
      { iDestruct (own_valid_2 with "Hauth Hres")
          as %[Hfoo _]%auth_valid_discrete.
        iPureIntro. revert Hfoo.
        rewrite /= left_id singleton_included=> -[R' [Hf HR]].
        revert Hf HR.
        rewrite /to_props_map lookup_fmap fmap_Some_equiv=>-[R'' [/= Hf ->]].
        intros [?%to_agree_inj | ?%to_agree_included]%Some_included;
          simplify_eq/=; exists R''; eauto. }
      destruct Hfoo as [R' [HReq Hf]].
      rewrite /all_props {1}(bi.big_sepM_lookup_acc _ fa s R'); last done.
      iDestruct "Hfa" as "[HR' Hfa]".
      assert ((▷ R)%I ≡ (▷ R')%I) as ->. by apply bi.later_proper.
      by iFrame.
  Qed.
End flock.