From stdpp Require Import fin_maps.
From iris.bi Require Import lib.fractional.
From iris.proofmode Require Import proofmode.
From clutch.prob_lang Require Import tactics lang notation.
From clutch.prob_lang.gwp Require Import gen_weakestpre.
From clutch.eris Require Export total_primitive_laws array_laws.
From iris.prelude Require Import options.

Section lifting.

  Context `{!erisGS Σ}.
  Implicit Types P Q : iProp Σ.
  Implicit Types Φ Ψ : val → iProp Σ.
  Implicit Types σ : state.
  Implicit Types v : val.
  Implicit Types l : loc.
  Implicit Types vs : list val.
  Implicit Types sz off : nat.

Lemma twp_allocN E v n s :
  (0 < n)%Z →
  [[{ True }]] AllocN (Val $ LitV $ LitInt $ n) (Val v) @ s; E
  [[{ l, RET LitV (LitLoc l);
          l ↦∗ replicate (Z.to_nat n) v }]].
  Proof.
    iIntros (? Φ) "_ HΦ".
    iApply twp_allocN_seq; auto; try lia.
    iIntros (l) "Hlm".
    iApply "HΦ".
    by iApply pointsto_seq_array.
  Qed.

  Lemma twp_allocN_vec E v n s :
    (0 < n)%Z →
    [[{ True }]]
AllocN #n v @ s; E
                          [[{ l, RET #l; l ↦∗ vreplicate (Z.to_nat n) v }]].
  Proof.
    iIntros (? Φ) "_ HΦ".
    iApply (twp_allocN with "[//] [HΦ]"); try lia.
    iIntros (l) "Hl".
    iApply "HΦ". by rewrite vec_to_list_replicate.
  Qed.

Lemma twp_load_offset E l dq off vs v s :
  vs !! off = Some v →
  [[{ l ↦∗{dq} vs }]] ! #(l +ₗ off) @ s; E [[{ RET v; l ↦∗{dq} vs }]].
Proof.
  iIntros (Hlookup Φ) "Hl HΦ".
  iDestruct (update_array l _ _ _ _ Hlookup with "Hl") as "[Hl1 Hl2]".
  iApply (twp_load with "Hl1").
  iIntros "Hl1".
  iApply "HΦ".
  iDestruct ("Hl2" $! v) as "Hl2". rewrite list_insert_id; last done.
  iApply "Hl2".
  iApply "Hl1".
Qed.

Lemma twp_load_offset_vec E l dq sz (off : fin sz) (vs : vec val sz) s :
  [[{ l ↦∗{dq} vs }]] ! #(l +ₗ off) @ s; E [[{ RET vs !!! off; l ↦∗{dq} vs }]].
  Proof. apply twp_load_offset. by apply vlookup_lookup. Qed.

Lemma twp_store_offset E l off vs v s :
  is_Some (vs !! off) →
  [[{ l ↦∗ vs }]] #(l +ₗ off) <- v @ s; E [[{ RET #(); l ↦∗ <[off:=v]> vs }]].
Proof.
  iIntros ([w Hlookup] Φ) "Hl HΦ".
  iDestruct (update_array l _ _ _ _ Hlookup with "Hl") as "[Hl1 Hl2]".
  iApply (twp_store with "Hl1").
  iIntros "Hl1".
  iApply "HΦ". iApply "Hl2". iApply "Hl1".
Qed.

Lemma twp_store_offset_vec E l sz (off : fin sz) (vs : vec val sz) v s :
  [[{ l ↦∗ vs }]] #(l +ₗ off) <- v @ s; E [[{ RET #(); l ↦∗ vinsert off v vs }]].
Proof.
  setoid_rewrite vec_to_list_insert. apply twp_store_offset.
  eexists. by apply vlookup_lookup.
Qed.

End lifting.

Global Typeclasses Opaque array.

Section eris_twp_genwp.

  Context `{!erisGS Σ}.

  Local Definition eris_twp_gwp_mixin :
    GenWpMixin
      false
      (λ E e Φ, WP e @ twp_default; E [{ Φ }])%I
      (λ l dq v, (l ↦{dq} v)%I).
  Proof.
    constructor.
    - intros l dq v. apply _.
    - iIntros. by iApply tgl_wp_value.
    - intros. apply tgl_wp_fupd.
    - intros. apply (tgl_wp_bind (ectxi_language.fill _)).
    - intros. by iApply total_lifting.twp_pure_step_fupd.
    - iIntros. by iApply twp_alloc.
    - iIntros. by iApply (twp_load with "[$]").
    - iIntros. by iApply (twp_store with "[$]").
  Qed.

  Canonical Structure eris_twp_genwp : GenWp Σ :=
    {| gwp_laters := false;
       gwp :=
         λ E e Φ, WP e @ twp_default; E [{ Φ }];
       gwp_pointsto :=
         λ l dq v, (l ↦{dq} v)%I;
       gwp_mixin := eris_twp_gwp_mixin |}%I.

End eris_twp_genwp.