(** 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.
From iris.algebra Require Import auth excl frac.
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;
}.

Class flockG Σ :=
  FlockG {
    flock_stateG :> inG Σ (authR (optionUR (exclR lockstateC)));
    flock_cinvG :> cinvG Σ;
    flock_lockG :> lockG Σ
  }.

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 flock_inv (γ : flock_name) (R : iProp Σ) : iProp Σ :=
    (∃ s : lockstate,
       own (flock_state_name γ) (● (Excl' s)) ∗
       match s with
       | Locked q =>
         cinv_own (flock_cinv_name γ) (q/2) ∗
         locked (flock_lock_name γ)
       | Unlocked => R
       end)%I.

  Definition is_flock (γ : flock_name) (lk : val) (R : iProp Σ) : iProp Σ :=
    (cinv (N .@ "inv") (flock_cinv_name γ) (flock_inv γ R) ∗
     is_lock (N .@ "lock") (flock_lock_name γ) lk
         (own (flock_state_name γ) (◯ (Excl' Unlocked))))%I.

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

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

  Definition flocked (γ : flock_name) (q : frac) : iProp Σ :=
    (own (flock_state_name γ) (◯ (Excl' (Locked q)))
     ∗ cinv_own (flock_cinv_name γ) (q/2))%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 newlock_cancel_spec (R : iProp Σ):
    {{{ R }}} newlock #() {{{ lk γ, RET lk; is_flock γ lk R ∗ unflocked γ 1 }}}.
  Proof.
    iIntros (Φ) "HR HΦ". rewrite -wp_fupd.
    iMod (own_alloc (● (Excl' Unlocked) ⋅ ◯ (Excl' Unlocked)))
      as (γ_state) "[Hauth Hfrag]"; first done.
    iApply (newlock_spec (N.@"lock") (own γ_state (◯ (Excl' Unlocked))) with "Hfrag").
    iNext. iIntros (lk γ_lock) "#Hlock".
    iMod (cinv_alloc_strong ∅ _ (N.@"inv")
       (fun γ => flock_inv (Build_flock_name γ_lock γ γ_state) R) with "[-HΦ]") as (γ_cinv) "(_ & #Hcinv & Hown)".
    { iNext. rewrite /flock_inv /=. iIntros (γ _).
      iExists Unlocked. by iFrame. }
    pose (γ := (Build_flock_name γ_lock γ_cinv γ_state)).
    iModIntro. iApply ("HΦ" $! lk γ). iFrame "Hown".
    rewrite /is_flock. by iFrame "Hlock".
  Qed.

  Lemma acquire_cancel_spec γ π lk (R : iProp Σ) :
    {{{ is_flock γ lk R ∗ unflocked γ π }}}
      acquire lk
    {{{ RET #(); flocked γ π ∗ ▷ R }}}.
  Proof.
    iIntros (Φ) "[Hl Hunfl] 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|]) "Hstate".
    - iDestruct "Hstate" as ">(Hstate & Hcown & Hlocked2)".
      iExFalso. iApply (locked_exclusive with "Hlocked Hlocked2").
    - iDestruct "Hstate" as "[>Hstate HR]".
      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]".
      iMod ("Hcl" with "[Hstate Hcown Hlocked]").
      { iNext. iExists (Locked π). iFrame. }
      iApply "HΦ". by iFrame.
  Qed.

  Lemma release_cancel_spec γ π lk R :
    {{{ is_flock γ lk R ∗ flocked γ π ∗ ▷ R }}}
      release lk
    {{{ RET #(); unflocked γ π }}}.
  Proof.
    iIntros (Φ) "(Hl & (Hstate' & Hflkd) & HR) 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|]) "Hstate"; last first.
    - iDestruct "Hstate" as "[>Hstate HR']".
      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 "Hstate" as ">(Hstate & Hcown & Hlocked)".
      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']".
      iMod ("Hcl" with "[Hstate HR]").
      { iNext. iExists Unlocked. iFrame. }
      iApply (release_spec with "[$Hlk $Hlocked $Hstate']").
      iModIntro. iNext. iIntros "_". iApply "HΦ".
      by iSplitL "Hcown".
  Qed.

  Lemma cancel_lock γ lk R :
    is_flock γ lk R -∗ unflocked γ 1 ={⊤}=∗ ▷ R.
  Proof.
    rewrite /is_flock /unflocked.
    iDestruct 1 as "(#Hcinv & #Hislock)". iIntros "Hcown".
    iMod (cinv_cancel_vs _ R with "Hcinv [] Hcown") as "$"; eauto.
    iIntros "[Hcown Hstate]".
    iDestruct "Hstate" as ([q|]) "(Hstate & HR)".
    - iDestruct "HR" as ">(Hcown2 & Hlocked)".
      iDestruct (cinv_own_1_l with "Hcown Hcown2") as %[].
    - by iFrame.
  Qed.
End flock.