From iris.proofmode Require Import tactics.
From iris_logrel Require Export logrel.
From iris_logrel.examples Require Import lock.

Definition CG_increment : val := λ: "x" "l" <>,
  acquire "l";;
  "x" <- #1 + !"x";;
  release "l".

Definition counter_read : val := λ: "x" <>, !"x".

Definition CG_counter : val := λ: <>,
  let: "l" := newlock #() in
  let: "x" := ref #0 in
  (CG_increment "x" "l", counter_read "x").

Definition FG_increment : val := λ: "v", rec: "inc" <> :=
  let: "c" := !"v" in
  if: CAS "v" "c" (#1 + "c")
  then #()
  else "inc" #().

Definition FG_counter : val := λ: <>,
  let: "x" := ref #0 in
  (FG_increment "x", counter_read "x").

Section CG_Counter.
  Context `{logrelG Σ}.
  Variable (Δ : list (prodC valC valC -n> iProp Σ)).
  Open Scope expr_scope.
  
  (* Coarse-grained increment *)  
  Lemma CG_increment_type Γ :
    typed Γ CG_increment (TArrow (Tref TNat) (TArrow LockType (TArrow TUnit TUnit))).
  Proof. solve_typed. Qed.

  Hint Resolve CG_increment_type : typeable.

  Lemma bin_log_related_CG_increment_r Γ K E1 E2 t τ (x l : loc) (n : nat) :
    nclose specN ⊆ E1 →
    (x ↦ₛ (#nv n) -∗ l ↦ₛ (#♭v false) -∗
       (x ↦ₛ (#nv S n) -∗ l ↦ₛ (#♭v false) -∗
         ({E1,E2;Δ;Γ} ⊨ t ≤log≤ fill K (Lit tt) : τ)) -∗
    {E1,E2;Δ;Γ} ⊨ t ≤log≤ fill K ((CG_increment $/ (LitV (Loc x)) $/ LitV (Loc l)) #())%E : τ)%I.
  Proof.
    iIntros (?) "Hx Hl Hlog".
    unfold CG_increment. unlock. simpl_subst/=.
    rel_rec_r.
    rel_apply_r (bin_log_related_acquire_r with "Hl"); eauto.
    iIntros "Hl".
    rel_rec_r.
    rel_load_r.
    rel_op_r.
    rel_store_r.
    rel_rec_r.
    rel_apply_r (bin_log_related_release_r with "Hl"); eauto.
    by iApply "Hlog".
  Qed.
 
  Lemma counter_read_type Γ :
    typed Γ counter_read (TArrow (Tref TNat) (TArrow TUnit TNat)).
  Proof. solve_typed. Qed.

  Hint Resolve counter_read_type : typeable.

  Lemma CG_counter_type Γ :
    typed Γ CG_counter (TArrow TUnit (TProd (TArrow TUnit TUnit) (TArrow TUnit TNat))).
  Proof. solve_typed. Qed.

  Hint Resolve CG_counter_type : typeable.

  (* Fine-grained increment *)
  Lemma FG_increment_type Γ :
    typed Γ FG_increment (TArrow (Tref TNat) (TArrow TUnit TUnit)).
  Proof. solve_typed. Qed.

  Hint Resolve FG_increment_type : typeable.

  Lemma bin_log_FG_increment_l Γ K E x n t τ :
    x ↦ᵢ (#nv n) -∗
    (x ↦ᵢ (#nv (S n)) -∗ {E,E;Δ;Γ} ⊨ fill K (Lit tt) ≤log≤ t : τ) -∗
    {E,E;Δ;Γ} ⊨ fill K (FG_increment #x #()) ≤log≤ t : τ.
  Proof.
    iIntros "Hx Hlog".
    iApply bin_log_related_wp_l.
    wp_bind (FG_increment #x).
    unfold FG_increment. unlock.
    iApply wp_rec; eauto.
    iNext. simpl.
    iApply wp_value; eauto. simpl. by rewrite decide_left.
    iApply wp_rec; eauto.
    iNext. simpl.
    wp_bind (Load (# x)).
    iApply (wp_load with "[Hx]"); auto. iNext.
    iIntros "Hx".
    iApply wp_rec; eauto.
    iNext. simpl.
    wp_bind (BinOp _ _ _).
    iApply (wp_nat_binop);eauto. iNext. iModIntro. simpl.    
    wp_bind (CAS _ _ _).    
    iApply (wp_cas_suc with "[Hx]"); auto.
    iNext. iIntros "Hx".
    iApply wp_if_true. iNext.
    iApply wp_value; auto.
    by iApply "Hlog".
  Qed.

  Lemma FG_counter_type Γ :
    typed Γ FG_counter (TArrow TUnit (TProd (TArrow TUnit TUnit) (TArrow TUnit TNat))).
  Proof. solve_typed. Qed.

  Hint Resolve FG_counter_type : typeable.

  Definition counterN : namespace := nroot .@ "counter".

  Definition counter_inv l cnt cnt' : iProp Σ :=
    (∃ n, l ↦ₛ (#♭v false) ∗ cnt ↦ᵢ (#nv n) ∗ cnt' ↦ₛ (#nv n))%I.

  Lemma bin_log_counter_read_r Γ E1 E2 K x n t τ
    (Hspec : nclose specN ⊆ E1) :
    x ↦ₛ (#nv n)
    -∗ (x ↦ₛ (#nv n)
         -∗ {E1,E2;Δ;Γ} ⊨ t ≤log≤ fill K (#n n)%E : τ)%I
    -∗ {E1,E2;Δ;Γ} ⊨ t ≤log≤ fill K ((counter_read $/ LitV (Loc x)) #())%E : τ.
  Proof.
    iIntros "Hx Hlog".
    unfold counter_read. unlock. simpl.
    rel_rec_r.
    rel_load_r.
    by iApply "Hlog".
  Qed.

  (* A logically atomic specification for
     a fine-grained increment with a baked in frame. *)
  (* Unfortunately, the precondition is not baked in the rule so you can only use it when your spatial context is empty *)
  Lemma bin_log_FG_increment_logatomic R Γ E1 E2 K x t τ  :
    □ (|={E1,E2}=> ∃ n, x ↦ᵢ (#nv n) ∗ R(n) ∗
       ((∃ n, x ↦ᵢ (#nv n) ∗ R(n)) ={E2,E1}=∗ True) ∧
        (∀ m, x ↦ᵢ (#nv (S m)) ∗ R(m) -∗ 
            {E2,E1;Δ;Γ} ⊨ fill K (Lit tt) ≤log≤ t : τ))
    -∗ ({E1,E1;Δ;Γ} ⊨ fill K ((FG_increment $/ LitV (Loc x)) #())%E ≤log≤ t : τ).
  Proof.
    iIntros "#H".
    unfold FG_increment. unlock. simpl.
    iLöb as "IH".
    rel_rec_l.
    iPoseProof "H" as "H2". (* lolwhat *)
    rel_load_l_atomic.
    iMod "H" as (n) "[Hx [HR Hrev]]".  iModIntro.
    iExists (#nv n). iFrame. iNext. iIntros "Hx".
    iDestruct "Hrev" as "[Hrev _]".
    iApply fupd_logrel.
    iMod ("Hrev" with "[HR Hx]").
    { iExists _. iFrame. } iModIntro. simpl.
    rel_rec_l.
    rel_op_l.
    rel_cas_l_atomic.
    iMod "H2" as (n') "[Hx [HR HQ]]". iModIntro. simpl.
    destruct (decide (n = n')); subst.
    - iExists (#nv n'). iFrame.
      iSplitR; eauto. { iDestruct 1 as %Hfoo. exfalso. done. }
      iIntros "_ !> Hx". simpl.
      iDestruct "HQ" as "[_ HQ]".
      iSpecialize ("HQ" $! n' with "[Hx HR]"). { iFrame. }
      rel_if_true_l. done.
    - iExists (#nv n'). iFrame. 
      iSplitL; eauto; last first.
      { iDestruct 1 as %Hfoo. exfalso. simplify_eq. }
      iIntros "_ !> Hx". simpl.
      rel_if_false_l.
      iDestruct "HQ" as "[HQ _]".
      iMod ("HQ" with "[Hx HR]").
      { iExists _; iFrame. }
      iApply "IH".
  Qed.

  (* A similar atomic specification for the counter_read fn *)
  Lemma bin_log_counter_read_atomic_l R Γ E1 E2 K x t τ :
    □ (|={E1,E2}=> ∃ n, x ↦ᵢ (#nv n) ∗ R(n) ∗
       ((∃ n, x ↦ᵢ (#nv n) ∗ R(n)) ={E2,E1}=∗ True) ∧
        (∀ m, x ↦ᵢ (#nv m) ∗ R(m) -∗
            {E2,E1;Δ;Γ} ⊨ fill K (#n m) ≤log≤ t : τ))
    -∗ {E1,E1;Δ;Γ} ⊨ fill K ((counter_read $/ LitV (Loc x)) #())%E ≤log≤ t : τ.
  Proof.
    iIntros "#H".
    unfold counter_read. unlock. simpl.
    rel_rec_l.
    rel_load_l_atomic.
    iMod "H" as (n) "[Hx [HR Hfin]]". iModIntro.
    iExists _; iFrame "Hx". iNext.
    iIntros "Hx".
    iDestruct "Hfin" as "[_ Hfin]".
    iApply "Hfin". by iFrame.
  Qed.

  (* TODO: try to use with_lock rules *)
  Lemma FG_CG_increment_refinement l cnt cnt' Γ :
    inv counterN (counter_inv l cnt cnt') -∗
    {⊤,⊤;Δ;Γ} ⊨ FG_increment $/ LitV (Loc cnt) ≤log≤ CG_increment $/ LitV (Loc cnt') $/ LitV (Loc l) : TArrow TUnit TUnit.
  Proof.
    iIntros "#Hinv".
    iApply bin_log_related_arrow_val.
    { unfold FG_increment. unlock; simpl_subst. reflexivity. }
    { unfold CG_increment. unlock; simpl_subst. reflexivity. }
    { unfold FG_increment. unlock; simpl_subst/=. solve_closed. (* TODO: add a clause to the reflection mechanism that reifies a [lambdasubst] expression as a closed one *) }
    { unfold CG_increment. unlock; simpl_subst/=. solve_closed. }

    iAlways. iIntros (v v') "[% %]"; simpl in *; subst.
    iApply (bin_log_FG_increment_logatomic (fun n => (l ↦ₛ (#♭v false)) ∗ cnt' ↦ₛ #nv n)%I _ _ _ [] cnt with "[Hinv]").
    iAlways.
    iInv counterN as ">Hcnt" "Hcl". iModIntro.
    iDestruct "Hcnt" as (n) "[Hl [Hcnt Hcnt']]".
    iExists n. iFrame. clear n.
    iSplit.
    - iDestruct 1 as (n) "[Hcnt [Hl Hcnt']]".
      iMod ("Hcl" with "[-]").
      { iNext. iExists _. iFrame. }
      done.
    - iIntros (m) "[Hcnt [Hl Hcnt']]".
      iApply (bin_log_related_CG_increment_r _ [] with "[Hcnt'] [Hl]"); auto. { solve_ndisj.  }
      iIntros "Hcnt' Hl".
      iMod ("Hcl" with "[-]").
      { iNext. iExists _. iFrame. }
      simpl.
      by rel_vals.
  Qed.

  Lemma counter_read_refinement l cnt cnt' Γ :
    inv counterN (counter_inv l cnt cnt') -∗
    {⊤,⊤;Δ;Γ} ⊨ counter_read $/ LitV (Loc cnt) ≤log≤ counter_read $/ LitV (Loc cnt') : TArrow TUnit TNat.
  Proof.
    iIntros "#Hinv".
    iApply bin_log_related_arrow_val.
    { unfold counter_read. unlock. simpl. reflexivity. }
    { unfold counter_read. unlock. simpl. reflexivity. }
    { unfold counter_read. unlock. simpl. solve_closed. }
    { unfold counter_read. unlock. simpl. solve_closed. }
    iAlways. iIntros (v v') "[% %]"; simpl in *; subst.
    iApply (bin_log_counter_read_atomic_l (fun n => (l ↦ₛ (#♭v false)) ∗ cnt' ↦ₛ #nv n)%I _ _ _ [] cnt with "[Hinv]").
    iAlways. iInv counterN as (n) "[>Hl [>Hcnt >Hcnt']]" "Hclose". 
    iModIntro. 
    iExists n. iFrame "Hcnt Hcnt' Hl". clear n.
    iSplit.
    - iDestruct 1 as (n) "(Hcnt & Hl & Hcnt')". iApply "Hclose".
      iNext. iExists n. by iFrame.
    - iIntros (m) "(Hcnt & Hl & Hcnt') /=".
      iApply (bin_log_counter_read_r _ _ _ [] with "Hcnt'").
      { solve_ndisj. }
      iIntros "Hcnt'".
      iMod ("Hclose" with "[Hl Hcnt Hcnt']"); simpl.
      { iNext. iExists _. by iFrame. }
      rel_vals. simpl. eauto.
  Qed.

  Lemma FG_CG_counter_refinement :
    {⊤,⊤;Δ;∅} ⊨ FG_counter ≤log≤ CG_counter :
          TArrow TUnit (TProd (TArrow TUnit TUnit) (TArrow TUnit TNat)).
  Proof.
    unfold FG_counter, CG_counter.
    iApply bin_log_related_arrow; try by (unlock; eauto).
    iAlways. iIntros (? ?) "/= ?"; simplify_eq/=.
    unlock. rel_rec_l. rel_rec_r.
    rel_apply_r bin_log_related_newlock_r; auto.
    iIntros (l) "Hl".
    rel_rec_r.
    rel_alloc_r as cnt' "Hcnt'".
    rel_alloc_l as cnt "Hcnt". simpl.

    rel_rec_l.
    rel_rec_r.

    (* establishing the invariant *)
    iAssert (counter_inv l cnt cnt')
      with "[Hl Hcnt Hcnt']" as "Hinv".
    { iExists _. by iFrame. }
    iMod (inv_alloc counterN with "[Hinv]") as "#Hinv"; trivial.

    iApply (bin_log_related_pair _ with "[]").
    - rel_rec_l.
      unfold CG_increment. unlock.
      rel_rec_r.
      rel_rec_r.
      replace (λ: <>, acquire # l ;; #cnt' <- #1 + (! #cnt');; release # l)%E
        with (CG_increment $/ LitV (Loc cnt') $/ LitV (Loc l))%E; last first.
      { unfold CG_increment. unlock; simpl_subst/=. reflexivity. }
      iApply (FG_CG_increment_refinement with "Hinv").
    - rel_rec_l.
      rel_rec_r.
      iApply (counter_read_refinement with "Hinv").
  Qed.
End CG_Counter.

Theorem counter_ctx_refinement :
  ∅ ⊨ FG_counter ≤ctx≤ CG_counter :
         TArrow TUnit (TProd (TArrow TUnit TUnit) (TArrow TUnit TNat)).
Proof.
  eapply (logrel_ctxequiv logrelΣ); [solve_closed.. | intros ].
  apply FG_CG_counter_refinement.
Qed.