From iris_c.c_translation Require Export monad.
From iris_c.lib Require Export U.
From iris_c.c_translation Require Export proofmode.
From iris.heap_lang Require Export assert.
From iris_c.lib Require Import mset flock list.

Notation "♯ l" := (c_ret (LitV l%Z%V)) (at level 8, format "♯ l").
Notation "♯ₗ l" := (c_ret (cloc_to_val l)) (at level 8, format "♯ₗ l") : expr_scope.

Definition c_alloc : val := λ: "x1" "x2",
  "vv" ←ᶜ "x1" |||ᶜ "x2" ;;ᶜ
  let: "n" := Fst "vv" in
  let: "v" := Snd "vv" in
  assert: (#0 < "n");;
  c_atomic_env (λ: <>, SOME (ref (SOME (#true, lreplicate "n" "v")), #0)).
Notation "'allocᶜ' ( e1 , e2 )" :=
  (c_alloc e1%E e2%E) (at level 10, e1, e2 at level 99) : expr_scope.

Definition c_free_check : val :=
  rec: "go" "env" "l" "n" :=
  if: "n" = #0 then #() else
  let: "n" := "n" - #1 in
  assert: (mset_member ("l", "n") "env" = #false);;
  "go" "env" "l" "n".

Definition c_free : val := λ: "x",
  "v" ←ᶜ "x";;ᶜ
  c_atomic_env (λ: "env",
    match: "v" with
      NONE => assert: #false (* null pointer *)
    | SOME "li" =>
       let: "l" := Fst "li" in
       let: "i" := Snd "li" in
       match: !"l" with
         NONE => assert: #false (* location already freed, double free *)
       | SOME "kll" =>
          let: "k" := Fst "kll" in
          let: "ll" := Snd "kll" in
          (* Make sure its not a block scoped variable *)
          assert: ("k" = #true);;
          (* We need to make sure `i = 0` and that all `0 ... length of block`
          are unlocked. *)
          assert: ("i" = #0);;
          let: "n" := llength "ll" in
          c_free_check "env" "l" "n";;
         "l" <- NONE
       end
    end
  ).
Notation "'freeᶜ' ( e )" :=
  (c_free e%E) (at level 10, e at level 99) : expr_scope.

Definition c_store : val := λ: "x1" "x2",
  "vv" ←ᶜ "x1" |||ᶜ "x2" ;;ᶜ
  c_atomic_env (λ: "env",
    match: Fst "vv" with
      NONE => assert: #false (* null pointer *)
    | SOME "li" =>
       let: "l" := Fst "li" in
       let: "i" := Snd "li" in
       let: "v" := Snd "vv" in
       mset_add ("l", "i") "env" ;;
       match: !"l" with
         NONE => assert: #false (* store after free *)
       | SOME "kll" => "l" <- SOME (Fst "kll", linsert "i" "v" (Snd "kll")) ;; "v"
       end
    end
  ).
Notation "e1 =ᶜ e2" := (c_store e1%E e2%E) (at level 80) : expr_scope.

Definition c_load : val := λ: "x",
  "v" ←ᶜ "x";;ᶜ
  c_atomic_env (λ: "env",
    match: "v" with
      NONE => assert: #false (* null pointer *)
    | SOME "li" =>
       let: "l" := Fst "li" in
       let: "i" := Snd "li" in
       assert: (mset_member ("l", "i") "env" = #false);;
       match: !"l" with
         NONE => assert: #false (* load after free *)
       | SOME "kll" => llookup "i" (Snd "kll")
       end
    end
  ).
Notation "∗ᶜ e" :=
  (c_load e)%E (at level 20, right associativity) : expr_scope.

Notation "'skipᶜ'" := (c_ret #()).

Definition c_seq_bind : val := λ: "x" "f",
  "a" ←ᶜ "x" ;;ᶜ
  c_atomic_env (λ: "env", mset_clear "env") ;;ᶜ
  "f" "a".
Notation "x ←ᶜ e1 ;ᶜ e2" :=
  (c_seq_bind e1%E (λ: x, e2)%E)%E
  (at level 100, e1 at next level, e2 at level 200, right associativity,
   format "'[' x  ←ᶜ  '[' e1 ']' ;ᶜ  '/' e2 ']'") : expr_scope.
Notation "e1 ;ᶜ e2" :=
  (c_seq_bind e1%E (λ: <>, e2)%E)%E
  (at level 100, e2 at level 200,
   format "'[' '[hv' '[' e1 ']'  ;ᶜ ']'  '/' e2 ']'") : expr_scope.

Definition c_mut_bind : val := λ: "x" "f",
  "v" ←ᶜ "x" ;ᶜ
  "l" ←ᶜ c_atomic_env (λ: <>, ref (SOME (#false, lreplicate #1 "v"))) ;;ᶜ
  "b" ←ᶜ "f" (SOME ("l", #0)) ;;ᶜ
  c_atomic_env (λ: <>, "l" <- NONE) ;;ᶜ
  c_ret "b".
Notation "x ←mutᶜ e1 ;ᶜ e2" :=
  (c_mut_bind e1 (λ: x, e2))%E
  (at level 100, e1 at next level, e2 at level 200, right associativity,
   format "'[' x  ←mutᶜ  '[' e1 ']' ;ᶜ  '/' e2 ']'") : expr_scope.

Definition c_if : val := λ: "cnd" "e1" "e2",
  (* sequenced binds needed here; there should be sequence point after the conditional *)
  "c" ←ᶜ "cnd" ;ᶜ
  if: "c" then "e1" #() else "e2" #().
Notation "'ifᶜ' ( cnd ) { e1 } 'elseᶜ' { e2 }" :=
  (c_if cnd%E (λ: <>, e1)%E (λ: <>, e2)%E)
  (at level 10, cnd, e1, e2 at level 99,
   format "'[v' 'ifᶜ'  ( cnd )  {  '/  ' '[' e1 ']'  '/' }  'elseᶜ'  {  '/  ' '[' e2 ']'  '/' } ']'") : expr_scope.

Definition c_while: val :=
  rec: "while" "cnd" "bdy" :=
    ifᶜ ("cnd" #()) { "bdy" #() ;ᶜ "while" "cnd" "bdy" }
    elseᶜ { skipᶜ }.
Notation "'whileᶜ' ( cnd ) { e }" := (c_while (λ: <>, cnd)%E (λ: <>, e)%E)
  (at level 10, cnd, e at level 99,
   format "'[v' 'whileᶜ'  ( cnd )  {  '/  ' '[' e ']'  '/' } ']'") : expr_scope.

(* A version of while with value lambdas, this is an artifact because of the way
heap_lang works in Coq *)
Notation "'whileVᶜ' ( cnd ) { e }" := (c_while (LamV <> cnd) (LamV <> e))
  (at level 10, cnd, e at level 99,
   format "'[v' 'whileVᶜ'  ( cnd )  {  '/  ' '[' e ']' '/'  } ']'") : expr_scope.

Definition c_fun (f : expr) : val := λ: "arg",
  (* sequence point at the end of a function *)
  "v" ←ᶜ f "arg" ;ᶜ
  c_ret "v".
(* TODO: Similar notation for recursive functions *)
Notation "'λᶜ' x , e" := (c_fun (λ: x, e)%V)
  (at level 200, x at level 1, e at level 200,
   format "'[' 'λᶜ'  x ,  '/  ' e ']'") : val_scope.

Definition c_call : val := λ: "f" "arg",
  (* sequence point before a function call *)
  "fa" ←ᶜ ("f" |||ᶜ "arg");ᶜ
  c_atomic (λ: <>, (Fst "fa") (Snd "fa")).
Notation "'callᶜ' f a" :=
  (c_call f a)%E
  (at level 10, f, a at level 9,
   format "'callᶜ'  f  a") : expr_scope.

Definition c_un_op (op : un_op) : val := λ: "x",
  "v" ←ᶜ "x" ;;ᶜ c_ret (UnOp op "v").

Inductive cbin_op :=
  | CBinOp : bin_op → cbin_op
  | PtrPlusOp : cbin_op
  | PtrLtOp : cbin_op.

Definition c_ptr_plus : val := λ: "x" "y",
  (* all binds should be non-sequenced *)
  "vv" ←ᶜ ("x" |||ᶜ "y");;ᶜ
  match: Fst "vv" with
    NONE => assert #false (* null pointer *)
  | SOME "li" => c_ret (SOME (Fst "li", Snd "vv" + Snd "li"))
  end.

Definition c_ptr_lt : val := λ: "x" "y",
  (* all binds should be non-sequenced *)
  "pq" ←ᶜ ("x" |||ᶜ "y");;ᶜ
  match: Fst "pq" with
    NONE => assert #false (* null pointer *)
  | SOME "p" =>
     match: Snd "pq" with
       NONE => assert #false (* null pointer *)
     | SOME "q" =>
        if: Fst "p" = Fst "q"
        then c_ret (Snd "p" < Snd "q") else c_ret #false
     end
  end.

Definition c_bin_op (op : cbin_op) : val :=
  (* all binds should be non-sequenced *)
  match op with
  | CBinOp op' => λ: "x1" "x2",
     "vv" ←ᶜ "x1" |||ᶜ "x2" ;;ᶜ
     c_ret (BinOp op' (Fst "vv") (Snd "vv"))
  | PtrPlusOp => c_ptr_plus
  | PtrLtOp   => c_ptr_lt
  end.
Notation "e1 +ᶜ e2" := (c_bin_op (CBinOp PlusOp) e1%E e2%E) (at level 50) : expr_scope.
Notation "e1 -ᶜ e2" := (c_bin_op (CBinOp MinusOp) e1%E e2%E) (at level 35) : expr_scope.
Notation "e1 *ᶜ e2" := (c_bin_op (CBinOp MultOp) e1%E e2%E) (at level 40) : expr_scope.
Notation "e1 ≤ᶜ e2" := (c_bin_op (CBinOp LeOp) e1%E e2%E) (at level 70) : expr_scope.
Notation "e1 <ᶜ e2" := (c_bin_op (CBinOp LtOp) e1%E e2%E) (at level 70) : expr_scope.
Notation "e1 ==ᶜ e2" := (c_bin_op (CBinOp EqOp) e1%E e2%E) (at level 70) : expr_scope.
Notation "e1 !=ᶜ e2" := (c_un_op NegOp (c_bin_op (CBinOp EqOp) e1%E e2%E)) (at level 70): expr_scope.
Notation "~ᶜ e" := (c_un_op NegOp e%E) (at level 20, right associativity) : expr_scope.

Notation "e1 +∗ᶜ e2" := (c_bin_op PtrPlusOp e1%E e2%E) (at level 50) : expr_scope.
Notation "e1 <∗ᶜ e2" := (c_bin_op PtrLtOp e1%E e2%E)  (at level 70) : expr_scope.

Definition int_of_val (v : val) : option Z :=
  match v with LitV (LitInt x) => Some x | _ => None end.

Definition cbin_op_eval (op : cbin_op) (v1 v2 : val) : option val :=
  match op with
  | CBinOp op' => bin_op_eval op' v1 v2
  | PtrPlusOp =>
     cl ← cloc_of_val v1;
     o  ← int_of_val v2;
     Some (cloc_to_val (cl +∗ o))
  | PtrLtOp =>
     cl1 ← cloc_of_val v1;
     cl2 ← cloc_of_val v2;
     Some #(cloc_lt cl1 cl2)
  end.

Definition c_pre_bin_op (op : cbin_op) : val := λ: "x" "y",
  (* all binds should be non-sequenced *)
  "lv" ←ᶜ ("x" |||ᶜ "y");;ᶜ
  c_atomic (λ: <>,
    "ov" ←ᶜ ∗ᶜ (c_ret (Fst "lv"));;ᶜ
    c_ret (Fst "lv") =ᶜ c_bin_op op (c_ret "ov") (c_ret (Snd "lv"));;ᶜ
    c_ret "ov").
Notation "e1 +=ᶜ e2" := (c_pre_bin_op (CBinOp PlusOp) e1%E e2%E) (at level 80) : expr_scope.
Notation "e1 +∗=ᶜ e2" := (c_pre_bin_op PtrPlusOp e1%E e2%E) (at level 80) : expr_scope.

Section proofs.
  Context `{cmonadG Σ}.

  Lemma cwp_alloc R Φ Ψ1 Ψ2 e1 e2 :
    CWP e1 @ R {{ Ψ1 }} -∗
    CWP e2 @ R {{ Ψ2 }} -∗
    (∀ v1 v2, Ψ1 v1 -∗ Ψ2 v2 -∗ ∃ n : nat,
              ⌜ v1 = #n ⌝ ∧ ⌜ n ≠ 0%nat ⌝ ∧
              ∀ cl, block_info cl true n -∗ cl ↦C∗ replicate n v2 -∗ Φ (cloc_to_val cl)) -∗
    CWP allocᶜ(e1, e2) @ R {{ Φ }}.
  Proof.
    iIntros "H1 H2 HΦ".
    cwp_apply (cwp_wp with "H2"); iIntros (v2) "H2".
    cwp_apply (cwp_wp with "H1"); iIntros (v1) "H1".
    cwp_lam. cwp_pures.
    iApply cwp_bind. iApply (cwp_par with "H1 H2").
    iIntros "!>" (w1 w2) "H1 H2 !>". cwp_pures.
    iDestruct ("HΦ" with "H1 H2") as (n -> ?) "HΦ".
    cwp_apply wp_assert. wp_op. rewrite bool_decide_true; last lia.
    iSplit; first done. iNext. cwp_pures.
    iApply cwp_atomic_env.
    iIntros (env). iDestruct 1 as (X σ HX) "[Hlocks Hσ]". iIntros "$". wp_pures.
    wp_apply (lreplicate_spec with "[//]"); iIntros (ll Hll).
    wp_alloc l as "Hl".
    iMod (full_locking_heap_alloc_upd with "Hσ Hl")
      as (?) "(Hσ & Hinfo & Hl)"; first done.
    { by destruct n. }
    wp_pures. iIntros "!>". rewrite cloc_to_val_eq replicate_length.
    iSplitL "Hlocks Hσ".
    - iExists X, _. iFrame.
      iIntros "!%". intros w Hw. destruct (HX _ Hw) as (cl&Hcl&Hin).
      exists cl; split; first done. by rewrite locked_locs_alloc_heap.
    - iApply ("HΦ" $! (CLoc l 0) with "Hinfo Hl").
  Qed.

  Lemma cwp_free R Φ e :
    CWP e @ R {{ v, ∃ cl ws, ⌜ v = cloc_to_val cl ⌝ ∧
                             block_info cl true (length ws) ∗ cl ↦C∗ ws ∗ Φ #() }} -∗
    CWP freeᶜ(e) @ R {{ Φ }}.
  Proof.
    iIntros "H".
    cwp_apply (cwp_wp with "H"); iIntros (v) "H". cwp_lam. cwp_pures.
    iApply cwp_bind. iApply (cwp_wand with "H"). clear v.
    iIntros (v). iDestruct 1 as (cl ws ->) "(Hinfo & Hcl & HΦ)". cwp_pures.
    iApply cwp_atomic_env. iIntros (env) "Henv HR". wp_pures.
    rewrite cloc_to_val_eq. wp_pures.
    iDestruct "Henv" as (X σ HX) "[Hlocks Hσ]".
    iAssert ⌜ ∀ i : nat, is_Some (ws !! i) → (cl +∗ i) ∉ locked_locs σ⌝%I as %Hlocked.
    { iIntros (i [w Hi]). rewrite /mapsto_list.
      iDestruct (big_sepL_lookup with "Hcl") as "H"; first done.
      by iApply (full_locking_heap_unlocked with "[$]"). }
    iMod (full_locking_heap_free_upd with "Hσ Hinfo Hcl")
      as (ll Hoff Hl) "[Hl Hclose]".
    wp_load. wp_pures. rewrite Hoff.
    wp_apply wp_assert; wp_pures; iSplit; first done. iNext.
    wp_apply wp_assert; wp_pures; iSplit; first done. iNext.
    wp_pures. wp_apply (llength_spec with "[//]"); iIntros "_"; wp_pures.
    iAssert (∀ Ψ (n : nat), ⌜ n ≤ length ws ⌝ →
      (is_mset env X -∗ Ψ #()) -∗ WP c_free_check env #(cloc_base cl) #n {{ Ψ }})%I
      with "[Hlocks]" as "Hcheck".
    { iIntros (Ψ n Hn) "HΨ". iInduction n as [|n] "IH" forall (Ψ Hn).
      { wp_lam; wp_pures. by iApply "HΨ". }
      wp_lam; wp_pures. wp_apply wp_assert.
      rewrite Nat2Z.inj_succ Z.sub_1_r Z.pred_succ.
      wp_apply (mset_member_spec with "Hlocks"); iIntros "Hlocks"; case_bool_decide.
      { destruct (HX (#(cloc_base cl), #n)%V) as (cl'&[= <-]&?); first done.
        destruct (Hlocked n); first by (apply lookup_lt_is_Some_2; lia).
        destruct cl; simplify_eq/=. by rewrite /cloc_plus /= Z.add_0_r. }
      wp_op. iSplit; first done. iNext; wp_pures.
      iApply ("IH" with "[%] Hlocks HΨ"). lia. }
    wp_apply ("Hcheck" with "[//]"); iIntros "Hlock". wp_pures. wp_store.
    iIntros "!> {$HΦ $HR}". iExists X, _.
    iFrame "Hlock". iSplit; last by iApply "Hclose".
    iPureIntro; intros w Hw. destruct (HX _ Hw) as (cl'&Hcl&Hin).
    exists cl';  split; first done. apply locked_locs_free_heap; first done.
    intros (?&?&?). destruct (Hlocked (Z.to_nat (cloc_offset cl'))).
    { apply lookup_lt_is_Some_2, Nat2Z.inj_lt. rewrite Z2Nat.id; lia. }
    destruct cl, cl'; simplify_eq/=.
    by rewrite /cloc_plus /= Z.add_0_r Z2Nat.id; last lia.
  Qed.

  Lemma cwp_store R Φ Ψ1 Ψ2 e1 e2 :
    CWP e1 @ R {{ Ψ1 }} -∗
    CWP e2 @ R {{ Ψ2 }} -∗
    (∀ v1 v2, Ψ1 v1 -∗ Ψ2 v2 -∗ ∃ cl w,
              ⌜ v1 = cloc_to_val cl ⌝ ∧ cl ↦C w ∗ (cl ↦C[LLvl] v2 -∗ Φ v2)) -∗
    CWP e1 =ᶜ e2 @ R {{ Φ }}.
  Proof.
    iIntros "H1 H2 HΦ".
    cwp_apply (cwp_wp with "H2"); iIntros (v2) "H2".
    cwp_apply (cwp_wp with "H1"); iIntros (v1) "H1".
    cwp_lam. cwp_pures.
    iApply cwp_bind. iApply (cwp_par with "H1 H2").
    iIntros "!>" (w1 w2) "H1 H2 !>". cwp_pures.
    iDestruct ("HΦ" with "H1 H2") as (cl w ->) "[Hl HΦ]".
    iApply cwp_atomic_env.
    iIntros (env). iDestruct 1 as (X σ HX) "[Hlocks Hσ]". iIntros "HR".
    iDestruct (full_locking_heap_unlocked with "Hl Hσ") as %Hw1.
    iDestruct (mapsto_offset_non_neg with "Hl") as %?.
    assert ((#(cloc_base cl), #(cloc_offset cl))%V ∉ X) as HclX.
    { intros Hcl. destruct (HX _ Hcl) as (cl'&[=]%cloc_to_of_val&?). naive_solver. }
    iMod (full_locking_heap_store_upd with "Hσ Hl") as (k ll vs Hl Hi) "[Hl Hclose]".
    wp_pures. rewrite cloc_to_val_eq. wp_pures.
    wp_apply (mset_add_spec with "[$]"); first done.
    iIntros "Hlocks /="; wp_pures.
    wp_load; wp_pures.
    iEval (rewrite -(Z2Nat.id (cloc_offset cl)) //).
    wp_apply (linsert_spec with "[//]"); [eauto|]. iIntros (ll' Hl').
    iApply wp_fupd. wp_store.
    iMod ("Hclose" $! _ LLvl with "[//] Hl") as "[Hσ Hl]".
    iIntros "!> !> {$HR}". iSplitL "Hlocks Hσ"; last by iApply "HΦ".
    iExists ({[(#(cloc_base cl), #(cloc_offset cl))%V]} ∪ X), _.
    iFrame "Hσ Hlocks". iPureIntro. rewrite locked_locs_lock. set_solver.
  Qed.

  Lemma cwp_load_exists_frac R Φ e :
    CWP e @ R {{ v, ∃ cl q w, ⌜ v = cloc_to_val cl ⌝ ∧
                              cl ↦C{q} w ∗ (cl ↦C{q} w -∗ Φ w) }} -∗
    CWP ∗ᶜe @ R {{ Φ }}.
  Proof.
    iIntros "H".
    cwp_apply (cwp_wp with "H"); iIntros (v) "H". cwp_lam. cwp_pures.
    iApply cwp_bind. iApply (cwp_wand with "H"). clear v.
    iIntros (v). iDestruct 1 as (cl q w ->) "[Hl HΦ]". cwp_pures.
    iDestruct (mapsto_offset_non_neg with "Hl") as %?.
    iApply cwp_atomic_env. iIntros (env) "Henv HR".
    iDestruct "Henv" as (X σ HX) "[Hlocks Hσ]".
    iDestruct (full_locking_heap_unlocked with "Hl Hσ") as %Hv.
    assert ((#(cloc_base cl), #(cloc_offset cl))%V ∉ X) as HclX.
    { intros Hcl. destruct (HX _ Hcl) as (?&[=]%cloc_to_of_val&?); naive_solver. }
    iMod (full_locking_heap_load_upd with "Hσ Hl") as (k ll vs Hl Hi) "[Hl Hclose]".
    wp_pures. rewrite cloc_to_val_eq. wp_pures.
    wp_apply wp_assert. wp_apply (mset_member_spec with "Hlocks"); iIntros "Hlocks /=".
    rewrite bool_decide_false //.
    wp_op. iSplit; first done. iNext; wp_seq.
    wp_load; wp_match.
    iEval (rewrite -(Z2Nat.id (cloc_offset cl)) //).
    wp_apply (llookup_spec with "[//]"); [done|]; iIntros "_".
    iDestruct ("Hclose" with "Hl") as "[Hσ Hl]".
    iIntros "!> {$HR}". iSplitL "Hlocks Hσ"; last by iApply "HΦ".
    iExists X, _. by iFrame.
  Qed.

  Lemma cwp_load R Φ q e :
    CWP e @ R {{ v, ∃ cl w, ⌜ v = cloc_to_val cl ⌝ ∧
                            cl ↦C{q} w ∗ (cl ↦C{q} w -∗ Φ w) }} -∗
    CWP ∗ᶜe @ R {{ Φ }}.
  Proof.
    iIntros "H". iApply cwp_load_exists_frac.
    cwp_apply (cwp_wand with "H").
    iIntros (v). iDestruct 1 as (cl w ->) "[H1 H2]"; eauto with iFrame.
  Qed.

  (* Internal spec: breaks the abstraction/notation *)
  Lemma cwp_seq_bind' R Φ e (f: val) :
    CWP e @ R {{ v, U (CWP f v @ R {{ Φ }}) }} -∗
    CWP c_seq_bind e f @ R {{ Φ }}.
  Proof.
    iIntros "H".
    cwp_apply (cwp_wp with "H"); iIntros (v) "H". cwp_lam. cwp_pures.
    iApply cwp_bind. iApply (cwp_wand with "H"). iIntros (w) "H".
    cwp_pures. iApply cwp_bind.
    iApply cwp_atomic_env. iIntros (env) "Henv $". iApply wp_fupd.
    iDestruct "Henv" as (X σ _) "[Hlocks Hσ]".
    wp_lam. wp_apply (mset_clear_spec with "Hlocks"); iIntros "Hlocks".
    rewrite U_eq. iDestruct "H" as (us) "[Hus H]".
    iInduction us as [|[ul [uq uv]] us] "IH" forall (σ); simpl.
    - iModIntro. iSplitR "H".
      + iExists ∅, σ. by iFrame.
      + iNext. cwp_lam. by iApply "H".
    - iDestruct "Hus" as "[Hu Hus]".
      iDestruct (full_locking_heap_locked_present with "Hu Hσ") as %[z Hz].
      iMod (full_locking_heap_unlock with "Hσ Hu") as "[Hσ Hu]".
      iApply ("IH" with "Hus [H Hu] Hσ Hlocks").
      iIntros "Hus". iApply "H". by iFrame.
  Qed.

  Lemma cwp_seq_bind R Φ x e1 e2 :
    CWP e1 @ R {{ v, U (CWP subst' x v e2 @ R {{ Φ }}) }} -∗
    CWP x ←ᶜ e1 ;ᶜ e2 @ R {{ Φ }}.
  Proof.
    iIntros "H". cwp_pures. iApply cwp_seq_bind'.
    iApply (cwp_wand with "H"); iIntros (v) "H !>". by cwp_lam.
  Qed.

  (* Internal spec: breaks the abstraction/notation *)
  Lemma cwp_mut_bind' R Φ e (f: val) :
    CWP e @ R {{ v, U (∀ cl,
      cl ↦C v -∗
      CWP f (cloc_to_val cl) @ R {{ w, ∃ v', cl ↦C v' ∗ Φ w }}) }} -∗
    CWP c_mut_bind e f @ R {{ Φ }}.
  Proof.
    iIntros "H".
    cwp_apply (cwp_wp with "H"); iIntros (ev) "H". cwp_lam. cwp_pures.
    iApply cwp_seq_bind'; iApply (cwp_wand with "H"); iIntros (v) "H !>".
    cwp_pures. iApply cwp_bind.
    cwp_apply cwp_atomic_env; iIntros (env) "Henv $". iApply wp_fupd.
    iDestruct "Henv" as (X σ HX) "[Hlocks Hσ]". wp_pures.
    wp_apply (lreplicate_spec 1 with "[//]"); iIntros (ll Hll).
    wp_alloc l as "Hl".
    iMod (full_locking_heap_alloc_upd with "Hσ Hl") as (?) "(Hσ & Hinfo & Hl)"=> //=.
    iDestruct "Hl" as "[Hl _]".
    iIntros "!> !>". iSplitL "Hlocks Hσ".
    { iExists X, _. iFrame.
      iIntros "!%". intros w Hw. destruct (HX _ Hw) as (cl&Hcl&Hin).
      exists cl; split; first done. by rewrite locked_locs_alloc_unlocked. }
    iSpecialize ("H" with "Hl"). rewrite cloc_to_val_eq /=.
    cwp_pures. iApply cwp_bind. cwp_pures. iApply (cwp_wand with "H").
    iIntros (w). iDestruct 1 as (v') "[Hl H]". cwp_pures. iApply cwp_bind.
    cwp_apply cwp_atomic_env; iIntros (env') "Henv $". iApply wp_fupd.
    iDestruct "Henv" as (X' σ' HX') "[Hlock Hσ]". wp_pures.
    iDestruct (full_locking_heap_unlocked with "Hl Hσ") as %?.
    iMod (full_locking_heap_free_upd _ _ [_] with "Hσ Hinfo [$Hl //]")
      as (ll' _ _) "[Hll Hclose] /=".
    wp_store. iIntros "!> !>". iSplitL "Hlock Hll Hclose".
    { iExists X', _. iFrame "Hlock". iSplit; last by iApply "Hclose".
      iPureIntro; intros w' Hw'. destruct (HX' _ Hw') as (cl'&Hcl&Hin).
      exists cl';  split; first done. rewrite locked_locs_delete. set_solver. }
    cwp_pures. iApply cwp_ret. by iApply wp_value.
  Qed.

  Lemma cwp_mut_bind R Φ x e1 e2 :
    CWP e1 @ R {{ v, U (∀ cl,
      cl ↦C v -∗
      CWP subst' x (cloc_to_val cl) e2 @ R {{ w, ∃ v', cl ↦C v' ∗ Φ w }}) }} -∗
    CWP x ←mutᶜ e1 ;ᶜ e2 @ R {{ Φ }}.
  Proof.
    iIntros "H". cwp_pures. iApply cwp_mut_bind'.
    iApply (cwp_wand with "H"); iIntros (v) "H !>"; iIntros (cl) "Hcl".
    cwp_lam. by iApply "H".
  Qed.

  Lemma cwp_if R Φ c e1 e2 :
    CWP c @ R {{ v, (⌜v = #true⌝ ∧ U (CWP e1 @ R {{ Φ }})) ∨
                    (⌜v = #false⌝ ∧ U (CWP e2 @ R {{ Φ }})) }} -∗
    CWP ifᶜ (c) { e1 } elseᶜ { e2 } @ R {{ Φ }}.
  Proof.
    iIntros "H". rewrite /c_if -lock. cwp_pures.
    cwp_apply (cwp_wp with "H"). iIntros (v) "H". cwp_pures.
    iApply cwp_seq_bind'. iApply (cwp_wand with "H").
    iIntros (v') "[[-> ?] | [-> ?]] !>"; by cwp_pures.
  Qed.

  Lemma cwp_while R Φ c e :
    CWP whileVᶜ (c) { e } @ R {{ Φ }} -∗
    CWP whileᶜ (c) { e } @ R {{ Φ }}.
  Proof. iIntros "H". by cwp_pures. Qed.

  Lemma cwp_whileV R Φ c e :
    (* The later is crucial for Löb induction *)
    ▷ CWP c @ R {{ v,
        ⌜v = #true⌝ ∧ U (CWP e @ R {{ _, U (CWP whileVᶜ (c) { e } @ R {{ Φ }})}})
      ∨ ⌜v = #false⌝ ∧ U (Φ #()) }} -∗
    CWP whileVᶜ (c) { e } @ R {{ Φ }}.
  Proof.
    iIntros "H". cwp_lam. cwp_pures. rewrite /c_if -lock. cwp_pures.
    cwp_apply (cwp_wp with "H"). iIntros (v) "H". cwp_lam. cwp_pures.
    iApply cwp_seq_bind'. iApply (cwp_wand with "H").
    iIntros (v') "[[-> H] | [-> H]] !>".
    - cwp_pures. iApply cwp_seq_bind'.
      iApply (cwp_wand with "H"); iIntros (w) "H !>". by cwp_lam.
    - cwp_pures. iApply cwp_ret. by iApply wp_value.
  Qed.

  Lemma cwp_whileV_inv I R Φ c e :
    I -∗
    □ (I -∗ CWP c @ R {{ v, (⌜v = #false⌝ ∧ U (Φ #())) ∨
                            (⌜v = #true⌝ ∧ U (CWP e @ R {{ _, U I }})) }}) -∗
    CWP whileVᶜ (c) { e } @ R {{ Φ }}.
  Proof.
    iIntros "HI #Hinv". iLöb as "IH".
    iApply cwp_whileV. iNext.
    iSpecialize ("Hinv" with "HI"). iApply (cwp_wand with "Hinv").
    iIntros (v) "[[-> H]|[-> H]] /="; first by auto.
    iLeft. iSplit; first done. iModIntro.
    iApply (cwp_wand with "H"); iIntros (_) "HI !>". by iApply "IH".
  Qed.

  Lemma cwp_while_inv I R Φ c e :
    I -∗
    □ (I -∗ CWP c @ R {{ v, (⌜v = #false⌝ ∧ U (Φ #())) ∨
                            (⌜v = #true⌝ ∧ U (CWP e @ R {{ _, U I }})) }}) -∗
    CWP whileᶜ (c) { e } @ R {{ Φ }}.
  Proof.
    iIntros "HI #Hinv". iApply cwp_while. by iApply (cwp_whileV_inv with "HI Hinv").
  Qed.

  Lemma cwp_fun R Φ e mx v :
    CWP subst' mx v e @ R {{ λ w, U (Φ w) }} -∗
    CWP (λᶜ mx, e)%V v @ R {{ Φ }}.
  Proof.
    iIntros "H". cwp_lam. iApply cwp_seq_bind; simpl. cwp_lam.
    iApply (cwp_wand with "H"); iIntros (w) "H !>".
    by iApply cwp_ret; iApply wp_value.
  Qed.

  Lemma cwp_call R Ψ1 Ψ2 Φ ef ea :
    CWP ef @ R {{ Ψ1 }} -∗
    CWP ea @ R {{ Ψ2 }} -∗
    (∀ f a, Ψ1 f -∗ Ψ2 a -∗ U (R -∗ ▷ CWP f a {{ v, R ∗ Φ v }})) -∗
    CWP callᶜ ef ea @ R {{ Φ }}.
  Proof.
    iIntros "H1 H2 H".
    cwp_apply (cwp_wp with "H2"); iIntros (vf) "H2".
    cwp_apply (cwp_wp with "H1"); iIntros (va) "H1".
    cwp_lam. cwp_pures.
    iApply cwp_seq_bind'. iApply (cwp_par with "H1 H2").
    iIntros "!>" (f a) "H1 H2 !>". iSpecialize ("H" with "H1 H2"); iModIntro.
    cwp_pures.
    iApply cwp_atomic. iIntros "HR". iSpecialize ("H" with "HR").
    iExists True%I. iModIntro; iSplit; first done. cwp_pures.
    iApply (cwp_wand with "H"); eauto.
  Qed.

  Lemma cwp_un_op R Φ op e :
    CWP e @ R {{ v, ∃ w, ⌜un_op_eval op v = Some w⌝ ∧ Φ w }} -∗
    CWP c_un_op op e @ R {{ Φ }}.
  Proof.
    iIntros "H".
    cwp_apply (cwp_wp with "H"); iIntros (v) "HΦ"; cwp_lam; cwp_pures.
    iApply cwp_bind.
    iApply (cwp_wand with "HΦ"); iIntros (w) "HΦ"; cwp_lam.
    iDestruct "HΦ" as (w0) "[% H]".
    iApply cwp_ret. by wp_op.
  Qed.

  Lemma cwp_ptr_plus R Φ Ψ2 e1 e2 :
    CWP e2 @ R {{ Ψ2 }} -∗
    CWP e1 @ R {{ v1, ∀ v2, Ψ2 v2 -∗ ∃ cl (n : Z),
                       ⌜ v1 = cloc_to_val cl ⌝ ∗
                       ⌜ v2 = #n ⌝ ∗
                       Φ (cloc_to_val (cl +∗ n)) }} -∗
    CWP c_ptr_plus e1 e2 @ R {{ Φ }}.
  Proof.
    iIntros "He2 HΦ".
    cwp_apply (cwp_wp with "He2"); iIntros (a2) "Ha2".
    cwp_apply (cwp_wp with "HΦ"); iIntros (a1) "Ha1". cwp_lam; cwp_pures.
    iApply cwp_bind. iApply (cwp_par with "Ha1 Ha2").
    iIntros "!>" (v1 v2) "Hv1 Hv2 !>". cwp_pures.
    iDestruct ("Hv1" with "Hv2") as (cl n -> ->) "HΦ /=".
    rewrite cloc_to_val_eq.
    cwp_pures. iApply cwp_ret. by iApply wp_value.
  Qed.

  Lemma cwp_ptr_lt R Φ Ψ1 e1 e2 :
    CWP e1 @ R {{ Ψ1 }} -∗
    CWP e2 @ R {{ v2, ∀ v1, Ψ1 v1 -∗ ∃ p q,
                       ⌜ v1 = cloc_to_val p ⌝ ∗
                       ⌜ v2 = cloc_to_val q ⌝ ∗
                       Φ #(cloc_lt p q) }} -∗
    CWP c_ptr_lt e1 e2 @ R {{ Φ }}.
  Proof.
    iIntros "He1 HΦ".
    cwp_apply (cwp_wp with "HΦ"); iIntros (a2) "Ha2".
    cwp_apply (cwp_wp with "He1"); iIntros (a1) "Ha1". cwp_lam; cwp_pures.
    iApply cwp_bind. iApply (cwp_par with "Ha1 Ha2").
    iIntros "!>" (v1 v2) "Hv1 Hv2 !>". cwp_let.
    iDestruct ("Hv2" with "Hv1") as ([pl pi] [ql qi] -> ->) "HΦ /=".
    rewrite cloc_to_val_eq /cloc_lt /=. cwp_pures.
    case_bool_decide as Hp; subst.
    - rewrite (bool_decide_true (#ql = #ql)) //. cwp_pures.
      iApply cwp_ret. by iApply wp_value.
    - rewrite /= bool_decide_false; last congruence.
      cwp_if. iApply cwp_ret. by iApply wp_value.
  Qed.

  Lemma cwp_bin_op R Φ Ψ1 Ψ2 op e1 e2 :
    CWP e1 @ R {{ Ψ1 }} -∗ CWP e2 @ R {{ Ψ2 }} -∗
    (∀ v1 v2, Ψ1 v1 -∗ Ψ2 v2 -∗ ∃ w, ⌜cbin_op_eval op v1 v2 = Some w⌝ ∧ Φ w) -∗
    CWP c_bin_op op e1 e2 @ R {{ Φ }}.
  Proof.
    iIntros "H1 H2 HΦ".
    destruct op as [op'| |].
    - cwp_apply (cwp_wp with "H2"); iIntros (v2) "HΨ2".
      cwp_apply (cwp_wp with "H1"); iIntros (v1) "HΨ1". cwp_lam; cwp_pures.
      iApply cwp_bind.
      iApply (cwp_par Ψ1 Ψ2 with "HΨ1 HΨ2").
      iNext. iIntros (w1 w2) "HΨ1 HΨ2 !>". cwp_pures.
      iDestruct ("HΦ" with "HΨ1 HΨ2") as (w0 ?) "H". iApply cwp_ret. by wp_op.
    - iApply (cwp_ptr_plus with "H2").
      iApply (cwp_wand with "H1").
      iIntros (v1) "HΨ1"; iIntros (v2) "HΨ2".
      iDestruct ("HΦ" with "HΨ1 HΨ2") as (w Hop) "HΦ"; simpl in *.
      destruct (cloc_of_val v1) as [cl|] eqn:Hcl; simplify_eq/=.
      destruct (int_of_val v2) as [o|] eqn:Ho; simplify_eq/=.
      iExists cl,o. iFrame.
      rewrite -(cloc_to_of_val v1 cl) //.
      by destruct v2; repeat (case_match || simplify_eq/=).
    - iApply (cwp_ptr_lt with "H1").
      iApply (cwp_wand with "H2").
      iIntros (v2) "HΨ2"; iIntros (v1) "HΨ1".
      iDestruct ("HΦ" with "HΨ1 HΨ2") as (w Hop) "HΦ"; simpl in *.
      destruct (cloc_of_val v1) as [cl1|] eqn:Hcl1; simplify_eq/=.
      destruct (cloc_of_val v2) as [cl2|] eqn:Hcl2; simplify_eq/=.
      iExists cl1, cl2. iFrame.
      rewrite -(cloc_to_of_val v1 cl1) // -(cloc_to_of_val v2 cl2) //.
  Qed.

  Lemma cwp_pre_bin_op R Φ Ψ1 Ψ2 e1 e2 op :
    CWP e1 @ R {{ Ψ1 }} -∗ CWP e2 @ R {{ Ψ2 }} -∗
    (∀ v1 v2, Ψ1 v1 -∗ Ψ2 v2 -∗
      ∃ cl v w, ⌜ v1 = cloc_to_val cl ⌝ ∧
                cl ↦C v ∗
                ⌜ cbin_op_eval op v v2 = Some w ⌝ ∗
                (cl ↦C[LLvl] w -∗ Φ v)) -∗
    CWP c_pre_bin_op op e1 e2 @ R {{ Φ }}.
  Proof.
    iIntros "He1 He2 HΦ".
    cwp_apply (cwp_wp with "He2"); iIntros (a2) "Ha2".
    cwp_apply (cwp_wp with "He1"); iIntros (a1) "Ha1". cwp_lam; cwp_pures.
    iApply cwp_bind. iApply (cwp_par with "Ha1 Ha2"). iNext.
    iIntros (v1 v2) "Hv1 Hv2 !>". cwp_pures.
    iApply cwp_atomic.
    iIntros "$ !>". iDestruct ("HΦ" with "Hv1 Hv2") as (cl v w ->) "(Hl & % & HΦ)".
    simplify_eq/=. iExists True%I. iSplit; first done. cwp_pures.
    iApply cwp_bind. iApply cwp_load. iApply cwp_ret. iApply wp_value.
    iExists cl, v; iFrame. iSplit; first done.
    iIntros "Hl". cwp_pures. iApply cwp_bind.
    iApply (cwp_store _ _
      (λ v', ⌜v' = cloc_to_val cl⌝)%I (λ v', ⌜v' = w⌝)%I with "[] [] [-]").
    - cwp_proj. iApply cwp_ret; by wp_value_head.
    - iApply (cwp_bin_op _ _ (λ v', ⌜v' = v⌝)%I (λ v', ⌜v' = v2⌝)%I);
        try (try cwp_proj; iApply cwp_ret; by wp_value_head).
      iIntros (? ? -> ->); eauto.
    - iIntros (? ? -> ->).
      iExists _, _; iFrame. iSplit; first done.
      iIntros "?". cwp_seq. iApply cwp_ret; iApply wp_value.
      iIntros "_". by iApply "HΦ".
  Qed.
End proofs.

(* Make sure that we only use the provided rules and don't break the abstraction *)
Typeclasses Opaque c_alloc c_store c_load c_un_op c_bin_op c_seq_bind c_if c_while c_call.
Global Opaque c_alloc c_store c_load c_un_op c_bin_op c_seq_bind c_if c_while c_call.