From iris.base_logic.lib Require Import invariants own.
From iris.heap_lang Require Import proofmode notation par assert.

(** Each thread is the same: an incr in a critical section *)

Definition check : val :=
  λ: "r",
    let: "x" := !"r" in
    let: "y" := !"r" in
    assert:"x" ≤ "y".

Definition loop_check : val :=
  rec: "loop" "r" :=
    check "r";;
    "loop" "r".

Definition loop_incr : val :=
  rec: "loop" "r" :=
    FAA "r" #1;;
    "loop" "r".

Definition main : expr :=
  let: "r" := ref #0 in
  loop_check "r" ||| loop_incr "r".


(** Resource algebra laws *)

(** since [max] idempotent, we can spawn [◯a] from [●a] *)

Lemma authoritative_gives_fragment a :
  ● (MaxNat a) ~~> ● (MaxNat a) ⋅ ◯ (MaxNat a).
Proof.
  apply auth_update_alloc.
  apply local_update_unital_discrete.
  intros b Vk.
  rewrite left_id.
  intros <-. split; auto.
  rewrite max_nat_op.
  f_equal.
  lia.
Qed.

(** [valid (●a⋅◯b)] implies [b ≼ a] i.e. [b <= a] *)

Lemma fragment_leq_authoritative a b :
  valid (● MaxNat a ⋅ ◯ MaxNat b) → a >= b.
Proof.
  (* also: by intros [?%max_nat_included _]%auth_both_valid_discrete. *)
  intros Hvalid.
  apply auth_both_valid_discrete in Hvalid.
  destruct Hvalid as [Hvalid meh].
  apply max_nat_included in Hvalid.
  apply Hvalid.
Qed.


Section Proofs.

Context `{!heapGS Σ, spawnG Σ}.
Context `{inG Σ (authR max_natO)}.

Definition mcountInv r γ : iProp Σ := TODO.

Definition frag γ k := own γ (◯ (MaxNat k)).

(** Proof of the left thread *)

Lemma read_spec N r γ n :
  {{{ inv N (mcountInv r γ) ∗ own γ (◯ (MaxNat n)) }}}
    ! #r
  {{{ k, RET #k; ⌜k >= n⌝ ∗ own γ (◯ (MaxNat k)) }}}.
Proof.
  iIntros (post) "(#Hinv & #Hfrag) Hpost".
  iInv "Hinv" as (k) "(Hr & Hauth)" "Hclose".
  wp_load.

  (* before closing the invariant, extract the knowledge [k >= n] *)
  iCombine "Hauth" "Hfrag" gives %Hvalid%fragment_leq_authoritative.

  (* also before closing the invariant, we  the knowledge [k >= n] *)
  iPoseProof (own_update _ _ _ (authoritative_gives_fragment k) with "Hauth") as "Hboth".
  iMod "Hboth".
  iDestruct "Hboth" as "(Hauth & #Hfrag')".

  (* closing the invariant *)
  iMod ("Hclose" with "[Hr Hauth]"). now iExists k; iFrame.
  iModIntro.
  iApply "Hpost". iFrame "Hfrag'". auto.
Qed.

Lemma read_spec_no_frag_in_pre N r γ :
  {{{ inv N (mcountInv r γ) }}}
    ! #r
  {{{ k, RET #k; own γ (◯ (MaxNat k)) }}}.
Proof.
  iIntros (post) "#Hinv Hpost".
  iInv "Hinv" as (a) "(Hr & Hauth)" "Hclose".
  wp_load.
  iMod (own_update _ _ _ (authoritative_gives_fragment a) with "Hauth")
    as "(Hauth & #Hfrag')".
  iMod ("Hclose" with "[Hr Hauth]"). now iExists a; iFrame.
  iModIntro.
  iApply "Hpost". iFrame "Hfrag'".
Qed.

(* we could remove the [◯ _] in precondition since it is deducible from ● *)
Lemma check_spec N r γ :
  {{{ inv N (mcountInv r γ) ∗ own γ (◯ (MaxNat 0)) }}}
    check #r
  {{{ RET #(); True }}}.
Proof.
  iIntros (post) "(#Hinv & #Hfrag) Hpost".
  unfold check.
  wp_pures.
  wp_apply (read_spec with "[$]") as (x) "(%Hx & Hx)".
  wp_pures.
  wp_apply (read_spec with "[$]") as (y) "(%Hy & Hy)".
  wp_pures.
  wp_apply wp_assert.
  wp_pures.
  case_bool_decide. 2:lia.
  iModIntro. iSplit; auto.
  iApply "Hpost".
  auto.
Qed.

Lemma incr_spec r ι γ n :
  {{{ inv ι (mcountInv r γ) ∗ own γ (◯ (MaxNat n)) }}}
    FAA #r #1
  {{{ k, RET #k; own γ (◯ (MaxNat (k + 1))) ∗ ⌜k >= n⌝ }}}.
Proof.
  iIntros (post) "(Hinv & Hfrag) Hpost".
  iInv "Hinv" as (a) "(Hr & Hauth)" "Hclose".
  wp_faa.

  (* before closing the invariant, extract the knowledge [k >= n] *)
  iCombine "Hauth" "Hfrag" gives %Hvalid%fragment_leq_authoritative.

    
  assert (● (MaxNat a) ⋅ ◯ (MaxNat n) ~~> ● (MaxNat (a + 1)) ⋅ ◯ (MaxNat (a + 1))) as Hfpu.
  { apply auth_update.
    apply local_update_unital_discrete.
    intros [i] Vk.
    intros [=->].
    split; auto.
    rewrite max_nat_op. f_equal.
    lia. }

  iPoseProof (own_update _ _ _ Hfpu with "[Hauth Hfrag]") as "Hboth".
  by rewrite own_op; iFrame.
  iMod "Hboth".
  iDestruct "Hboth" as "(Hauth & #Hfrag)".
  iMod ("Hclose" with "[Hr Hauth]").
  { iExists (a + 1).
    replace ((a + 1)%nat : Z) with (Z.add a 1) by lia.
    iFrame. }
  iModIntro.
  iApply "Hpost". iFrame "Hfrag". auto.
Qed.

Definition mcountN : namespace :=
  nroot .@ "mcount".

Lemma main_spec : {{{ True }}} main {{{ RET #42; False }}}.
Proof.
  iIntros (post) "Htrue Hpost".
  unfold main.
  wp_alloc r as "Hr".
  wp_pures.
  iMod (own_alloc (● (MaxNat 0) ⋅ ◯ (MaxNat 0))) as (γ) "(Hauth & #Hfrag)".
  { apply auth_both_valid_2.
    constructor.
    unfold "≼".
    exists (MaxNat 0).
    compute.
    reflexivity. }

  iMod (inv_alloc mcountN _ (mcountInv r γ) with "[Hr Hauth]") as "#Hinv".
  by iExists 0; iFrame.

  wp_apply (par_spec (λ _, ⌜False⌝%I) (λ _, ⌜False⌝%I)) as (? ?) "(Hfalse & Hfalse')".

  - wp_pure.
    iLöb as "IH".
    wp_rec.
    wp_apply (check_spec with "[$]") as "Htrue".
    wp_pures.
    auto.

  - wp_pure.
    iLöb as "IH".
    wp_rec.
    wp_apply (incr_spec with "[$]") as (k) "A".
    wp_pures.
    auto.

  - iExFalso.
    auto.
Qed.


End Proofs.