(* modified from https://gitlab.mpi-sws.org/iris/tutorial-popl21/-/blob/master/exercises/ex_04_parallel_add.v *)
(**
In this exercise we prove that adding 1 to a number twice in parallel
actually increases it by 2, using ghost state.
*)
From iris.algebra Require Import excl_auth frac_auth numbers.
From iris.base_logic.lib Require Import invariants.
From iris.heap_lang Require Import lib.par proofmode notation.
From exercises Require Import ex_03_spinlock.

(* Previously we proved:

  Lemma parallel_incr_spec (l : loc) (n : Z):
    {{{ l ↦ #n }}} parallel_incr #l
    {{{ m, RET #m; ⌜(n ≤ m)%Z⌝ }}}.

but we couldn't actually prove that it returns n + 2.

*)

(* And in fact, it sometimes doesn't!

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

  (("l" <- !"l" + #1) ||| ("l" <- !"l" + #1)) -->
  (("l" <- n + #1) ||| ("l" <- !"l" + #1)) -->
  (("l" <- n + #1) ||| ("l" <- n + #1)) -->
  (* stores n + 1 twice *)
*)

(** Here are two well-synchronized versions of parallel_incr. We should be able
    to prove that both of them increase the value of the input by 2! *)
Definition parallel_incr_atomic : val := λ: "l",
    ((FAA "l" #1) ||| (FAA "l" #1));;
    !"l".

Definition parallel_incr_locked : val := λ: "l",
    let: "lock" := newlock #() in
    ((acquire "lock";; "l" <- !"l" + #1;; release "lock") |||
     (acquire "lock";; "l" <- !"l" + #1;; release "lock"));;
    acquire "lock";; let: "res" := !"l" in release "lock";; "res".


Section proof.
  Context `{!heapGS Σ, !spawnG Σ, !inG Σ (excl_authR ZO)}.

(** To prove that the program returns 2 exactly,
we need a piece of ghost state: integer ghost variables.

Whereas we previously abstracted over an arbitrary "ghost state" [Σ] in the
proofs, we now need to make sure that we can use integer ghost variables. For
this, we add the type class constraint:

  inG Σ (excl_authR ZO)

*)

(* Ghost variables (technically "exclusive authoritative" ghost state) break into
   two pieces: ●E and ◯E. *)

  (** Some helping lemmas for ghost state that we need in the proof. In actual
  proofs we tend to inline simple lemmas like these, but they are here to
  make things easier to understand. *)
  Lemma ghost_var_alloc n :
    ⊢ |==> ∃ γ, own γ (●E n) ∗ own γ (◯E n).
  Proof.
    iMod (own_alloc (●E n ⋅ ◯E n)) as (γ) "[??]".
    - by apply excl_auth_valid.
    - by eauto with iFrame.
  Qed.

  Lemma ghost_var_agree γ n m :
    own γ (●E n) -∗ own γ (◯E m) -∗ ⌜ n = m ⌝.
  Proof.
    iIntros "Hγ● Hγ◯".
    by iDestruct (own_valid_2 with "Hγ● Hγ◯") as %?%excl_auth_agree_L.
  Qed.

  Lemma ghost_var_update γ n' n m :
    own γ (●E n) -∗ own γ (◯E m) ==∗ own γ (●E n') ∗ own γ (◯E n').
  Proof.
    iIntros "Hγ● Hγ◯".
    iMod (own_update_2 _ _ _ (●E n' ⋅ ◯E n') with "Hγ● Hγ◯") as "[$$]".
    { by apply excl_auth_update. }
    done.
  Qed.



  Definition incr_inv_2 n (r : loc) (γ1 γ2 : gname) : iProp Σ :=
    (∃ n1 n2 : Z, r ↦ #(n + n1 + n2)
            ∗ own γ1 (●E n1) ∗ own γ2 (●E n2))%I.



  (** *Exercise*: finish the missing cases of the proof. *)
  Lemma parallel_incr_spec_2 l (n : Z) :
    {{{ l ↦ #n }}} parallel_incr_atomic #l {{{ RET #(n + 2); True }}}.
  Proof.
    iIntros (Φ) "Hl Post".
    unfold parallel_incr_atomic. wp_lam.
    iMod (ghost_var_alloc 0%Z) as (γ1) "[Hγ1● Hγ1◯]".
    iMod (ghost_var_alloc 0%Z) as (γ2) "[Hγ2● Hγ2◯]".
    iMod (inv_alloc nroot _ (incr_inv_2 n l γ1 γ2) with "[Hl Hγ1● Hγ2●]") as "#Hinv".
    { (* exercise *) admit. }
    wp_smart_apply (wp_par (fun _ => own γ1 (◯E 1%Z))
                           (fun _ => own γ2 (◯E 1%Z))
                with "[Hγ1◯] [Hγ2◯]").
    - iInv "Hinv" as "inv" "Hclose".
      unfold incr_inv_2 at 2.
      iMod "inv" as (n1 n2) "(Hl & Hγ1● & Hγ2●)".
      wp_faa.
      iDestruct (ghost_var_agree with "Hγ1● Hγ1◯") as %->.
      iMod (ghost_var_update γ1 1%Z with "Hγ1● Hγ1◯") as "[Hγ1● Hγ1◯]".
      iMod ("Hclose" with "[- Hγ1◯]"); [|auto].
      unfold incr_inv_2 at 2.
      iExists 1%Z, n2. iFrame "Hγ1● Hγ2●". replace (n + 1 + n2)%Z with (n + 0 + n2 + 1)%Z by lia.
      auto.
    - (* exercise *)
      admit.
    - (* exercise *)
      admit.
  Admitted.

  (* We should be able to prove the same spec for the locked version: *)
  Lemma parallel_incr_spec_2' l (n : Z) :
    {{{ l ↦ #n }}} parallel_incr_locked #l {{{ RET #(n + 2); True }}}.
  Proof.
    (* exercise *)
  Admitted.

End proof.



(** We can also do this using just one piece of ghost state that can
    be shared between threads. *)
Section proof3.
  Context `{!heapGS Σ, !spawnG Σ, !inG Σ (frac_authR natR)}.

  Definition incr_inv_3 (n : Z) (r : loc) (γ : gname) : iProp Σ :=
    (∃ m : nat, r ↦ #(n + m) ∗ own γ (●F m))%I.
  (* ●F is fractional ghost state, ◯ can be split between threads
     We only need the total contributions -- it doesn't matter which thread does them! *)

  (* Other kinds of ghost state: exclusive tokens, maps from keys to values,
     histories of operations, state machines for protocols *)

  (** *Exercise*: finish the missing cases of the proof. *)
  Lemma parallel_incr_spec_3 l (n : Z) :
    {{{ l ↦ #n }}} parallel_incr_atomic #l {{{ RET #(n + 2); True }}}.
  Proof.
    iIntros (Φ) "Hl Post".
    unfold parallel_incr_atomic. wp_lam.
    iMod (own_alloc (●F 0 ⋅ ◯F 0)) as (γ) "[Hγ● [Hγ1◯ Hγ2◯]]".
    { by apply auth_both_valid. }
    iMod (inv_alloc nroot _ (incr_inv_3 n l γ) with "[Hl Hγ●]") as "#Hinv".
    { (* exercise *) admit. }
    wp_smart_apply (wp_par (λ _, own γ (◯F{1/2} 1)) (λ _, own γ (◯F{1/2} 1))
                with "[Hγ1◯] [Hγ2◯]").
    - iInv "Hinv" as (m) ">[Hl Hγ●]" "Hclose".
      wp_faa.
      iMod (own_update_2 _ _ _ (●F (m+1) ⋅ ◯F{1/2}1) with "Hγ● Hγ1◯") as "[Hγ● Hγ1◯]".
      { rewrite (comm plus).
        by apply frac_auth_update, (op_local_update_discrete m 0 1). }
      iMod ("Hclose" with "[Hl Hγ●]"); [|by auto].
      iExists _. iFrame. replace (n + (m + 1)%nat)%Z with (n + m + 1)%Z by lia. auto.
    - (* exercise *)
      admit.
    - (* exercise *)
      admit.
  Admitted.

End proof3.