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

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

Instance lockstate_inhabited : Inhabited lockstate := populate Unlocked.

Record flock_name := {
  flock_lock_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.
Canonical Structure gnameC := leibnizC gname.

Class flockG Σ :=
  FlockG {
    flock_stateG :> inG Σ (authR (optionUR (exclR lockstateC)));
    flock_lockG  :> lockG Σ;
    flock_savedProp :> savedPropG Σ;
    flock_tokens :> inG Σ fracR;
    flock_props_active :> inG Σ (authR (optionUR (exclR (gmapC prop_id (prodC fracC (prodC gnameC gnameC))))));
    flock_props :> inG Σ (authR (gmapUR prop_id (prodR fracR (agreeR (prodC gnameC gnameC)))))
  }.

Definition flockΣ : gFunctors :=
  #[GFunctor (authR (optionUR (exclR lockstateC)))
   ;lockΣ
   ;savedPropΣ
   ;GFunctor fracR
   ;GFunctor (authR (optionUR (exclR (gmapC prop_id (prodC fracC (prodC gnameC gnameC))))))
   ;GFunctor (authR (gmapUR prop_id (prodR fracR (agreeR (prodC gnameC gnameC)))))%CF].

Instance subG_flockΣ Σ : subG flockΣ Σ → flockG Σ.
Proof. solve_inG. Qed.

Section flock.
  Context `{heapG Σ, flockG Σ}.
  Variable N : namespace.
  Definition flockN := N.@"flock".

  Definition to_props_map (f : gmap prop_id (gname * gname))
    : gmapUR prop_id (prodR fracR (agreeR (prodC gnameC gnameC))) :=
    (λ x, (1%Qp, to_agree (x.1, x.2))) <$> f.

  Lemma to_props_map_singleton_included fp i q ρ:
    {[i := (q, to_agree ρ)]} ≼ to_props_map fp → fp !! i = Some ρ.
  Proof.
    rewrite singleton_included=> -[[q' av] []].
    rewrite /to_props_map lookup_fmap fmap_Some_equiv => -[v' [Hi [/= -> ->]]].
    move=> /Some_pair_included_total_2 [_] /to_agree_included /leibniz_equiv_iff -> //.
    rewrite Hi. by destruct v'.
  Qed.

  Definition from_active (f : gmap prop_id (frac * (gname * gname)))
    : gmap prop_id (gname * gname) := fmap snd f.

  Lemma from_active_empty : from_active ∅ = ∅.
  Proof. by rewrite /from_active fmap_empty. Qed.

  Definition all_props (f : gmap prop_id (gname*gname)) : iProp Σ :=
    ([∗ map] i ↦ ρ ∈ f, ∃ R, saved_prop_own ρ.1 R ∗ R)%I.

  Definition all_tokens (f : gmap prop_id (frac * (gname*gname))) : iProp Σ :=
    ([∗ map] i ↦ ρ ∈ f, own ρ.2.2 ρ.1)%I.

  Definition flock_inv (γ : flock_name) : iProp Σ :=
    (∃ (s : lockstate)
       (fp : gmap prop_id (gname * gname))
       (fa : gmap prop_id (frac * (gname * gname))),
       (** fa -- active propositions, fp -- pending propositions *)
       ⌜fp ##ₘ from_active fa⌝ ∗
       own (flock_state_name γ) (● (Excl' s)) ∗
       own (flock_props_name γ) (● to_props_map (fp ∪ from_active fa)) ∗
       own (flock_props_active_name γ) (◯ Excl' fa) ∗
       all_props fp ∗
       match s with
       | Locked =>
         locked (flock_lock_name γ) ∗
         all_tokens fa
       | Unlocked => own (flock_props_active_name γ) (● Excl' ∅)
       end)%I.

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

  Definition flocked
    (γ : flock_name) (f : gmap prop_id (frac * (gname * gname))) : iProp Σ :=
    (own (flock_state_name γ) (◯ (Excl' Locked))
     ∗ own (flock_props_active_name γ) (● Excl' f)
     ∗ all_props (from_active f))%I.

  Definition flock_res (γ : flock_name) (s : prop_id) (R : iProp Σ) (π : Qp) : iProp Σ :=
    (∃ ρ, saved_prop_own ρ.1 R ∗ own ρ.2 π
     ∗ own (flock_props_name γ) (◯ {[s := (π, to_agree ρ)]}))%I.

  Lemma flock_res_op (γ : flock_name) (s : prop_id) (R : iProp Σ) (π1 π2 : frac) :
    flock_res γ s R (π1+π2) ⊣⊢ flock_res γ s R π1 ∗ flock_res γ s R π2.
  Proof.
    rewrite /flock_res. iSplit.
    - iDestruct 1 as (ρ) "(#Hsaved & Hρ & Hs)".
      iDestruct "Hρ" as "[Hρ1 Hρ2]".
      iDestruct "Hs" as "[Hs1 Hs2]".
      iSplitL "Hρ1 Hs1"; iExists _; by iFrame.
    - iIntros "[H1 H2]".
      iDestruct "H1" as (ρ) "(#Hsaved & Hρ1 & Hs1)".
      iDestruct "H2" as (ρ') "(_ & Hρ2 & Hs2)".
      iCombine "Hs1 Hs2" as "Hs".
      iDestruct (own_valid with "Hs")
        as %Hfoo%auth_valid_discrete.
      simpl in Hfoo. apply singleton_valid in Hfoo.
      destruct Hfoo as [_ Hfoo%agree_op_inv'].
      fold_leibniz. rewrite -!Hfoo.
      iCombine "Hρ1 Hρ2" as "Hρ".
      rewrite !frac_op' agree_idemp.
      iExists ρ. by iFrame.
  Qed.

  Global Instance flock_res_fractional γ s R : Fractional (flock_res γ s R).
  Proof. intros p q. apply flock_res_op. Qed.

  Global Instance flock_res_as_fractional γ s R π :
    AsFractional (flock_res γ s R π) (flock_res γ s R) π.
  Proof. split. done. apply _. Qed.

  Lemma flock_res_single_alloc γ lk R E :
    ↑flockN ⊆ E →
    is_flock γ lk -∗ ▷ R
    ={E}=∗ ∃ s, flock_res γ s R 1.
  Proof.
    iIntros (?) "Hl HR". rewrite /is_flock. iDestruct "Hl" as "(#Hcinv & #Hlk)".
    iMod (saved_prop_alloc R) as (ρ1) "#Hρ1".
    iMod (own_alloc 1%Qp) as (ρ2) "Hρ2"; first done.
    iInv (flockN.@"inv") as (s fp fa) "(>% & >Hstate & >Hauth & >Hfactive & Hfp & Hrest)" "Hcl".
    pose (i := (fresh ((dom (gset prop_id) (fp ∪ from_active fa))))).
    assert (i ∉ (dom (gset prop_id) (fp ∪ from_active fa))) as Hs.
    { apply is_fresh. }
    iMod (own_update  with "Hauth") as "Hauth".
    { apply (auth_update_alloc _ (to_props_map (<[i := (ρ1,ρ2)]> fp ∪ from_active fa))
                               {[ i := (1%Qp, to_agree (ρ1, ρ2)) ]}).
      rewrite -insert_union_l.
      rewrite /to_props_map /= fmap_insert.
      apply alloc_local_update; last done.
      apply (not_elem_of_dom (to_props_map (fp ∪ from_active fa)) i (D:=gset prop_id)).
      by rewrite /to_props_map dom_fmap. }
      iDestruct "Hauth" as "[Hauth Hres]".
      iExists i, (ρ1,ρ2). iFrame "Hres Hρ1 Hρ2".
      iApply ("Hcl" with "[-]").
      iNext. iExists _,_,_. iFrame. iFrame "Hrest".
      rewrite /all_props big_sepM_insert; last first.
      + apply (not_elem_of_dom _ i (D:=gset prop_id)).
          revert Hs. rewrite dom_union_L not_elem_of_union. set_solver. 
      + iFrame "Hfp". iSplitR; last by eauto.
        iPureIntro. admit. (* TODO: map disjoint *)
  Admitted.

  Lemma flock_res_single_dealloc γ lk i R E :
    ↑flockN ⊆ E →
    is_flock γ lk -∗ flock_res γ i R 1
    ={E}=∗ ∃ R', ▷ R' ∗ ▷ ▷ (R ≡ R').
  Proof.
    iIntros (?) "Hl HR". rewrite /is_flock. iDestruct "Hl" as "(#Hcinv & #Hlk)".
    iInv (flockN.@"inv") as ([|] fp fa) "(>% & >Hstate & >Hauth & >Hfactive & Hfp & Hrest)" "Hcl".
    - iDestruct "Hrest" as ">[Hlocked Htokens]".
      iDestruct "HR" as ([ρ1 ρ2]) "(#Hρ1 & Hρ2 & HR)". simpl.
      iDestruct (own_valid_2 with "Hauth HR")
        as %[Hfoo%to_props_map_singleton_included _]%auth_valid_discrete_2.     
      iAssert (⌜fa !! i = None⌝)%I with "[-]" as %Hbar.
      { remember (fa !! i) as fai. destruct fai as [[π [ρ'1 ρ'2]]|]; last done.
        iExFalso.
        assert (fp !! i = None) as Hbar.
        { apply lookup_union_Some_raw in Hfoo.
          destruct Hfoo as [Hfoo | [? ?]]; last done.
          specialize (H2 i). revert H2. rewrite Hfoo.
          rewrite /from_active lookup_fmap -Heqfai /=. done. }
        apply lookup_union_Some in Hfoo; last done.
        destruct Hfoo as [Hfoo | Hfa].
        - exfalso. rewrite Hfoo in Hbar. inversion Hbar.
        - revert Hfa. rewrite /from_active lookup_fmap -Heqfai /= =>Hfa.
          rewrite /all_tokens (big_sepM_lookup _ fa i); last done.
          simplify_eq/=.
          iDestruct (own_valid with "Htokens") as %Hbaz1%frac_valid'.
          iDestruct (own_valid_2 with "Hρ2 Htokens") as %Hbaz2.
          rewrite frac_op' in Hbaz2. exfalso. apply Hbaz2. admit. (* TODO: Qp *) }
      assert (fp !! i = Some (ρ1, ρ2)) as Hbaz.
      { apply lookup_union_Some in Hfoo; last done.
        destruct Hfoo as [? | Hfoo]; first done.
        exfalso. revert Hfoo. rewrite  /from_active lookup_fmap Hbar. done. }

      iMod (own_update_2 with "Hauth HR") as "Hauth".
      { apply auth_update_dealloc, (delete_singleton_local_update _ i _). }

      rewrite /all_props (big_sepM_delete _ fp i (ρ1,ρ2)); last done.
      iDestruct "Hfp" as "[HR' Hfp]".
      iDestruct "HR'" as (R') "[Hsaved HR']".
      iDestruct (saved_prop_agree with "Hρ1 Hsaved") as "Hfoo".
      iMod ("Hcl" with "[-HR' Hfoo Hρ2]") as "_".
      { iNext. iExists Locked, (delete i fp),fa. iFrame.
        iSplit. admit. (* TODO: map_disjoint *)
        rewrite /to_props_map -fmap_delete.
        rewrite delete_union (delete_notin (from_active fa)); first iFrame.
        by rewrite /from_active lookup_fmap Hbar. }
      iModIntro. iExists R'. iFrame.
    - iDestruct "Hrest" as ">Haactive".
      iAssert (⌜fa = ∅⌝)%I with "[-]" as %->.
      { iDestruct (own_valid_2 with "Haactive Hfactive")
          as %[Hfoo%Excl_included _]%auth_valid_discrete_2.
        iPureIntro. by unfold_leibniz. }
      rewrite from_active_empty right_id.
      iDestruct "HR" as ([ρ1 ρ2]) "(#Hρ1 & Hρ2 & HR)". simpl.
      iDestruct (own_valid_2 with "Hauth HR")
        as %[Hfoo%to_props_map_singleton_included _]%auth_valid_discrete_2.    
      iMod (own_update_2 with "Hauth HR") as "Hauth".
      { apply auth_update_dealloc, (delete_singleton_local_update _ i _). }
      rewrite /all_props (big_sepM_delete _ fp i (ρ1,ρ2)); last done.
      iDestruct "Hfp" as "[HR' Hfp]".
      iDestruct "HR'" as (R') "[Hsaved HR']".
      iDestruct (saved_prop_agree with "Hρ1 Hsaved") as "Hfoo".
      iMod ("Hcl" with "[-HR' Hfoo Hρ2]") as "_".
      { iNext. iExists Unlocked, (delete i fp),∅. iFrame.
        rewrite !from_active_empty right_id.
        rewrite /to_props_map fmap_delete. iFrame "Hauth".
        iPureIntro. apply map_disjoint_empty_r. }
      iModIntro. iExists R'. iFrame.
  Admitted.

  Lemma newlock_cancel_spec :
    {{{ True }}} newlock #() {{{ lk γ, RET lk; is_flock γ lk }}}.
  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".
    pose (γ := (Build_flock_name γ_lock γ_state γ_props γ_ac_props)).
    iMod (inv_alloc (flockN.@"inv") _ (flock_inv γ) with "[-HΦ]") as "#Hinv".
    { iNext. iExists Unlocked, ∅, ∅. rewrite from_active_empty right_id. iFrame.
      iSplit; eauto. by rewrite /all_props big_sepM_empty. }
    iModIntro. iApply ("HΦ" $! lk γ with "[-]").
    rewrite /is_flock. by iFrame "Hlock".
  Qed.

  Lemma acquire_cancel_spec γ π lk i R :
    {{{ is_flock γ lk ∗ flock_res γ i R π }}}
      acquire lk
    {{{ RET #(); ∃ R', R' ∗ ▷ (R ≡ R') ∗ (▷ R' ={⊤}=∗ flocked γ ∅ ∗ flock_res γ i R π) }}}.
  Proof.
    iIntros (Φ) "(Hl & HRres) HΦ".
    rewrite /is_flock. iDestruct "Hl" as "(#Hinv & #Hlk)".
    iApply (acquire_spec with "Hlk").
    iNext. iIntros "[Hlocked Hunlk]".
    iInv (flockN.@"inv") as ([|] fp fa) "(>% & >Hstate & >Hauth & >Hfactive & Hfp & Hrest)" "Hcl".
    - iDestruct "Hrest" as "(>Hlocked2 & Htokens)".
      iExFalso. iApply (locked_exclusive with "Hlocked Hlocked2").
    - iDestruct "Hrest" as ">Haactive".
      iAssert (⌜fa = ∅⌝)%I with "[-]" as %->.
      { iDestruct (own_valid_2 with "Haactive Hfactive")
          as %[Hfoo%Excl_included _]%auth_valid_discrete_2.
        iPureIntro. by unfold_leibniz. }
      rewrite from_active_empty right_id.
      iDestruct "HRres" as ([ρ1 ρ2]) "(#Hρ1 & Hρ2 & HR)". simpl.

      (* Unlocked ~~> Locked *)
      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]".

      (* (i,ρ) ∈ fp *)
      iDestruct (own_valid_2 with "Hauth HR")
        as %[Hfoo%to_props_map_singleton_included _]%auth_valid_discrete_2.    

      rewrite /all_props (big_sepM_delete _ fp i (ρ1,ρ2)); last done.
      iDestruct "Hfp" as "[HR' Hfp]".
      iDestruct "HR'" as (R') "[Hsaved HR']".
      iDestruct (saved_prop_agree with "Hρ1 Hsaved") as "#Hfoo".
      
      (* move (i,ρ) to the set of active propositions *)
      iMod (own_update_2 with "Haactive Hfactive") as "[Haactive Hfactive]".
      { apply (auth_update _ _ (Excl' {[ i := (π, (ρ1, ρ2)) ]})
                               (Excl' {[ i := (π, (ρ1, ρ2)) ]})).
        apply option_local_update.
        by apply exclusive_local_update. }
      
      iMod ("Hcl" with "[-HΦ HR' Hfoo Haactive Hflkd HR]") as "_".
      { iNext. iExists Locked,(delete i fp),{[i := (π, (ρ1, ρ2))]}.
        iFrame. iSplitR. admit. (* TODO: map_disjoint *)
        iSplitL "Hauth".
        - rewrite /from_active map_fmap_singleton /=.
          assert (fp = (delete i fp ∪ {[i := (ρ1, ρ2)]})) as <-.
          { apply map_eq=>j. destruct (decide (i = j)) as [->|?].
            - remember (fp !! j) as fpj. destruct fpj.
              + symmetry. apply lookup_union_Some.
                admit. (* TODO map disjoint *)
                right. simplify_eq/=. apply lookup_singleton.
              + symmetry. apply lookup_union_None. naive_solver.
            - remember (fp !! j) as fpj. destruct fpj.
              + symmetry. apply lookup_union_Some.
                admit. (* TODO map disjoint *)
                left. simplify_eq/=. rewrite lookup_delete_ne; eauto.
              + symmetry. apply lookup_union_None. split; eauto.
                * rewrite lookup_delete_ne; eauto.
                * rewrite lookup_singleton_ne; eauto. }
          iFrame.
        - rewrite /all_tokens big_sepM_singleton //. }
      iModIntro. (* iApply ("HΦ" $! {[ i := (π, (ρ1, ρ2)) ]}). *)
      iApply "HΦ".
      iNext. iExists R'. iFrame "HR' Hfoo". iIntros "HR'".
      clear Hfoo H1 fp. rewrite /flocked /flock_res.
      iInv (flockN.@"inv") as ([|] fp fa) "(>% & >Hstate & >Hauth & >Hfactive & Hfp & Hrest)" "Hcl"; last first.
      + iDestruct (own_valid_2 with "Hstate Hflkd")
          as %[Hfoo%Excl_included _]%auth_valid_discrete_2.
        fold_leibniz. inversion Hfoo.
      + iDestruct "Hrest" as ">[Hlocked Htokens]".
        
        iDestruct (own_valid_2 with "Haactive Hfactive")
          as %[Hfoo%Excl_included _]%auth_valid_discrete_2.
        fold_leibniz. simplify_eq/=.

        iAssert (⌜fp !! i = None⌝)%I with "[-]" as %Hbar.
        { remember (fp !! i) as fpi. destruct fpi as [[ρ'1 ρ'2]|]; last done.
          exfalso. specialize (H1 i). revert H1. rewrite -Heqfpi.
          rewrite /from_active lookup_fmap lookup_singleton. done. }

        iMod (own_update_2 with "Haactive Hfactive") as "[Haactive Hfactive]".
        { apply (auth_update _ _ (Excl' ∅)
                                 (Excl' ∅)).
          apply option_local_update. by apply exclusive_local_update. }
        
        iMod ("Hcl" with "[HR' Hlocked Hstate Hauth Hfactive Hfp]") as "_".
        { iNext. iExists Locked,(<[i := (ρ1,ρ2)]>fp),∅.
          iSplitR.
          { rewrite from_active_empty. iPureIntro. apply map_disjoint_empty_r. }
          iFrame "Hstate Hfactive Hlocked".
          iSplitL "Hauth".
          - rewrite from_active_empty right_id /from_active map_fmap_singleton.
            rewrite insert_union_singleton_r; done.
          - rewrite /all_tokens /all_props big_sepM_empty. iSplitL; last done.
            rewrite big_sepM_insert; last done.
            iFrame "Hfp". iExists R'. iFrame. simpl.
            iRewrite "Hfoo" in "Hρ1". by iFrame. }
        iModIntro. iFrame. rewrite /all_props from_active_empty big_sepM_empty.
        iSplitR; first done.
        rewrite /all_tokens big_sepM_singleton. simpl.
        iExists (ρ1, ρ2). by iFrame.
  Admitted.

  Lemma release_cancel_spec γ lk :
    {{{ is_flock γ lk ∗ flocked γ ∅ }}}
      release lk
    {{{ RET #(); True }}}.
  Proof.
    iIntros (Φ) "(#Hl & (Hflkd & Haactive & Hfa)) HΦ". rewrite -fupd_wp.
    rewrite /is_flock. iDestruct "Hl" as "(#Hinv & #Hlk)".
    iInv (flockN.@"inv") as ([|] fp fa) "(>% & >Hstate & >Hauth & >Hfactive & Hfp & Hrest)" "Hcl"; last first.
    - iDestruct (own_valid_2 with "Hstate Hflkd")
          as %[Hfoo%Excl_included _]%auth_valid_discrete_2.
      fold_leibniz. inversion Hfoo.
    - iDestruct "Hrest" as ">[Hlocked Htokens]".        
      iDestruct (own_valid_2 with "Haactive Hfactive")
        as %[Hfoo%Excl_included _]%auth_valid_discrete_2.
      fold_leibniz. simplify_eq/=.

      (* Locked ~~> Unlocked *)
      iMod (own_update_2 with "Hstate Hflkd") as "Hstate".
      { apply (auth_update _ _ (Excl' Unlocked) (Excl' Unlocked)).
        apply option_local_update.
        by apply exclusive_local_update. }
      iDestruct "Hstate" as "[Hstate Hunflkd]".

      iMod ("Hcl" with "[Hstate Hauth Haactive Hfactive Hfp]") as "_".
      { iNext. iExists Unlocked,fp,∅.
        iSplitR; eauto. iFrame. }
      iApply (release_spec with "[$Hlk $Hlocked $Hunflkd]").
      done.
  Qed.

End flock.