From iris.heap_lang Require Export proofmode notation.
From iris_c.lib Require Export locking_heap.
From iris.heap_lang Require Import adequacy spin_lock assert par.
From iris.algebra Require Import frac.
From iris_c.lib Require Import mset flock.

(* M A := ref (list loc) → Mutex → A *)

(* A → M A *)
Definition c_ret : val := λ: "a" <> <>, "a".

(* (A → M B) → M A → M B *)
Definition c_bind : val := λ: "x" "f" "env" "l",
  let: "a" := "x" "env" "l" in
  "f" "a" "env" "l".

Notation "x ←ᶜ y ;;ᶜ z" :=
  (c_bind y (λ: x, z))%E
  (at level 100, y at next level, z at level 200, right associativity) : expr_scope.
Notation "y ;;ᶜ z" := (c_bind y (λ: <>, z))%E
  (at level 100, z at level 200, right associativity) : expr_scope.

(* M A → A *)
Definition c_run : val := λ: "x",
  let: "env" := mset_create #() in
  let: "l" := newlock #() in
  "x" "env" "l".

(* M A → M A *)
Definition c_atomic : val := λ: "x" "env" "l",
  acquire "l";;
  let: "k" := newlock #() in
  let: "a" := "x" #() "env" "k" in
  release "l";;
  "a".

(* (ref (list loc) → A) → M A *)
Definition c_atomic_env : val := λ: "f" "env" "l",
  acquire "l";;
  let: "a" := "f" "env" in
  release "l";;
  "a".

(* M A → M B → M (A * B) *)
Definition c_par : val := λ: "x" "y" "env" "l",
  "x" "env" "l" ||| "y" "env" "l".
Notation "e1 |||ᶜ e2" := (c_par e1 e2)%E (at level 50) : expr_scope.

Definition cmonadN := nroot .@ "amonad".

Class cmonadG (Σ : gFunctors) := AMonadG {
  aheapG :> heapG Σ;
  aflockG :> flockG Σ;
  alocking_heapG :> locking_heapG Σ;
  aspawnG :> spawnG Σ
}.

Section cwp.
  Context `{cmonadG Σ}.

  Definition env_inv (env : val) : iProp Σ :=
    (∃ (X : gset val) (σ : gmap cloc (lvl * val)),
      ⌜ ∀ v, v ∈ X → ∃ cl, cloc_of_val (SOMEV v) = Some cl ∧ cl ∈ locked_locs σ ⌝ ∧
      is_mset env X ∗
      full_locking_heap σ)%I.

  Definition flock_resources (γ : flock_name) (I : gmap prop_id lock_res) :=
    ([∗ map] i ↦ X ∈ I, flock_res cmonadN γ i X)%I.

  (** DF: The outer `WP` here is needed to be able to perform some reductions inside a heap_lang context.
      Without this, the `cwp_cwp` rule is not provable.

      My intuitive explanation: we want to preform some reductions to `e` until it is actually a value that is a monadic computation.
      In some sense it is a form of CPSing on a logical level.
      But I still cannot precisely state why is it needed.
 *)
  Definition cwp_def (e : expr)
      (R : iProp Σ) (Φ : val → iProp Σ) : iProp Σ :=
    WP e {{ ev,
      ∀ (γ : flock_name) (env : val) (l : val) (I : gmap prop_id lock_res),
        is_flock cmonadN γ l -∗
        flock_resources γ I -∗
        ([∗ map] X ∈ I, res_prop X) ≡ (env_inv env ∗ R) -∗
        WP ev env l {{ v, Φ v ∗ flock_resources γ I }}
    }}%I.
  Definition cwp_aux : seal (@cwp_def). by eexists. Qed.
  Definition cwp := unseal cwp_aux.
  Definition cwp_eq : @cwp = @cwp_def := seal_eq cwp_aux.
End cwp.

Notation "'CWP' e @ R {{ Φ } }" := (cwp e%E R%I Φ)
  (at level 20, e, Φ at level 200, only parsing) : bi_scope.
Notation "'CWP' e {{ Φ } }" := (cwp e%E True%I Φ)
  (at level 20, e, Φ at level 200, only parsing) : bi_scope.

Notation "'CWP' e @ R {{ v , Q } }" := (cwp e%E R%I (λ v, Q))
  (at level 20, e, Q at level 200,
   format "'[' 'CWP'  e  '/' '[          ' @  R  {{  v ,  Q  } } ']' ']'") : bi_scope.
Notation "'CWP' e {{ v , Q } }" := (cwp e%E True%I (λ v, Q))
  (at level 20, e, Q at level 200,
   format "'[' 'CWP'  e  '/' '[   ' {{  v ,  Q  } } ']' ']'") : bi_scope.

Section cwp_rules.
  Context `{cmonadG Σ}.

  Lemma cwp_wp R Φ Ψ e :
    CWP e @ R {{ Φ }} -∗
    (∀ v : val, CWP v @ R {{ Φ }} -∗ Ψ v) -∗
    WP e {{ Ψ }}.
  Proof.
    iIntros "Hwp H". rewrite cwp_eq /=. iApply (wp_wand with "Hwp").
    iIntros (v) "Hwp". iApply "H". by iApply wp_value'.
  Qed.

  Lemma wp_cwp_bind R Φ K e :
    WP e {{ v, CWP (fill K (of_val v)) @ R {{ Φ }} }} -∗
    CWP fill K e @ R {{ Φ }}.
  Proof. rewrite cwp_eq. by apply: wp_bind. Qed.

  Lemma wp_cwp_bind_inv R Φ K e :
    CWP fill K e @ R {{ Φ }} -∗
    WP e {{ v, CWP fill K (of_val v) @ R {{ Φ }} }}.
  Proof. rewrite cwp_eq. by apply: wp_bind_inv. Qed.

  Lemma cwp_insert_res e Φ R1 R2 :
    ▷ R1 -∗
    CWP e @ (R1 ∗ R2) {{ v, ▷ R1 ={⊤}=∗ Φ v }} -∗
    CWP e @ R2 {{ Φ }}.
  Proof.
    iIntros "HR1 Hcwp". rewrite cwp_eq.
    iApply (wp_wand with "Hcwp").
    iIntros (v) "HΦ".
    iIntros (γ env l I) "#Hflock Hres #Heq".
    iMod (flock_res_alloc_strong _ (dom (gset prop_id) I) with "Hflock HR1") as (j ρ) "[% Hres']"; first done.
    pose (I' := <[j:= (LockRes R1 1 ρ)]>I).
    assert (I !! j = None) by by eapply not_elem_of_dom.
    iSpecialize ("HΦ" $! _ env l I' with "Hflock [Hres Hres'] []").
    { rewrite /flock_resources /I'.
      rewrite big_sepM_insert //. iFrame. }
    { rewrite big_sepM_insert // /=. iRewrite "Heq".
      rewrite (assoc _ R1 _ R2).
      rewrite (comm _ R1 (env_inv env)).
      rewrite -(assoc _ _ R1 R2). done. }
    iApply wp_fupd.
    iApply (wp_wand with "HΦ").
    iIntros (w) "[HΦ HI]". rewrite /flock_resources /I'.
    rewrite big_sepM_insert /=; last done. iDestruct "HI" as "[HR1 $]".
    iMod (flock_res_dealloc with "Hflock HR1") as "HR"; try done.
    by iApply "HΦ".
  Qed.

  Lemma cwp_fupd_wand e Φ Ψ R :
    CWP e @ R {{ Φ }} -∗
    (∀ v, Φ v ={⊤}=∗ Ψ v) -∗
    CWP e @ R {{ Ψ }}.
  Proof.
    iIntros "Hwp H". rewrite cwp_eq.
    iApply (wp_wand with "Hwp"); iIntros (v) "HΦ".
    iIntros (γ env l I) "#Hflock Hres #Heq". iApply wp_fupd.
    iApply (wp_wand with "[HΦ Hres]"). iApply ("HΦ" with "Hflock Hres Heq").
    iIntros (w) "[HΦ $]". by iApply "H".
  Qed.

  Lemma cwp_fupd e Φ R :
    CWP e @ R {{ v, |={⊤}=> Φ v }} -∗ CWP e @ R {{ Φ }}.
  Proof. iIntros "Hwp". iApply (cwp_fupd_wand with "Hwp"); auto. Qed.

  Lemma fupd_cwp e Φ R :
    (|={⊤}=> CWP e @ R {{ v, Φ v }}) -∗ CWP e @ R {{ Φ }}.
  Proof. rewrite cwp_eq. by iIntros ">Hwp". Qed.

  Lemma cwp_wand e Φ Ψ R :
    CWP e @ R {{ Φ }} -∗
    (∀ v, Φ v -∗ Ψ v) -∗
    CWP e @ R {{ Ψ }}.
  Proof.
    iIntros "Hwp H". iApply (cwp_fupd_wand with "Hwp"); iIntros (v) "HΦ !>".
    by iApply "H".
  Qed.

  Lemma cwp_pure K φ n e1 e2 R Φ :
    PureExec φ n e1 e2 →
    φ →
    ▷^n CWP (fill K e2) @ R {{ Φ }} -∗
    CWP (fill K e1) @ R {{ Φ }}.
  Proof.
    iIntros (? Hφ) "Hcwp". iApply wp_cwp_bind. wp_pure _.
    by iApply wp_cwp_bind_inv.
  Qed.

  Lemma cwp_ret e R Φ :
    WP e {{ Φ }} -∗ CWP c_ret e @ R {{ Φ }}.
  Proof.
    iIntros "Hwp". rewrite cwp_eq /cwp_def. wp_apply (wp_wand with "Hwp").
    iIntros (v) "HΦ". wp_lam. wp_pures.
    iIntros (γ env l I) "#Hlock Hres #Heq". wp_pures. iFrame.
  Qed.

  Lemma cwp_bind (f : val) (e : expr) R Φ :
    CWP e @ R {{ ev, CWP f ev @ R {{ Φ }} }} -∗
    CWP c_bind e f @ R {{ Φ }}.
  Proof.
    iIntros "Hwp". rewrite cwp_eq /cwp_def.
    wp_apply (wp_wand with "Hwp"). iIntros (ev) "Hwp".
    wp_lam. wp_pures.
    iIntros (γ env l I) "#Hflock Hres #Heq". wp_pures. wp_bind (ev env l).
    iApply (wp_wand with "[Hwp Hres]"). iApply ("Hwp" with "Hflock Hres Heq").
    iIntros (w) "[Hwp Hres]". wp_let. wp_apply (wp_wand with "Hwp").
    iIntros (v) "H". iApply ("H" with "Hflock Hres Heq").
  Qed.

  Lemma cwp_atomic (ev : val) R Φ :
    (R -∗ ▷ ∃ R', R' ∗ CWP ev #() @ R' {{ w, R' -∗ R ∗ Φ w }}) -∗
    CWP c_atomic ev @ R {{ Φ }}.
  Proof.
    iIntros "Hwp". rewrite cwp_eq /cwp_def. wp_lam. wp_pures.
    iIntros (γ env l I) "#Hlock1 Hres #Heq1". wp_pures.
    wp_apply (acquire_flock_spec with "[$]").
    iIntros "Hfl".
    iMod (flocked_inv_open with "Hfl") as "[HI Hcl]"; first done.
    iRewrite "Heq1" in "HI".
    iDestruct "HI" as "[Henv HR]".
    wp_pures; simpl.
    iDestruct ("Hwp" with "HR") as (Q) "[HQ Hwp]".
    wp_apply (newflock_spec cmonadN); first done.
    iIntros (k γ') "#Hlock2".
    iMod (flock_res_alloc_strong _ ∅ _ _ (env_inv env ∗ Q)%I with "Hlock2 [$HQ $Henv]") as (s ρ) "[_ Hres]"; first done.
    wp_let.
    wp_apply (wp_wand with "Hwp"); iIntros (ev') "Hwp". wp_bind (ev' _ _).
    iApply (wp_wand with "[Hwp Hres]").
    - iApply ("Hwp" $! _ _ _ {[s:=(LockRes _ 1 ρ)]}  with "Hlock2 [Hres] []").
      + rewrite /flock_resources big_sepM_singleton //.
      + rewrite big_sepM_singleton //.
    - iIntros (w) "[HR Hres]".
      rewrite /flock_resources big_sepM_singleton /=.
      iMod (flock_res_dealloc with "Hlock2 Hres") as "[Henv HQ]"; try done.
      wp_let.
      iDestruct ("HR" with "HQ") as "[HR HΦ]".
      iMod ("Hcl" with "[HR Henv]") as "Hflocked".
      { iNext. iRewrite "Heq1". iFrame. }
      wp_apply (release_cancel_spec with "[$Hlock1 $Hflocked]").
      iIntros "$". wp_pures. iFrame.
  Qed.

  Lemma cwp_atomic_env (ev : val) R Φ :
    (∀ env, env_inv env -∗ R -∗
      WP ev env {{ w, ▷ (env_inv env ∗ R ∗ Φ w) }}) -∗
    CWP c_atomic_env ev @ R {{ Φ }}.
  Proof.
    iIntros "Hwp". rewrite cwp_eq /cwp_def. wp_lam. wp_pures.
    iIntros (γ env l I) "#Hlock Hres #Heq". wp_pures.
    wp_apply (acquire_flock_spec with "[$]").
    iIntros "Hfl".
    iMod (flocked_inv_open with "Hfl") as "[HI Hcl]"; first done.
    iRewrite "Heq" in "HI".
    iDestruct "HI" as "[Henv HR]".
    wp_pures; simpl.
    iSpecialize ("Hwp" with "Henv HR").
    wp_apply (wp_wand with "Hwp").
    iIntros (w) "[Henv [HR HΦ]]". wp_pures.
    iRewrite "Heq" in "Hcl".
    iMod ("Hcl" with "[$HR $Henv]") as "Hflocked".
    wp_apply (release_cancel_spec with "[$Hlock $Hflocked]").
    iIntros "$". wp_pures. iFrame.
  Qed.

  Lemma cwp_par Ψ1 Ψ2 e1 e2 R Φ :
    CWP e1 @ R {{ Ψ1 }} -∗
    CWP e2 @ R {{ Ψ2 }} -∗
    ▷ (∀ w1 w2, Ψ1 w1 -∗ Ψ2 w2 -∗ ▷ Φ (w1,w2)%V) -∗
    CWP e1 |||ᶜ e2 @ R {{ Φ }}.
  Proof.
    iIntros "Hwp1 Hwp2 HΦ". rewrite cwp_eq /cwp_def.
    wp_apply (wp_wand with "Hwp2").
    iIntros (ev2) "Hwp2".
    wp_apply (wp_wand with "Hwp1").
    iIntros (ev1) "Hwp1". wp_lam. wp_pures.
    iIntros (γ env l I) "#Hlock Hres #Heq". wp_pures.
    pose (I' := fmap (λ X, LockRes (res_prop X) (res_frac X/2) (res_name X)) I).
    iAssert (flock_resources γ I' ∗ flock_resources γ I')%I with "[Hres]" as "[Hres1 Hres2]".
    { rewrite /flock_resources -big_sepM_sepM.
      rewrite /I' big_sepM_fmap /=.
      iApply (big_sepM_mono with "Hres"). iIntros (k x Hk). simpl.
      by rewrite -flock_res_op Qp_div_2. }
    iApply (par_spec (λ v, Ψ1 v ∗ flock_resources γ I')%I
                     (λ v, Ψ2 v ∗ flock_resources γ I')%I
      with "[Hwp1 Hres1] [Hwp2 Hres2]").
    - wp_lam. iApply ("Hwp1" with "Hlock Hres1").
      by rewrite /I' big_sepM_fmap /=.
    - wp_lam. iApply ("Hwp2" with "Hlock Hres2").
      by rewrite /I' big_sepM_fmap /=.
    - iNext. iIntros (w1 w2) "[[HΨ1 Hres1] [HΨ2 Hres2]]".
      iAssert (flock_resources γ I)%I with "[Hres1 Hres2]" as "$".
      { iCombine "Hres1 Hres2" as "Hres".
        rewrite /flock_resources -big_sepM_sepM.
        rewrite /I' big_sepM_fmap /=.
        iApply (big_sepM_mono with "Hres"). iIntros (k x Hk). simpl.
          by rewrite -flock_res_op Qp_div_2. }
      iApply ("HΦ" with "[$] [$]").
  Qed.

  Global Instance frame_cwp p R' e R Φ Ψ :
    (∀ v, Frame p R (Φ v) (Ψ v)) →
    Frame p R (CWP e @ R' {{ Φ }}) (CWP e @ R' {{ Ψ }}).
  Proof.
    rewrite /Frame. iIntros (HR) "[HR H]". iApply (cwp_wand with "H").
    iIntros (v) "H". iApply HR; iFrame.
  Qed.

  Global Instance is_except_0_cwp R e Φ : IsExcept0 (CWP e @ R {{ Φ }}).
  Proof. rewrite /IsExcept0. iIntros "H". iApply fupd_cwp. by iMod "H". Qed.

  Global Instance elim_modal_bupd_cwp p R e P Φ :
    ElimModal True p false (|==> P) P (CWP e @ R {{ Φ }}) (CWP e @ R {{ Φ }}).
  Proof.
    rewrite /ElimModal bi.intuitionistically_if_elim; iIntros (_) "[HP HR]".
    iApply fupd_cwp. iMod "HP". by iApply "HR".
  Qed.

  Global Instance elim_modal_fupd_wp p R e P Φ :
    ElimModal True p false (|={⊤}=> P) P (CWP e @ R {{ Φ }}) (CWP e @ R {{ Φ }}).
  Proof.
    rewrite /ElimModal bi.intuitionistically_if_elim; iIntros (_) "[HP HR]".
    iApply fupd_cwp. iMod "HP". by iApply "HR".
  Qed.

  Global Instance add_modal_fupd_wp R e P Φ :
    AddModal (|={⊤}=> P) P (CWP e @ R {{ Φ }}).
  Proof. rewrite /AddModal. iIntros "[>HP H]". by iApply "H". Qed.
End cwp_rules.

Section cwp_run.
  Context `{heapG Σ, flockG Σ, spawnG Σ, locking_heapPreG Σ}.

  Lemma cwp_run (ev : val) Φ :
    (∀ `{cmonadG Σ}, CWP ev {{ w, Φ w }}) -∗
    WP c_run ev {{ Φ }}.
  Proof.
    iIntros "Hwp". wp_lam.
    wp_bind (mset_create #()). iApply mset_create_spec; first done.
    iNext. iIntros (env) "Henv". wp_let.
    iMod locking_heap_init as (?) "Hσ".
    pose (amg := AMonadG Σ _ _ _ _).
    iSpecialize ("Hwp" $! amg). rewrite cwp_eq /cwp_def.
    wp_apply (newflock_spec cmonadN); first done.
    iIntros (k γ') "#Hlock". iApply wp_fupd.
    iMod (flock_res_alloc_strong _ ∅ _ _ (env_inv env)%I
        with "Hlock [Henv Hσ]") as (s ρ) "[_ Hres]"; first done.
    { iNext. iExists ∅, ∅. iFrame. iPureIntro; set_solver. }
    wp_let.
    iMod (wp_value_inv with "Hwp") as "Hwp".
    iApply (wp_wand with "[Hwp Hres]").
    - iApply ("Hwp" $! _ _ _ {[s := LockRes _ 1 ρ]} with "Hlock [Hres] []").
      + rewrite /flock_resources big_sepM_singleton //.
      + by rewrite big_sepM_singleton right_id.
    - iIntros (w) "[HΦ Hres]".
      rewrite /flock_resources big_sepM_singleton /=.
      by iMod (flock_res_dealloc with "Hlock Hres") as "Henv".
  Qed.
End cwp_run.

(* Make sure that we only use the provided rules and don't break the abstraction *)
Typeclasses Opaque c_ret c_bind c_run c_atomic c_atomic_env c_par.
Opaque c_ret c_bind c_run c_atomic c_atomic_env c_par.

(* Definition locking_heapΣ : gFunctors := *)
(*   #[heapΣ; GFunctor (auth.authR locking_heapUR)]. *)

(* Instance subG_locking_heapG {Σ} : subG locking_heapΣ Σ → locking_heapPreG Σ. *)
(* Proof. solve_inG. Qed. *)

(* Definition cwp_adequacy Σ R s v σ φ : *)
(*   (R -∗ (∀ `{locking_heapG Σ}, cwp (of_val v) R (λ w, R -∗ ⌜φ w⌝)))%I → *)
(*   adequate MaybeStuck (a_run v) σ φ. *)
(*   (∀ `{heapG Σ}, WP e @ s; ⊤ {{ v, ⌜φ v⌝ }}%I) → *)
(* Proof. *)
(*   intros Hwp; eapply (wp_adequacy _ _); iIntros (?) "". *)
(*   iMod (gen_heap_init σ) as (?) "Hh". *)
(*   iModIntro. iExists gen_heap_ctx. iFrame "Hh". *)
(*   iApply (Hwp (HeapG _ _ _)). *)
(* Qed. *)