From iris.program_logic Require Import ownership wsat.
From iris.program_logic Require Export pviewshifts invariants.
From iris.proofmode Require Import pviewshifts invariants.

Local Hint Extern 10 (_ ≤ _) ⇒ omega.
Local Hint Extern 100 (_ ⊥ _) ⇒ set_solver.
Local Hint Extern 100 (_ ∉ _) ⇒ set_solver.
Local Hint Extern 10 (✓{_} _) ⇒
  repeat match goal with
  | H : wsat _ _ _ _ _ |- _ ⇒ apply wsat_valid in H; last omega
  end; solve_validN.


Program Definition psvs_def {Λ Σ} (E1 E2 : coPset) (P : iProp Λ Σ) : iProp Λ Σ :=
  {| uPred_holds n r1 rl := ∀ k Ef σ rf rfl,
       0 < k ≤ n → E1 ∪ E2 ⊥ Ef →
       wsat k (E1 ∪ Ef) σ (r1 ⋅ rf) (rl ⋅ rfl) →
       ∃ r2 rl2, P k r2 rl2 ∧ wsat k (E2 ∪ Ef) σ (r2 ⋅ rf) (rl2 ⋅ rfl) ∧ rl ⤳_(k) rl2 |}.
Next Obligation.
  intros Λ Σ E1 E2 P n r1 rl1 r2 rl2 HP [r3 Hr2] Hrl k Ef σ rf rfl ?? Hwsat; simpl in ×.
  edestruct (HP k Ef σ (r3 ⋅ rf) rfl) as (r'&rl'&?&?&?); eauto.
  - rewrite ?(assoc op) -(dist_le _ _ _ _ Hr2); last lia.
    by rewrite (dist_le _ _ _ _ Hrl); last lia.
  - ∃ (r' ⋅ r3), rl'; rewrite -assoc -(dist_le _ _ _ _ Hrl) //; last lia.
    split_and!; auto. apply uPred_mono with r' rl'; eauto using cmra_includedN_l.
Qed.
Next Obligation. naive_solver. Qed.

Definition psvs_aux : { x | x = @psvs_def }. by eexists. Qed.
Definition psvs := proj1_sig psvs_aux.
Definition psvs_eq : @psvs = @psvs_def := proj2_sig psvs_aux.

Arguments psvs {_ _} _ _ _%I.
Instance: Params (@psvs) 4.

Notation "|={ E1 , E2 }=>> Q" := (psvs E1 E2 Q%I)
  (at level 99, E1, E2 at level 50, Q at level 200,
   format "|={ E1 , E2 }=>> Q") : uPred_scope.
Notation "|={ E }=>> Q" := (psvs E E Q%I)
  (at level 99, E at level 50, Q at level 200,
   format "|={ E }=>> Q") : uPred_scope.
Notation "|==>> Q" := (psvs ⊤ ⊤ Q%I)
  (at level 99, Q at level 200, format "|==>> Q") : uPred_scope.

Section psvs.
Context {Λ : language} {Σ : iFunctor}.
Implicit Types P Q : iProp Λ Σ.
Implicit Types m : iGst Λ Σ.

Lemma psvs_zero E1 E2 P r rl: psvs_def E1 E2 P 0 r rl.
Proof. intros ??????. exfalso. omega. Qed.

Global Instance psvs_ne E1 E2 n : Proper (dist n ==> dist n) (@psvs Λ Σ E1 E2).
Proof.
  rewrite psvs_eq.
  intros P Q HPQ; split⇒ n' r1 rl1 ???; simpl; split; intros HP rf rfl k Ef σ ???;
    destruct (HP rf rfl k Ef σ) as (r2&rl2&?&?&?); auto;
    ∃ r2, rl2; split_and?; auto; try apply HPQ; eauto.
Qed.
Global Instance psvs_proper E1 E2 : Proper ((≡) ==> (≡)) (@psvs Λ Σ E1 E2).
Proof. apply ne_proper, _. Qed.

Lemma psvs_stepP E m ml (P: iGst Λ Σ → iGst Λ Σ → Prop):
  m # ml ~~>>: P → (ownG m ★ ownGl ml)%I ⊢ (|={E}=>> ∃ m' ml', ⧆■ P m' ml' ★ ownG m' ★ ownGl ml') .
Proof.
  rewrite psvs_eq. intros Hup.
  uPred.unseal; split;
  intros [|n] r rl ? ? (r1&r2&rl1&rl2&Hr&Hrl&(HownG&HaffG)%ownG_spec&HownGl%ownGl_spec)
         [|k] Ef σ rf rfl ???; try lia.
  destruct (wsat_update_gst' k (E ∪ Ef) σ r rf rl rfl m ml P) as (m'&ml'&?&?&?); eauto.
  {
    rewrite (dist_le _ _ _ _ Hr); last lia.
    etransitivity.
    - apply cmra_includedN_le with (S n); first apply HownG; auto.
    - apply cmra_includedN_l; auto.
  }
  {
    rewrite (dist_le _ _ _ _ Hrl); last lia.
    rewrite (dist_le _ _ _ _ HaffG) -?(dist_le _ _ _ _ HownGl) //= ?left_id //=; lia.
  }
  ∃ (update_gst m' r), (update_gst ml' rl).
  split_and!; auto.
  ∃ m'. ∃ ml'.
    ∃ (∅: iRes Λ Σ), (update_gst m' r), (∅: iRes Λ Σ) , (update_gst ml' rl);
    split_and!; rewrite ?left_id; eauto.
    - by split.
    - ∃ (update_gst m' r), (∅: iRes Λ Σ) , (∅: iRes Λ Σ), (update_gst ml' rl); split_and!;
      rewrite ?left_id ?right_id; auto.
      × apply ownG_spec; split; auto.
      × apply ownGl_spec.
        rewrite /update_gst.
        rewrite (dist_le _ _ _ _ Hrl); last lia.
        rewrite ?(dist_le _ _ _ _ HaffG); last lia.
        rewrite -(dist_le _ _ _ _ HownGl); last lia.
        simpl. rewrite ?left_id; auto.
Qed.

Lemma psvs_mono E1 E2 P Q : (P ⊢ Q) → (|={E1,E2}=>> P) ⊢ |={E1,E2}=>> Q.
Proof.
  rewrite psvs_eq. intros HPQ; split⇒ n r rl ? ? HP k Ef σ rf rfl ???.
  destruct (HP k Ef σ rf rfl) as (r2&rl2&?&?&?); eauto.
  ∃ r2, rl2; split_and!; eauto using uPred_in_entails.
Qed.
Lemma pvs_psvs E1 E2 E3 P :
  E2 ⊆ E1 ∪ E3 → (|={E1,E2}=> |={E2,E3}=>> P) ⊢ (|={E1,E3}=>> P).
Proof.
  rewrite psvs_eq pvs_eq. intros ?; split⇒ n r1 rl ? ? HP1 k Ef σ rf rfl ???.
  destruct (HP1 k Ef σ rf rfl) as (r2&HP2&?); auto.
Qed.
Lemma psvs_pvs E1 E2 E3 P :
  E2 ⊆ E1 ∪ E3 → (|={E1,E2}=>> |={E2,E3}=> P) ⊢ (|={E1,E3}=>> P).
Proof.
  rewrite psvs_eq pvs_eq. intros ?; split⇒ n r1 rl ? ? HP1 k Ef σ rf rfl ???.
  destruct (HP1 k Ef σ rf rfl) as (r2&rl2&HP2&?&?); auto.
  destruct (HP2 k Ef σ rf rfl) as (r3&HP3&?); eauto.
Qed.
Lemma psvs_mask_frame E1 E2 Ef P :
  Ef ⊥ E1 ∪ E2 → (|={E1,E2}=>> P) ⊢ (|={E1 ∪ Ef,E2 ∪ Ef}=>> P).
Proof.
  rewrite psvs_eq. intros ?; split⇒ n r rl ? ? HP k Ef' σ rf rfl ???.
  destruct (HP k (Ef∪Ef') σ rf rfl) as (r'&rl'&?&?); rewrite ?(assoc_L _); eauto.
  by ∃ r', rl'; rewrite -(assoc_L _).
Qed.
Lemma psvs_frame_r E1 E2 P Q : ((|={E1,E2}=>> P) ★ ⧆Q) ⊢ (|={E1,E2}=>> P ★ ⧆Q).
Proof.
  rewrite psvs_eq.
  uPred.unseal; split; intros n r rl ? ? (r1&r2&rl1&rl2&Hr&Hrl&HP&HQ) k Ef σ rf rfl ???.
  destruct (HP k Ef σ (r2 ⋅ rf) (rl2 ⋅ rfl)) as (r'&rl'&?&?&?); eauto.
  {
    rewrite ?assoc -(dist_le _ _ _ _ Hr); last lia.
    rewrite -(dist_le _ _ _ _ Hrl); last lia; auto.
  }
  ∃ (r' ⋅ r2), (rl' ⋅ rl2); split_and!.
  - ∃ r', r2, rl', rl2. split_and?; eauto using dist_le.
    apply uPred_closed with n; auto.
  - by rewrite -?assoc.
  - rewrite (dist_le _ _ _ _ Hrl); last lia.
    destruct HQ as (_&Haff). rewrite (dist_le _ _ _ _ Haff) ?right_id; auto.
Qed.

Import uPred.
Global Instance psvs_mono' E1 E2 : Proper ((⊢) ==> (⊢)) (@psvs Λ Σ E1 E2).
Proof. intros P Q; apply psvs_mono. Qed.
Global Instance psvs_flip_mono' E1 E2 :
  Proper (flip (⊢) ==> flip (⊢)) (@psvs Λ Σ E1 E2).
Proof. intros P Q; apply psvs_mono. Qed.
Lemma pvs_psvs' E P : (|={E}=> |={E}=>> P) ⊢ (|={E}=>> P).
Proof. apply pvs_psvs; set_solver. Qed.
Lemma psvs_pvs' E P : (|={E}=>> |={E}=> P) ⊢ (|={E}=>> P).
Proof. apply psvs_pvs; set_solver. Qed.
Lemma psvs_frame_l E1 E2 P Q : (⧆P ★ |={E1,E2}=>> Q) ⊢ (|={E1,E2}=>> ⧆P ★ Q).
Proof. rewrite !(comm _ (⧆P)%I); apply psvs_frame_r. Qed.
Lemma psvs_const_l E1 E2 P Q : (⧆P ★ |={E1,E2}=>> Q) ⊢ (|={E1,E2}=>> ⧆P ★ Q).
Proof. rewrite !(comm _ (⧆P)%I); apply psvs_frame_r. Qed.
Lemma psvs_wand_l E1 E2 P Q : (⧆(P -★ Q) ★ (|={E1,E2}=>> P)) ⊢ (|={E1,E2}=>> Q).
Proof. by rewrite psvs_frame_l affine_elim wand_elim_l. Qed.
Lemma psvs_wand_r E1 E2 P Q : ((|={E1,E2}=>> P) ★ ⧆(P -★ Q)) ⊢ (|={E1,E2}=>> Q).
Proof. by rewrite psvs_frame_r affine_elim wand_elim_r. Qed.
Lemma psvs_mask_frame' E1 E1' E2 E2' P :
  E1' ⊆ E1 → E2' ⊆ E2 → E1 ∖ E1' = E2 ∖ E2' →
  (|={E1',E2'}=>> P) ⊢ (|={E1,E2}=>> P).
Proof.
  intros HE1 HE2 HEE.
  rewrite (psvs_mask_frame _ _ (E1 ∖ E1')); last set_solver.
  by rewrite {2}HEE -!union_difference_L.
Qed.

Lemma psvs_mask_frame_mono E1 E1' E2 E2' P Q :
  E1' ⊆ E1 → E2' ⊆ E2 → E1 ∖ E1' = E2 ∖ E2' →
  (P ⊢ Q) → (|={E1',E2'}=>> P) ⊢ (|={E1,E2}=>> Q).
Proof. intros HE1 HE2 HEE →. by apply psvs_mask_frame'. Qed.

Lemma psvs_mask_weaken E1 E2 P : E1 ⊆ E2 → (|={E1}=>> P) ⊢ (|={E2}=>> P).
Proof. auto using psvs_mask_frame'. Qed.

Lemma psvs_step E m ml m' ml' :
  (cmra_step_update m ml m' ml') → (ownG m ★ ownGl ml)%I ⊢ (|={E}=>> ownG m' ★ ownGl ml').
Proof.
  intros; rewrite (psvs_stepP E _ _ (λ x xl, m' = x ∧ ml' = xl)); last apply cmra_step_stepP; auto.
  apply psvs_mono.
  apply uPred.exist_elim⇒ m''.
  apply uPred.exist_elim ⇒ ml''.
  apply uPred.pure_elim_sep_l⇒Heq.
  by destruct Heq as (->&->).
Qed.

End psvs.

Definition psvs_fsa {Λ Σ} : FSA Λ Σ () := λ E Φ, (|={E}=>> Φ ())%I.
Global Arguments psvs_fsa _ _ _ _/.
 Instance psvs_fsa_prf {Λ Σ}: AffineFrameShiftAssertion True (@psvs_fsa Λ Σ).
Proof.
  rewrite /psvs_fsa.
  split; auto using psvs_mask_frame_mono.
  - intros. rewrite pvs_psvs; last set_solver. rewrite psvs_pvs; last set_solver. auto.
  - intros. rewrite pvs_psvs; last set_solver. rewrite psvs_pvs; last set_solver. auto.
  - intros. rewrite -{1}(uPred.affine_affine P) psvs_frame_r.
    by setoid_rewrite uPred.affine_elim.
Qed.