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.tachis Require Export primitive_laws.
From iris.prelude Require Import options.

Definition array `{!tachisGS Σ F} (l : loc) (dq : dfrac) (vs : list val) : iProp Σ :=
  [∗ list] i ↦ v ∈ vs, (l +ₗ i) ↦{dq} v.

(*
Notation "l ↦∗{ dq } vs" := (array l dq vs)
  (at level 20, dq custom dfrac at level 1, format "l  ↦∗{ dq } vs") : bi_scope.
*)

Notation "l ↦∗{ dq } vs" := (array l dq vs)
  (at level 20, format "l ↦∗{ dq } vs") : bi_scope.

Notation "l ↦∗ v" := (l ↦∗{ DfracOwn 1 } v)%I
                      (at level 20, format "l ↦∗ v") : bi_scope.

Section lifting.

Context `{!tachisGS Σ F}.
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.

Global Instance array_timeless l q vs : Timeless (array l q vs) := _.

Global Instance array_fractional l vs : Fractional (λ q, l ↦∗{DfracOwn q} vs)%I := _.
Global Instance array_as_fractional l q vs :
  AsFractional (l ↦∗{DfracOwn q} vs) (λ q, l ↦∗{DfracOwn q} vs)%I q.
Proof. split; done || apply _. Qed.

Lemma array_nil l dq : l ↦∗{dq} [] ⊣⊢ emp.
Proof. by rewrite /array. Qed.

Lemma array_singleton l dq v : l ↦∗{dq} [v] ⊣⊢ l ↦{dq} v.
Proof. by rewrite /array /= right_id loc_add_0. Qed.

Lemma array_app l dq vs ws :
  l ↦∗{dq} (vs ++ ws) ⊣⊢ l ↦∗{dq} vs ∗ (l +ₗ length vs) ↦∗{dq} ws.
Proof.
  rewrite /array big_sepL_app.
  setoid_rewrite Nat2Z.inj_add.
  by setoid_rewrite loc_add_assoc.
Qed.

Lemma array_cons l dq v vs : l ↦∗{dq} (v :: vs) ⊣⊢ l ↦{dq} v ∗ (l +ₗ 1) ↦∗{dq} vs.
Proof.
  rewrite /array big_sepL_cons loc_add_0.
  setoid_rewrite loc_add_assoc.
  setoid_rewrite Nat2Z.inj_succ.
  by setoid_rewrite Z.add_1_l.
Qed.

Global Instance array_cons_frame l dq v vs R Q :
  Frame false R (l ↦{dq} v ∗ (l +ₗ 1) ↦∗{dq} vs) Q →
  Frame false R (l ↦∗{dq} (v :: vs)) Q | 2.
Proof. by rewrite /Frame array_cons. Qed.

Lemma update_array l dq vs off v :
  vs !! off = Some v →
  ⊢ l ↦∗{dq} vs -∗ ((l +ₗ off) ↦{dq} v ∗ ∀ v', (l +ₗ off) ↦{dq} v' -∗ l ↦∗{dq} <[off:=v']>vs).
Proof.
  iIntros (Hlookup) "Hl".
  rewrite -[X in (l ↦∗{_} X)%I](take_drop_middle _ off v); last done.
  iDestruct (array_app with "Hl") as "[Hl1 Hl]".
  iDestruct (array_cons with "Hl") as "[Hl2 Hl3]".
  assert (off < length vs) as H by (apply lookup_lt_is_Some; by eexists).
  rewrite length_take min_l; last by lia. iFrame "Hl2".
  iIntros (w) "Hl2".
  clear Hlookup. assert (<[off:=w]> vs !! off = Some w) as Hlookup.
  { apply list_lookup_insert_eq. lia. }
  rewrite -[in (l ↦∗{_} <[off:=w]> vs)%I](take_drop_middle (<[off:=w]> vs) off w Hlookup).
  iApply array_app. rewrite take_insert_ge; last by lia. iFrame.
  iApply array_cons.
  rewrite length_take min_l; last by lia. iFrame.
  rewrite drop_insert_lt; last by lia. done.
Qed.

Lemma pointsto_seq_array l dq v n :
  ([∗ list] i ∈ seq 0 n, (l +ₗ (i : nat)) ↦{dq} v) -∗
  l ↦∗{dq} replicate n v.
Proof.
  rewrite /array. iInduction n as [|n'] "IH" forall (l); simpl.
  { done. }
  iIntros "[$ Hl]". rewrite -fmap_S_seq big_sepL_fmap.
  setoid_rewrite Nat2Z.inj_succ. setoid_rewrite <-Z.add_1_l.
  setoid_rewrite <-loc_add_assoc. iApply "IH". done.
Qed.

Lemma wp_allocN E x v (z : Z) s :
  TCEq x (cost (AllocN #z v)) →
  (0 < z)%Z →
  {{{ if bool_decide (cost (AllocN #z v) = 0%R) then True else ⧖ x }}}
    AllocN #z v @ s; E
  {{{ l, RET #l; l ↦∗ replicate (Z.to_nat z) v }}}.
Proof.
  iIntros (? ? Φ) "Hx HΦ".
  iApply (wp_allocN_seq with "[$Hx]"); [lia|].
  iModIntro.
  iIntros (l) "Hlm".
  iApply "HΦ".
  by iApply pointsto_seq_array.
Qed.

Lemma wp_allocN_vec E x v (z : Z) s :
  TCEq x (cost (AllocN #z v)) →
  (0 < z)%Z →
  {{{ if bool_decide (cost (AllocN #z v) = 0%R) then True else ⧖ x }}}
    AllocN #z v @ s; E
  {{{ l, RET #l; l ↦∗ vreplicate (Z.to_nat z) v }}}.
Proof.
  iIntros (-> ? Φ) "Hx HΦ".
  iApply (wp_allocN with "[$] [HΦ]"); try lia.
  iModIntro.
  iIntros (l) "Hl".
  iApply "HΦ". by rewrite vec_to_list_replicate.
Qed.

Lemma wp_load_offset E x l dq off vs v s :
  TCEq x (cost (! #(l +ₗ off))%E) →
  vs !! off = Some v →
  {{{ (if bool_decide (cost (! #(l +ₗ off))%E = 0%R) then True else ⧖ x) ∗ ▷ l ↦∗{dq} vs }}}
    ! #(l +ₗ off) @ s; E
  {{{ RET v; l ↦∗{dq} vs }}}.
Proof.
  iIntros (-> Hlookup Φ) "[Hx >Hl] HΦ".
  iDestruct (update_array l _ _ _ _ Hlookup with "Hl") as "[Hl1 Hl2]".
  iApply (wp_load with "[$]").
  iModIntro.
  iIntros "Hl1".
  iApply "HΦ".
  iDestruct ("Hl2" $! v) as "Hl2". rewrite list_insert_id //.
  iApply "Hl2".
  iApply "Hl1".
Qed.

Lemma wp_load_offset_vec E x l dq sz (off : fin sz) (vs : vec val sz) s :
  TCEq x (cost (! #(l +ₗ off))%E) →
  {{{ (if bool_decide (cost (! #(l +ₗ off)) = 0%R) then True else ⧖ x) ∗ ▷ l ↦∗{dq} vs }}}
    ! #(l +ₗ off) @ s; E
  {{{ RET vs !!! off; l ↦∗{dq} vs }}}.
Proof. intros ->; apply wp_load_offset => //. by apply vlookup_lookup. Qed.

Lemma wp_store_offset E x l off vs v s :
  TCEq x (cost (#(l +ₗ off) <- v)%E) →
  is_Some (vs !! off) →
  {{{ (if bool_decide (cost (#(l +ₗ off) <- v) = 0%R) then True else ⧖ x) ∗ ▷ l ↦∗ vs }}}
    #(l +ₗ off) <- v @ s; E
  {{{ RET #(); l ↦∗ <[off:=v]> vs }}}.
Proof.
  iIntros (-> [w Hlookup] Φ) "[Hx >Hl] HΦ".
  iDestruct (update_array l _ _ _ _ Hlookup with "Hl") as "[Hl1 Hl2]".
  iApply (wp_store with "[$]").
  iModIntro.
  iIntros "Hl1".
  iApply "HΦ". iApply "Hl2". iApply "Hl1".
Qed.

Lemma wp_store_offset_vec E x l sz (off : fin sz) (vs : vec val sz) v s :
  TCEq x (cost (#(l +ₗ off) <- v)%E) →
  {{{ (if bool_decide (cost (#(l +ₗ off) <- v) = 0%R) then True else ⧖ x) ∗ ▷ l ↦∗ vs }}}
    #(l +ₗ off) <- v @ s; E
  {{{ RET #(); l ↦∗ vinsert off v vs }}}.
Proof.
  intros ->.
  setoid_rewrite vec_to_list_insert. apply wp_store_offset => //.
  eexists. by apply vlookup_lookup.
Qed.

End lifting.

Global Typeclasses Opaque array.