(* Requires Iris 3.0 commit 6b68b1716d0b3b89395995665625d165d2f6f6bd *)
From iris.program_logic Require Export weakestpre hoare.
From iris.base_logic.lib Require Export invariants.
From iris.heap_lang.lib Require Import par.
From iris.heap_lang Require Import proofmode notation.
From iris.algebra Require Import frac_auth.
Open Scope nat_scope.

Definition inc : val :=
  rec: "inc" "l" "n" :=
    let: "x" := !"l" in
    if: CAS "l" "x" ("x" + "n")
      then "x"
      else "inc" "l" "n".

Definition example : val := λ: <>,
  let: "l" := ref #0 in
  (inc "l" #2 ||| inc "l" #2);;
  !"l".

Section example.
  Context `{heapG Σ, spawnG Σ, inG Σ (frac_authR natR)}.
  Local Definition N := nroot .@ "example".
  Implicit Types n : nat.

  (* Rules for fractional ghost variables
     (proved from generic principles) *)
  Lemma frac_auth_alloc n :
    True -∗ |==> ∃ γ, own γ (●! n) ∗ own γ (◯!{1} n).
  Proof.
    by iMod (own_alloc (●! n ⋅ ◯! n)) as (γ) "[??]"; eauto with iFrame.
  Qed.

  Lemma frac_auth_update n n1 n2 q γ :
    own γ (●! n1) -∗
    own γ (◯!{q} n2) -∗ |==>
    own γ (●! (n1 + n)) ∗ own γ (◯!{q} (n2 + n)).
  Proof.
    iIntros "H H'". iMod (own_update_2 with "H H'") as "[$ $]"; last done.
    apply frac_auth_update, nat_local_update. lia.
  Qed.

  Lemma frac_auth_agree n n' γ :
    own γ (●! n) -∗ own γ (◯!{1} n') -∗ ⌜n = n'⌝.
  Proof.
    iIntros "H H'".
    by iDestruct (own_valid_2 with "H H'") as %->%frac_auth_agreeL.
  Qed.

  (* The invariant that we use *)
  Definition example_inv (γ : gname) (l : loc) : iProp Σ :=
    (∃ n : nat, own γ (●! n) ∗ l ↦ #n)%I.

  (* Proof of the threads *)
  Lemma inc_spec n n' γ l q :
    {{ inv N (example_inv γ l) ∗ own γ (◯!{q} n) }}
      inc #l #n'
    {{ _, own γ (◯!{q} (n + n')) }}.
  Proof.
    iIntros "!# [#Hinv Hγ]". iLöb as "IH".
    wp_rec; wp_let.
    wp_bind (! _)%E. iInv N as (m) ">[Hauth Hx]" "Hclose".
    wp_load.
    iMod ("Hclose" with "[Hauth Hx]") as "_".
    { iNext. rewrite {2}/example_inv. iExists m. iFrame. }
    iModIntro. wp_let. wp_op.
    wp_bind (CAS _ _ _).
    iInv N as (m') ">[Hauth Hx]" "Hclose".
    destruct (decide (m = m')) as [->|?].
    - wp_cas_suc.
      iMod (frac_auth_update n' with "Hauth Hγ") as "[Hauth Hγ]".
      iMod ("Hclose" with "[Hauth Hx]") as "_".
      { iNext; iExists (m' + n')%nat. rewrite Nat2Z.inj_add. iFrame. }
      iModIntro. wp_if. done.
    - wp_cas_fail; first naive_solver.
      iMod ("Hclose" with "[Hauth Hx]") as "_".
      { iNext. iExists m'. iFrame. }
      iModIntro. wp_if. by iApply "IH".
  Qed.

  (* Proof of the whole program *)
  Lemma example_spec :
    {{ True }} example #() {{ v, ⌜v = #4%nat⌝ }}.
  Proof.
    iIntros "!# _". wp_lam. wp_alloc l as "Hl". wp_let.
    iMod (frac_auth_alloc 0) as (γ) "[Hauth Hγ]".
    iDestruct "Hγ" as "[Hγ1 Hγ2]".
    iMod (inv_alloc _ _ (example_inv γ l) with "[Hl Hauth]") as "#Hinv".
    { iNext. iExists 0%nat. iFrame. }
    (**
      wp_par :

      WP e1 {{ Φ1 }} -∗
      WP e2 {{ Φ2 }} -∗
      (∀ v1, Φ v1 -∗ Φ v2 -∗ Ψ (v1,v2)) -∗
      WP e1 || e2 {{ Ψ }}
    **)
    wp_bind (_ ||| _)%E.
    wp_apply (wp_par (λ _, own γ (◯!{1/2} 2%nat))
                     (λ _, own γ (◯!{1/2} 2%nat)) with "[Hγ1] [Hγ2]").
    - iApply (inc_spec 0 2). by iFrame.
    - iApply (inc_spec 0 2). by iFrame.
    - iIntros (v1 v2) "[Hγ1 Hγ2] !>". iCombine "Hγ1 Hγ2" as "Hγ". simpl.
      wp_seq. iInv N as (m) ">[Hauth Hx]" "Hclose".
      iDestruct (frac_auth_agree with "Hauth Hγ") as %->.
      wp_load.
      iMod ("Hclose" with "[Hauth Hx]") as "_".
      { iNext. iExists 4%nat. iFrame. }
      auto.
  Qed.
End example.