(* iris_heap_lang/lib/atomic_heap.v *)
From iris.bi.lib Require Import fractional
From iris.proofmode Require Import proofmode
From iris.program_logic Require Export atomic
From iris.heap_lang Require Export derived_laws
From iris.heap_lang Require Import notation proofmode
From iris.prelude Require Import options

Class atomic_heap:= AtomicHeap {

Global Argumentsalloc : simpl never.

Global Argumentsfree : simpl never.

Global Argumentsload : simpl never.

Global Argumentsstore : simpl never.

Global Argumentscmpxchg : simpl never.

Global Argumentspointsto : simpl never.
Existing Class atomic_heapGS.

Global HintMode atomic_heapGS + + : typeclass_instances.

Global HintExtern 0 (atomic_heapGS _) => progress simpl : typeclass_instances.

Local NotationCAS e1 e2 e3 := (Snd (cmpxchg e1 e2 e3)).

Definition faa_atomic`{!atomic_heap} : val :=
  rec: "faa" "l" "n" :=
    let: "m" := load "l" in
    if: CAS "l" "m" ("m" + "n") then "m" else "faa" "l" "n".

Notation CASe1 e2 e3 := (Snd (cmpxchg e1 e2 e3)).

Notation FAAe1 e2 := (faa_atomic e1 e2).

Lemma cas_spec(l : loc) (w1 w2 : val) :
    val_is_unboxed w1 →
    ⊢ <<{ ∀∀ v, pointsto l (DfracOwn 1) v }>>
      CAS #l w1 w2 @ ∅
    <<{ if decide (v = w1)
        then pointsto l (DfracOwn 1) w2 else pointsto l (DfracOwn 1) v
      | RET #if decide (v = w1) then true else false }>>.
  Proof.
    iIntros (? Φ) "AU". awp_apply cmpxchg_spec; first done.
    iApply (aacc_aupd_commit with "AU"); first done.
    iIntros (v) "H↦". iAaccIntro with "H↦"; first by eauto with iFrame.
    iIntros "$ !> HΦ !>". wp_pures. done.
  Qed.

Lemma faa_spec(l : loc) (i2 : Z) :
    ⊢ <<{ ∀∀ i1 : Z, pointsto l (DfracOwn 1) #i1 }>>
      FAA #l #i2 @ ∅
    <<{ pointsto l (DfracOwn 1) #(i1 + i2) | RET #i1 }>>.
  Proof.
    iIntros (Φ) "AU". rewrite /faa_atomic. iLöb as "IH".
    wp_pures. awp_apply load_spec.
    iApply (aacc_aupd_abort with "AU"); first done.
    iIntros (i1) "H↦". iAaccIntro with "H↦"; first by eauto with iFrame.
    iIntros "$ !> AU !>". wp_pures.
    awp_apply cas_spec; first done.
    iApply (aacc_aupd with "AU"); first done.
    iIntros (m) "Hl".
    iAaccIntro with "Hl"; first by eauto with iFrame.
    iIntros "Hl"; destruct (decide (#m = #i1)); simplify_eq.
    - iModIntro. iRight. iFrame. iIntros "Hpost". iModIntro. by wp_pures.
    - iModIntro. iLeft. iFrame. iIntros "AU". iModIntro. wp_pure.
      by iApply "IH".
  Qed.

Definition primitive_alloc: val :=
  λ: "v", ref "v".

Definition primitive_free: val :=
  λ: "v", Free "v".

Definition primitive_load: val :=
  λ: "l", !"l".

Definition primitive_store: val :=
  λ: "l" "x", "l" <- "x".

Definition primitive_cmpxchg: val :=
  λ: "l" "e1" "e2", CmpXchg "l" "e1" "e2".

Lemma primitive_alloc_spec(v : val) :
    {{{ True }}} primitive_alloc v {{{ l, RET #l; l ↦ v }}}.
  Proof.
    iIntros (Φ) "_ HΦ". wp_lam. wp_alloc l. iApply "HΦ". done.
  Qed.

Lemma primitive_free_spec(l : loc) (v : val) :
    {{{ l ↦ v }}} primitive_free #l {{{ l, RET #l; True }}}.
  Proof.
    iIntros (Φ) "Hl HΦ". wp_lam. wp_free. iApply "HΦ". done.
  Qed.

Lemma primitive_load_spec(l : loc) :
    ⊢ <<{ ∀∀ (v : val) q, l ↦{q} v }>> primitive_load #l @ ∅
      <<{ l ↦{q} v | RET v }>>.
  Proof.
    iIntros (Φ) "AU". wp_lam.
    iMod "AU" as (v q) "[H↦ [_ Hclose]]".
    wp_load. iMod ("Hclose" with "H↦") as "HΦ". done.
  Qed.

Lemma primitive_store_spec(l : loc) (w : val) :
    ⊢ <<{ ∀∀ v, l ↦ v }>> primitive_store #l w @ ∅
      <<{ l ↦ w | RET #() }>>.
  Proof.
    iIntros (Φ) "AU". wp_lam. wp_let.
    iMod "AU" as (v) "[H↦ [_ Hclose]]".
    wp_store. iMod ("Hclose" with "H↦") as "HΦ". done.
  Qed.

Lemma primitive_cmpxchg_spec(l : loc) (w1 w2 : val) :
    val_is_unboxed w1 →
    ⊢ <<{ ∀∀ (v : val), l ↦ v }>>
      primitive_cmpxchg #l w1 w2 @ ∅
    <<{ if decide (v = w1) then l ↦ w2 else l ↦ v
      | RET (v, #if decide (v = w1) then true else false) }>>.
  Proof.
    iIntros (? Φ) "AU". wp_lam. wp_pures.
    iMod "AU" as (v) "[H↦ [_ Hclose]]".
    destruct (decide (v = w1)) as [Heq|Hne];
      [wp_cmpxchg_suc|wp_cmpxchg_fail];
    iMod ("Hclose" with "H↦") as "HΦ"; done.
  Qed.

Definition primitive_atomic_heap: atomic_heap :=
  {| atomic_heapGS _ := TCTrue;
     alloc_spec _ _ _ _ := primitive_alloc_spec;
     free_spec _ _ _ _ := primitive_free_spec;
     load_spec _ _ _ _ := primitive_load_spec;
     store_spec _ _ _ _ := primitive_store_spec;
     cmpxchg_spec _ _ _ _ := primitive_cmpxchg_spec;
     pointsto_persist _ _ _ _ := primitive_laws.pointsto_persist;
     pointsto_agree _ _ _ _ := primitive_laws.pointsto_agree |}.