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

Definition rand : val := λ: <>,
  let: "y" := ref #false
  in Fork ("y" <- #true);;
     !"y".

Definition lateChoice : val := λ: "x",
  "x" <- #0;; rand #().
Definition earlyChoice : val := λ: "x",
  let: "r" := rand #() in "x" <- #0;; "r".

Section Refinement.
  Context `{logrelG Σ}.

  Definition choiceN : namespace := nroot .@ "choice".

  Definition choice_inv y y' : iProp Σ :=
    (∃ f : bool, y ↦ᵢ #f ∗ y' ↦ₛ #f)%I.

  Lemma wp_rand :
    (WP rand #() {{ v, ⌜v = #true⌝ ∨ ⌜v = #false⌝}})%I.
  Proof.
    unfold rand. unlock.
    wp_rec.
    wp_alloc y as "Hy".
    iMod (inv_alloc choiceN _ (y ↦ᵢ #false ∨ y ↦ᵢ #true)%I with "[Hy]") as "#Hinv"; eauto.
    wp_rec.
    wp_bind (Fork _). iApply wp_fork. iNext.
    iSplitL.
    - iModIntro. wp_rec.
      iInv choiceN as "[Hy | Hy]" "Hcl"; iApply (wp_load with "Hy"); eauto; iNext;
        iIntros "Hy"; iMod ("Hcl" with "[Hy]"); eauto.
    - iInv choiceN as "[Hy | Hy]" "Hcl"; iApply (wp_store with "Hy"); eauto; iNext;
        iIntros "Hy"; iMod ("Hcl" with "[Hy]"); eauto.
  Qed.

  Lemma rand_l Δ Γ K ρ t τ :
    spec_ctx ρ -∗ (∀ b : bool, {Δ;Γ} ⊨ fill K #b ≤log≤ t : τ) -∗
    {Δ;Γ} ⊨ fill K (rand #()) ≤log≤ t : τ.
  Proof.
    iIntros  "#Hs Hlog".
    unfold rand. unlock. simpl.
    rel_rec_l.
    rel_alloc_l as y "Hy". simpl.
    rel_rec_l.
    iMod (inv_alloc choiceN _ (∃ b : bool, y ↦ᵢ #b)%I with "[Hy]") as "#Hinv".
    { iNext. eauto. }
    rel_fork_l. iModIntro.
    iSplitR.
    - iNext.
      iInv choiceN as (b) "Hy" "Hcl".
      iApply (wp_store with "Hy"); eauto. iNext. iIntros "Hy".
      iMod ("Hcl" with "[Hy]").
      { iNext. iExists true. by iFrame. }
      done.
    - simpl.
      rel_rec_l.
      rel_load_l_atomic.
      iInv choiceN as (b) "Hy" "Hcl".
      iExists #b. iFrame.
      iModIntro. iNext. iIntros "Hy".
      iMod ("Hcl" with "[Hy]").
      { iNext. iExists b. iFrame. }
      done. 
  Qed.

  Lemma rand_r (b : bool) Δ Γ E1 K ρ t τ :
    ↑specN ⊆ E1 →
    ↑choiceN ⊆ E1 →
    spec_ctx ρ -∗
    {E1;Δ;Γ} ⊨ t ≤log≤ fill K #b : τ -∗
    {E1;Δ;Γ} ⊨ t ≤log≤ fill K (rand #()) : τ.
  Proof.
    iIntros (??) "#Hs Hlog".
    unfold rand. unlock.
    rel_rec_r.
    rel_alloc_r as y "Hy".
    rel_rec_r.
    rel_fork_r as j "Hj".
    rel_rec_r.
    destruct b.
    - tp_store j.
      by rel_load_r.
    - by rel_load_r.
  Qed.

  Lemma lateChoice_l Δ Γ x v ρ t :
    spec_ctx ρ -∗ x ↦ᵢ v -∗
    (x ↦ᵢ #0 -∗ ∀ b : bool, {Δ;Γ} ⊨ #b ≤log≤ t : TBool) -∗
    {Δ;Γ} ⊨ lateChoice #x ≤log≤ t : TBool.
  Proof.
    iIntros "#Hs Hx Hlog".
    unfold lateChoice. unlock.
    rel_rec_l.
    rel_store_l. simpl.
    rel_rec_l.
    rel_apply_l (rand_l with "Hs"); eauto.
    by iApply "Hlog".
  Qed.
   
  Lemma prerefinement Δ Γ x x' n ρ :
    spec_ctx ρ -∗ x ↦ᵢ #n -∗ x' ↦ₛ #n -∗
    {Δ;Γ} ⊨ lateChoice #x ≤log≤ earlyChoice #x' : TBool.
  Proof.
    iIntros "#Hspec Hx Hx'".
    iApply (lateChoice_l with "Hspec Hx"). iIntros "Hx".
    iIntros (b).
    unfold earlyChoice. unlock.
    rel_rec_r.

    rel_apply_r (rand_r b with "Hspec"); eauto.
    rel_rec_r.
    rel_store_r. simpl.
    rel_rec_r.
    rel_vals; eauto.
  Qed.

  Lemma prerefinement2 Δ Γ x x' n ρ :
    spec_ctx ρ -∗ x ↦ᵢ #n -∗ x' ↦ₛ #n -∗
    {Δ;Γ} ⊨ earlyChoice #x ≤log≤ lateChoice #x' : TBool.
  Proof.
    iIntros "#Hspec Hx Hx'".
    unfold earlyChoice. unlock.
    rel_rec_l.
    rel_apply_l (rand_l with "Hspec"); eauto.
    iIntros (b).
    rel_rec_l.

    unfold lateChoice. unlock.
    rel_rec_r.
    rel_store_r. simpl.
    rel_rec_r.
    rel_apply_r (rand_r b with "Hspec"); eauto.

    rel_store_l. simpl.
    rel_rec_l.
    rel_vals; eauto.
  Qed.

End Refinement.