(* modified from https://gitlab.mpi-sws.org/iris/tutorial-popl21/-/blob/master/exercises/ex_03_spinlock.v,
   with pieces from https://gitlab.mpi-sws.org/iris/examples/-/blob/master/theories/lecture_notes/coq_intro_example_1.v *)

(**
This exercise covers the basics of concurrency in Iris.
*)
From iris.base_logic.lib Require Import invariants.
From iris.heap_lang Require Import proofmode notation.
From iris.heap_lang.lib Require Import par.

Section proof.
  Context `{!heapGS Σ}.
  Context `{!spawnG Σ}. (* enable parallelism *)

  Definition parallel_incr : val := λ: "l",
    (("l" <- !"l" + #1) ||| ("l" <- !"l" + #1));;
    !"l".

  (* What should this program's spec be? *)



























  Definition incr_inv (l : loc) (n : Z) :=
    (∃ (m : Z), ⌜(n ≤ m)%Z⌝ ∗ l ↦ #m)%I.




  Definition incrN : namespace := nroot .@ "incr".




  Lemma parallel_incr_spec (l : loc) (n : Z):
    {{{ l ↦ #n }}} parallel_incr #l
    {{{ m, RET #m; ⌜(n ≤ m)%Z⌝ }}}.
  Proof.
    iIntros (Φ) "Hl Hpost".
    unfold parallel_incr.
    wp_lam.

    (* establish the invariant incr_inv *)
    iMod (inv_alloc incrN _ (incr_inv l n) with "[Hl]") as "#inv".
    { unfold incr_inv.
      iExists n; iFrame.
      auto. }
    (* The invariant is persistent! (above the □ line) *)

    (* do the parallel operations, show that each one preserves
      the invariant *)
    wp_smart_apply (wp_par (fun _ => True)%I (fun _ => True)%I).
    (* invariant is persistent, so we don't need any postconditions *)
    + Fail wp_load.
      (* we can open the invariant around a single operation *)
      wp_bind (! #l)%E.
      iInv "inv" as "Hl" "Hclose".
      unfold incr_inv.
      iMod "Hl" as (m) "[% Hl]".
      wp_load.
      iMod ("Hclose" with "[Hl]").
      { iExists m; iFrame; auto. }
      iModIntro.
      wp_op.
      iInv "inv" as "Hl" "Hclose".
      unfold incr_inv.
      iMod "Hl" as (m') "[% Hl]".
      wp_store.
      iMod ("Hclose" with "[Hl]").
      { iExists (m + 1)%Z; iFrame; auto.
        iPureIntro. lia. }
      auto.
    + wp_bind (! #l)%E.
      iInv "inv" as "Hl" "Hclose".
      unfold incr_inv.
      iMod "Hl" as (m) "[% Hl]".
      wp_load.
      iMod ("Hclose" with "[Hl]").
      { iExists m; iFrame; auto. }
      iModIntro.
      wp_op.
      iInv "inv" as "Hl" "Hclose".
      unfold incr_inv.
      iMod "Hl" as (m') "[% Hl]".
      wp_store.
      iMod ("Hclose" with "[Hl]").
      { iExists (m + 1)%Z; iFrame; auto.
        iPureIntro. lia. }
      auto.
    + (* show that the invariant implies the postcondition *)
      (* exercise: complete the proof *)
  Admitted.





  (* spin lock *)

  Definition try_acquire : val := λ: "l",
    CAS "l" #false #true.

  Definition acquire : val :=
    rec: "acquire" "l" :=
      if: try_acquire "l" then #() else "acquire" "l".


  Definition release : val := λ: "l",
    "l" <- #false.


  Definition newlock : val := λ: <>,
    ref #false.


(*
    What should the specs for the lock functions be?

    {{{ ? }}} acquire l {{{ ? }}}
    {{{ ? }}} release l {{{ ? }}}
*)


(**
As shown, the spin lock is implemented as a reference to a Boolean. If the
Boolean is [true], it means some thread has acquired the lock/entered the
critical section. We initiate the lock with the value [false], which means that
the lock is initially in the unlocked state.

The most interesting function of the spin lock is [acquire], which checks if the
lock is in the unlocked state, and if not, changes the lock into [true]. Since
multiple threads could try to acquire the lock at the same time, we use the
[CAS l v w] (compare-and-set) instruction. This instruction atomically checks
if the value of the reference [l] is equal to [v], and if so, changes it into
[w] and returns [true]. If the value of [l] is unequal to [v], it leaves the
value of [l] as if, and returns [false].
[CAS l v w] is actually a short-hand for [Snd (CmpXchg l v w)], where [CmpXchg]
also returns the old value in [l] before the operation took place.
*)



















(**
We will prove the following lock specification in Iris:

  {{{ R }}} newlock #() {{{ lk, RET lk; is_lock lk R }}}.
  {{{ is_lock lk R }}} acquire lk {{{ RET #(); R }}}
  {{{ is_lock lk R ∗ R }}} release lk {{{ RET #(); True }}}.

Here, [is_lock lk R] is a representation predicate which says that a lock at
location [lk] guards the payload [R] described as an Iris proposition.
*)

  (** The invariant of the lock:

  - It owns the memory location [l ↦ #b], which contain a Boolean [b],
  - If the Boolean [b] is [false] (the lock is in the unlocked state),
    then the payload [R] of the lock holds.
  *)
  Definition lock_inv (l : loc) (R : iProp Σ) : iProp Σ :=
    (∃ b : bool, l ↦ #b ∗ if b then True else R)%I.

  (** Invariants in Iris are named by a *namespace* so that several invariants
  can be opened at the same time, while guaranteeing that no invariant is opened
  twice at the same time (which would be unsound). Here, this is irrelevant,
  since acquiring and releasing a lock only requires to open one invariant.

  The namespace [lockN] of the lock invariant:
  *)
  Definition lockN : namespace := nroot .@ "lock".
  Definition is_lock (lk : val) (R : iProp Σ) : iProp Σ :=
    (∃ l: loc, ⌜ lk = #l ⌝ ∧ inv lockN (lock_inv l R))%I.

  (** The main proofs. *)
  Lemma newlock_spec (R : iProp Σ):
    {{{ R }}} newlock #() {{{ lk, RET lk; is_lock lk R }}}.
  Proof.
    iIntros (Φ) "HR HΦ".
    wp_lam. wp_alloc l as "Hl".
    (** Use the Iris rule [inv_alloc] for allocating a lock. We put both the
    resources [HR : R] and the points-to [l ↦ #false] into the lock. *)
    iMod (inv_alloc lockN _ (lock_inv l R) with "[HR Hl]") as "#Hinv".
    { iNext. unfold lock_inv. iExists false. iFrame. }
    iModIntro. iApply "HΦ". unfold is_lock. iExists l. eauto.
  Qed.

  (** *Exercise*: finish the proof below. *)
  Lemma try_acquire_spec lk R :
    {{{ is_lock lk R }}} try_acquire lk
    {{{ b, RET #b; if b is true then R else True }}}.
  Proof.
    iIntros (Φ) "#Hl HΦ".
    unfold is_lock. iDestruct "Hl" as (l ->) "#Hinv".
    wp_rec.
    (** We have to "focus on" the atomic [CmpXchg] operation before we can
    open the invariant. *)
    wp_bind (CmpXchg _ _ _).
    (** Using the tactic [iInv] we open the invariant. *)
    iInv lockN as (b) "[Hl HR]" "Hclose".
    (** The post-condition of the WP is augmented with an *update* modality
    [ |={⊤ ∖ ↑lockN,⊤}=> ... ], which keeps track of the fact that we opened
    the invariant named [lockN]. We can introduce this update modality by
    closing the lock using the hypothesis:

      "Hclose" : ▷ lock_inv l R ={⊤ ∖ ↑lockN,⊤}=∗ True

    *)
    destruct b. (* is the lock held or not? *)
    - wp_cmpxchg_fail. iMod ("Hclose" with "[Hl]") as "_".
      { iNext. iExists true. iFrame. }
      iModIntro. wp_proj.
      (* exercise *) admit.
    - wp_cmpxchg_suc.
      (* exercise *)
  Admitted.

  (** *Exercise*: prove the spec of [acquire]. Since [acquire] is a recursive
  function, you should use the tactic [iLöb] for Löb induction. Use the tactic
  [wp_apply] to use [try_acquire_spec] when appropriate. *)
  Lemma acquire_spec lk R :
    {{{ is_lock lk R }}} acquire lk {{{ RET #(); R }}}.
  Proof.
    iIntros (Φ) "#Hl Hpost".
    iLöb as "IH".
    wp_rec.
    wp_apply (try_acquire_spec with "Hl").
    iIntros (b) "HR".
    destruct b.
    - wp_if_true.
      iApply "Hpost"; auto.
    - wp_if_false.
      wp_apply "IH".
      auto.
    (* done in class *)
  Qed.

  (** *Exercise*: prove the spec of [release]. At a certain point in this proof,
  you need to open the invariant. For this you can use:

    iInv lockN as (b) "[Hl HR]" "Hclose".

  in the same way as in the proof of [try_acquire]. You also need to close the
  invariant. *)
  Lemma release_spec lk R :
    {{{ is_lock lk R ∗ R }}} release lk {{{ RET #(); True }}}.
  Proof.
    (* exercise *)
  Admitted.


  (* Using locks: *)
  Lemma lock_test l :
    {{{ l ↦ #1 }}}
      let: "lk" := newlock #() in
      acquire "lk";; #l <- #3;; release "lk"
    {{{ RET #(); True }}}.
  Proof.
    iIntros (Φ) "Hl Hpost".
    wp_apply (newlock_spec (∃v : Z, l ↦ #v) with "[Hl]").
    { iExists 1. iApply "Hl". }
    iIntros (lk) "#Hlock".
    wp_let.
    wp_apply (acquire_spec with "Hlock").
    iIntros "H".
    iDestruct "H" as (v) "Hl".
    wp_seq.
    wp_store.
    wp_apply (release_spec with "[$Hlock Hl]").
    { iExists 3; auto. }
    iApply "Hpost".
    (* worked in class *)
  Qed.

  (** *Exercise*: Here is a version of parallel_incr that uses locks
  to synchronize the two threads: *)
  Definition parallel_incr_locked (lock : val) : val := λ: "l",
    ((acquire lock;; "l" <- !"l" + #1;; release lock) |||
     (acquire lock;; "l" <- !"l" + #1;; release lock));;
    acquire lock;; let: "res" := !"l" in release lock;; "res".

  (* Prove that it also ensures that "l" only increases. *)
  Lemma parallel_incr_locked_spec lock l (n : Z):
    {{{ is_lock lock (incr_inv l n) }}} parallel_incr_locked lock #l
    {{{ m, RET #m; ⌜(n ≤ m)%Z⌝ }}}.
  Proof.
    (* exercise *)
  Admitted.

(** *grad exercise* (incr_both_spec):
  Here is a well-synchronized function that adds (1, 2) to a tuple. *)
  Definition incr_both (lock : val) : val := λ: "p",
    acquire lock;; "p" <- (Fst !"p" + #1, Snd !"p" + #2);;
    release lock.

(* State an invariant that says that l points to a tuple whose second
   element is always two times the first element, and prove that
   incr_both preserves it. *)

  Definition two_times_inv (l : loc) : iProp Σ := True%I. (* change this *)

  Lemma incr_both_spec lock l :
    {{{ is_lock lock (two_times_inv l) }}} incr_both lock #l {{{ RET #(); True }}}.
  Proof.
    (* exercise *)
  Admitted.

End proof.