From iris.heap_lang Require Export proofmode notation.
From iris.heap_lang Require Import spin_lock assert par.
From iris.algebra Require Import frac auth.
From iris_c.lib Require Import locking_heap mset flock U.
From iris_c.c_translation Require Import monad lifting proofmode.

Definition a_alloc : val := λ: "x",
  a_bind (λ: "v", a_atomic_env (λ: <>, ref "v")) "x".

Definition a_store : val := λ: "x1" "x2",
  a_bind (λ: "vv",
    a_atomic_env (λ: "env",
      mset_add (Fst "vv") "env" ;;
      Fst "vv" <- Snd "vv" ;;
      Snd "vv"
    )
  ) (a_par "x1" "x2").

Definition a_load : val := λ: "x",
  a_bind (λ: "v",
    a_atomic_env (λ: "env",
      assert: (mset_member "v" "env" = #false);;
      !"v"
    )
  ) "x".

Definition a_un_op (op : un_op) : val := λ: "x",
  a_bind (λ: "v", a_ret (λ: <>, UnOp op "v")) "x".

Definition a_bin_op (op : bin_op) : val := λ: "x1" "x2",
  a_bind (λ: "vv",
    a_ret (BinOp op (Fst "vv") (Snd "vv"))
  ) (a_par "x1" "x2").

(* M () *)
(* The eta expansion is used to turn it into a value *)
Definition a_seq : val := λ: <>,
  a_atomic_env (λ: "env", mset_clear "env").

Definition a_seq_bind : val := λ: "e1" "e2",
  a_bind (λ: "x", a_seq #();;; "e2" "x") "e1".

Notation "e1 ;;;; e2" :=
  (a_seq_bind e1 (λ: <>, e2))%E (at level 80, right associativity).

Definition a_if : val := λ: "cnd" "e1" "e2",
  a_bind (λ: "c", if: "c" then "e1" else "e2") "cnd".

Definition a_while: val :=
  rec: "while" "cnd" "bdy" :=
    a_if "cnd" ("bdy" ;;;; "while" "cnd" "bdy") (a_ret #()).

Section proofs.
  Context `{locking_heapG Σ, heapG Σ, flockG Σ, spawnG Σ}.

  Lemma a_seq_spec R Φ :
    U (Φ #()) -∗
    awp (a_seq #()) R Φ.
  Proof.
    iIntros "HΦ". rewrite /a_seq. awp_lam.
    iApply awp_atomic_env.
    iIntros (env) "Henv HR".
    iApply wp_fupd.
    rewrite {2}/env_inv.
    iDestruct "Henv" as (X σ) "(HX & Hσ & Hls & Hlocks)".
    iDestruct "Hlocks" as %Hlocks.
    wp_let. iApply (mset_clear_spec with "HX").
    iNext. iIntros "HX".
    iDestruct "HΦ" as (us) "[Hus HΦ]".
    clear Hlocks.
    iInduction us as [|u us] "IH" forall (σ); simpl.
    - iModIntro. iFrame "HR". iSplitR "HΦ".
      + iExists ∅, σ. iFrame. iPureIntro.
        rewrite /correct_locks /set_Forall. set_solver.
      + by iApply "HΦ".
    - iDestruct "Hus" as "[Hu Hus]".
      iAssert (⌜σ !! u.1 = Some LLvl⌝%I) with "[Hσ Hu]" as %?.
      { rewrite mapsto_eq /mapsto_def.
        iDestruct "Hu" as "[Hu Hl]".
          by iDestruct (own_valid_2 with "Hσ Hl")
            as %[?%heap_singleton_included _]%auth_valid_discrete_2. }
      iMod (locking_heap_change_lock _ _ _ ULvl with "Hσ Hu") as "[Hσ Hu]".
      iApply ("IH" with "Hus [HΦ Hu] Hσ [Hls] HR HX").
      { iIntros "Hus". iApply "HΦ". by iFrame. }
      { rewrite -bi.big_sepM_insert_override; eauto. }
  Qed.

  Lemma a_sequence_spec R Φ (e1 e2 : expr) :
    AsVal e2 →
    awp e1 R (λ v, U (awp (e2 v) R Φ)) -∗
    awp (a_seq_bind e1 e2) R Φ.
  Proof.
    iIntros ([v2 <-%of_to_val]) "H".
    awp_apply (a_wp_awp with "H"); iIntros (v) "H".
    rewrite /a_seq_bind /=. do 2 awp_lam.
    iApply awp_bind. iApply (awp_wand with "H"). iIntros (w) "H".
    awp_lam. iApply awp_bind. iApply a_seq_spec.
    iModIntro. by awp_lam.
  Qed.

  Lemma a_load_spec R (l : loc) (v : val) Φ :
    l ↦U v -∗
    (l ↦U v -∗ Φ v) -∗
    awp (a_load (a_ret #l)) R Φ.
  Proof.
    unfold a_load. iIntros "Hv HΦ".
    rewrite /a_ret.
    do 2 awp_lam.
    iApply awp_bind.
    iApply awp_value.
    awp_let.
    iApply awp_atomic_env.
    iIntros (env) "Henv HR".
    rewrite {2}/env_inv.
    iDestruct "Henv" as (X σ) "(HX & Hσ & Hls & Hlocks)".
    iDestruct "Hlocks" as %Hlocks.
    iDestruct (locked_locs_unlocked with "Hv Hσ") as %Hl.
    assert (#l ∉ X).
    { unfold correct_locks in *. intros Hx. apply Hl.
      destruct (Hlocks _ Hx) as [l' [? Hl']]. by simplify_eq/=. }
    wp_let.
    wp_apply wp_assert.
    wp_apply (mset_member_spec #l env with "HX").
    iIntros "Henv /=". case_decide; first by exfalso. simpl.
    wp_op. iSplit; eauto. iNext. wp_seq.
    rewrite mapsto_eq /mapsto_def. iDestruct "Hv" as "[Hv Hl]". wp_load.
    iCombine "Hv Hl" as "Hv". iFrame "HR".
    iSplitR "HΦ Hv".
    - iExists X,σ. by iFrame.
    - by iApply "HΦ".
  Qed.

  Lemma a_alloc_spec R (ev : expr) (v : val) Φ :
    IntoVal ev v →
    (∀ l, l ↦U v -∗ Φ #l) -∗
    awp (a_alloc (a_ret ev)) R Φ.
  Proof.
    intros <-%of_to_val.
    unfold a_alloc. iIntros "HΦ".
    rewrite /a_ret.
    do 2 awp_lam.
    iApply awp_bind. iApply awp_value.
    awp_let.
    iApply awp_atomic_env.
    iIntros (env) "Henv HR".
    rewrite {2}/env_inv.
    iDestruct "Henv" as (X σ) "(HX & Hσ & Hls & Hlocks)".
    iDestruct "Hlocks" as %Hlocks.
    iApply wp_fupd.
    wp_let. wp_alloc l as "Hl".
    iAssert (⌜σ !! l = None⌝)%I with "[Hl Hls]" as %Hl.
    { remember (σ !! l) as σl. destruct σl; simplify_eq; eauto.
      iExFalso. rewrite (bi.big_sepM_lookup _ σ l _); last eauto.
      iDestruct "Hls" as (v') "Hl'".
      by iDestruct (mapsto_valid_2 l with "Hl Hl'") as %[]. }
    iDestruct "Hl" as "[Hl Hl']".
    iMod (locking_heap_alloc σ l ULvl v with "Hl' Hσ") as "[Hσ Hl']"; eauto.
    iModIntro. iFrame "HR". iSplitR "HΦ Hl'".
    - iExists X,(<[l:=ULvl]>σ). iFrame. iSplit.
      + rewrite bi.big_sepM_insert; eauto. iFrame. eauto.
      + iPureIntro. by rewrite locked_locs_alloc_unlocked.
    - by iApply "HΦ".
  Qed.

  Lemma a_store_spec R Φ Ψ1 Ψ2 e1 e2 :
    awp e1 R (λ v, ∃ l : loc, ⌜v = #l⌝ ∧ Ψ1 l)-∗
    awp e2 R Ψ2 -∗
    (∀ (l : loc) w,
      Ψ1 l -∗ Ψ2 w -∗ (∃ v, l ↦U v ∗ (l ↦L w -∗ Φ w))) -∗
    awp (a_store e1 e2) R Φ.
  Proof.
    iIntros "H1 H2 HΦ".
    awp_apply (a_wp_awp with "H1"); iIntros (v1) "H1". awp_lam.
    awp_apply (a_wp_awp with "H2"); iIntros (v2) "H2". awp_lam.
    iApply awp_bind. iApply (awp_par with "H1 H2").
    iIntros "!>" (w1 w2) "H1 H2 !>". iDestruct "H1" as (l ->) "H1". awp_lam.
    iDestruct ("HΦ" with "H1 H2") as (v) "[Hv HΦ]".
    iApply awp_atomic_env.
    iIntros (env) "Henv HR".
    rewrite {2}/env_inv.
    iDestruct "Henv" as (X σ) "(HX & Hσ & Hls & Hlocks)".
    iDestruct "Hlocks" as %Hlocks.
    iDestruct (locked_locs_unlocked with "Hv Hσ") as %Hl.
    assert (#l ∉ X).
    { unfold correct_locks in *. intros Hx. apply Hl.
      destruct (Hlocks _ Hx) as [l' [? Hl']]. by simplify_eq/=. }
    wp_let. wp_proj.
    wp_apply (mset_add_spec with "[$HX]"); eauto.
    iIntros "HX". wp_seq.
    iAssert (⌜σ !! l = Some ULvl⌝%I) with "[Hσ Hv]" as %?.
    { rewrite mapsto_eq /mapsto_def.
      iDestruct "Hv" as "[Hv Hl]".
      by iDestruct (own_valid_2 with "Hσ Hl")
        as %[?%heap_singleton_included _]%auth_valid_discrete_2. }
    iMod (locking_heap_change_lock _ _ ULvl LLvl with "Hσ Hv") as "[Hσ Hv]".
    do 2 wp_proj. rewrite mapsto_eq /mapsto_def.
    iDestruct "Hv" as "[Hv Hl]".
    rewrite (bi.big_sepM_lookup_acc _ _ l ULvl); last done.
    iDestruct "Hls" as "[Hl' Hls]".
    iDestruct "Hl'" as (?) "Hl'".
    iDestruct (mapsto_agree l (1/2) (1/2) with "Hl' Hv") as %->.
    iCombine "Hv Hl'" as "Hv".
    wp_store.
    iDestruct "Hv" as "[Hv Hl']".
    iSpecialize ("Hls" with "[Hl']"); eauto.
    wp_proj. iFrame "HR". iSplitR "HΦ Hl Hv".
    - iExists ({[#l]} ∪ X),(<[l:=LLvl]> σ). iFrame. iSplitL.
      + rewrite -bi.big_sepM_insert_override; eauto.
      + (* TODO: a separate lemma somewhere *)
        iPureIntro. rewrite locked_locs_lock.
        revert Hlocks. rewrite /correct_locks /set_Forall. set_solver.
    - iApply "HΦ". iFrame.
  Qed.

  Lemma a_if_true_spec R (e1 e2 : val) Φ :
     awp e1 R Φ -∗ awp (a_if (a_ret #true) e1 e2) R Φ.
  Proof.
    iIntros "HΦ".
    do 4 awp_lam.
    iApply awp_bind.
    iApply awp_value.
    awp_lam. by awp_if_true.
  Qed.

  Lemma a_if_false_spec R (e1 e2 : val) Φ :
     awp e2 R Φ -∗ awp (a_if (a_ret #false) e1 e2) R Φ.
  Proof.
    iIntros "HΦ".
    do 4 awp_lam.
    iApply awp_bind.
    iApply awp_value.
    awp_lam. by awp_if_false.
  Qed.

  Lemma a_if_spec R (e e1 e2 : val) Φ:
    awp e R (λ cnd, (⌜cnd = #true⌝ ∧ awp e1 R Φ) ∨ (⌜cnd = #false⌝ ∧ awp e2 R Φ)) -∗
    awp (a_if e e1 e2) R Φ.
  Proof.
    iIntros "HΦ". do 3 awp_lam.
    iApply awp_bind. iApply (awp_wand with "HΦ").
    iIntros (v) "[[% H] | [% H]]"; subst; by repeat awp_pure _.
  Qed.


End proofs.