(* tests/one_shot.v *)
From iris.algebra Require Import excl agree csum
From iris.proofmode Require Import proofmode
From iris.program_logic Require Export weakestpre
From iris.deprecated.program_logic Require Import hoare
From iris.heap_lang Require Export lang
From iris.heap_lang Require Import assert proofmode notation adequacy
From iris.heap_lang.lib Require Import par
From iris.prelude Require Import options

Definition one_shot_example: val := λ: <>,
  let: "x" := ref NONE in (

Definition one_shotR:= csumR (exclR unitO) (agreeR ZO).

Definition Pending: one_shotR := Cinl (Excl ()).

Definition Shot(n : Z) : one_shotR := Cinr (to_agree n).

Class one_shotGΣ := { #[local] one_shot_inG :: inG Σ one_shotR }.

Definition one_shotΣ: gFunctors := #[GFunctor one_shotR].

Global InstancesubG_one_shotΣ {Σ} : subG one_shotΣ Σ → one_shotG Σ.
Proof. solve_inG. Qed.

Local SetDefault Proof Using "Type*".
Context `{!heapGS Σ, !one_shotG Σ}.

Definition one_shot_inv(γ : gname) (l : loc) : iProp Σ :=
  (l ↦ NONEV ∗ own γ Pending ∨ ∃ n : Z, l ↦ SOMEV #n ∗ own γ (Shot n))%I.

Lemma wp_one_shot(Φ : val → iProp Σ) :
  (∀ f1 f2 : val,
    (∀ n : Z, □ WP f1 #n {{ w, ⌜w = #true⌝ ∨ ⌜w = #false⌝ }}) ∗
    □ WP f2 #() {{ (g : val), □ WP g #() {{ _, True }} }} -∗ Φ (f1,f2)%V)
  ⊢ WP one_shot_example #() {{ Φ }}.
Proof.
  iIntros "Hf /=". pose proof (nroot .@ "N") as N.
  wp_lam. wp_alloc l as "Hl".
  iMod (own_alloc Pending) as (γ) "Hγ"; first done.
  iMod (inv_alloc N _ (one_shot_inv γ l) with "[Hl Hγ]") as "#HN".
  { iNext. iLeft. by iSplitL "Hl". }
  wp_pures. iModIntro. iApply "Hf"; iSplit.
  - iIntros (n) "!>". wp_lam. wp_pures. wp_bind (CmpXchg _ _ _).
    iInv N as ">[[Hl Hγ]|H]"; last iDestruct "H" as (m) "[Hl Hγ]".
    + iMod (own_update with "Hγ") as "Hγ".
      { by apply cmra_update_exclusive with (y:=Shot n). }
      wp_cmpxchg_suc. iModIntro. iSplitL; last (wp_pures; by eauto).
      iNext; iRight; iExists n; by iFrame.
    + wp_cmpxchg_fail. iModIntro. iSplitL; last (wp_pures; by eauto).
      rewrite /one_shot_inv; eauto 10.
  - iIntros "!> /=". wp_lam. wp_bind (! _)%E.
    iInv N as ">Hγ".
    iAssert (∃ v, l ↦ v ∗ ((⌜v = NONEV⌝ ∗ own γ Pending) ∨
       ∃ n : Z, ⌜v = SOMEV #n⌝ ∗ own γ (Shot n)))%I with "[Hγ]" as "Hv".
    { iDestruct "Hγ" as "[[Hl Hγ]|Hl]"; last iDestruct "Hl" as (m) "[Hl Hγ]".
      + iExists NONEV. iFrame. eauto.
      + iExists (SOMEV #m). iFrame. eauto. }
    iDestruct "Hv" as (v) "[Hl Hv]". wp_load.
    iAssert (one_shot_inv γ l ∗ (⌜v = NONEV⌝ ∨ ∃ n : Z,
      ⌜v = SOMEV #n⌝ ∗ own γ (Shot n)))%I with "[Hl Hv]" as "[Hinv #Hv]".
    { 
...

Lemma ht_one_shot(Φ : val → iProp Σ) :
  ⊢ {{ True }} one_shot_example #()
    {{ ff,
      (∀ n : Z, {{ True }} Fst ff #n {{ w, ⌜w = #true⌝ ∨ ⌜w = #false⌝ }}) ∗
      {{ True }} Snd ff #() {{ g, {{ True }} g #() {{ _, True }} }}
    }}.
Proof.
  iIntros "!> _". iApply wp_one_shot. iIntros (f1 f2) "[#Hf1 #Hf2]"; iSplit.
  - iIntros (n) "!> _". wp_smart_apply "Hf1".
  - iIntros "!> _". wp_smart_apply (wp_wand with "Hf2"). by iIntros (v) "#? !> _".
Qed.

Definition client: expr :=
  let: "ff" := one_shot_example #() in
  (Fst "ff" #5 ||| let: "check" := Snd "ff" #() in "check" #()).

Lemma client_safe: ⊢ WP client {{ _, True }}.
  Proof using Type*.
    rewrite /client. wp_apply wp_one_shot. iIntros (f1 f2) "[#Hf1 #Hf2]".
    wp_let. wp_smart_apply wp_par.
    - wp_smart_apply "Hf1".
    - wp_proj. wp_bind (f2 _)%E. iApply wp_wand; first by iExact "Hf2".
      iIntros (check) "Hcheck". wp_pures. iApply "Hcheck".
    - auto.
  Qed.

Definition clientΣ: gFunctors := #[ heapΣ; one_shotΣ; spawnΣ ].

Lemma client_adequateσ : adequate NotStuck client σ (λ _ _, True).
Proof. apply (heap_adequacy clientΣ)=> ?. iIntros "_". iApply client_safe. Qed.