clutch.coneris.adequacy

From iris.proofmode Require Import base proofmode.
From iris.bi Require Export lib.fixpoint_mono big_op.
From iris.base_logic.lib Require Import ghost_map invariants fancy_updates.
From iris.algebra Require Import excl.
From iris.prelude Require Import options.

From clutch.prelude Require Import stdpp_ext iris_ext.
From clutch.con_prob_lang Require Import erasure notation.
From clutch.common Require Export con_language sch_erasable.
From clutch.base_logic Require Import error_credits.
From clutch.coneris Require Import weakestpre primitive_laws.
From clutch.prob Require Import distribution.
Import uPred.

Notation con_prob_lang_mdp := (con_lang_mdp con_prob_lang).

This file contains the two adequacy theorems of Coneris, wp_pgl_lim and wp_safety
Normal adequacy
Section adequacy.
  Context `{!conerisGS Σ}.

  Lemma step_fupd_fupdN_S n (P : iProp Σ) : ((|={}▷=>^(S n) P) ⊣⊢ (|={}=> |={}▷=>^(S n) P))%I.
  Proof. iSplit; iIntros; simpl; iApply fupd_idemp; iFrame. Qed.

  Lemma pgl_dbind' `{Countable A, Countable A'}
    (f : A distr A') (μ : distr A) (T : A' Prop) ε' n :
     0 <= ε' -∗
    ( a , |={}=>|={}▷=>^(n) pgl (f a) T ε') -∗
    |={}=> |={}▷=>^(n) pgl (dbind f μ) T (ε') : iProp Σ.
  Proof.
    iIntros (?) "H".
    iApply (step_fupdN_mono _ _ _ (⌜( a, pgl (f a) T ε')⌝)).
    { iIntros (?). iPureIntro.
      rewrite <-Rplus_0_l. eapply pgl_dbind; try done.
      by apply pgl_trivial. }
    iIntros (?) "/=".
    iApply "H".
  Qed.

  Lemma pgl_dbind_adv' `{Countable A, Countable A'}
    (f : A distr A') (μ : distr A) (T : A' Prop) ε' n :
     exists r, forall a, 0 <= ε' a <= r -∗
    ( a , |={}=> |={}▷=>^(n) pgl (f a) T (ε' a)) -∗
    |={}=> |={}▷=>^(n) pgl (dbind f μ) T ( Expval μ ε') : iProp Σ.
  Proof.
    iIntros (?) "H".
    iApply (step_fupdN_mono _ _ _ (⌜( a, pgl (f a) T (ε' a))⌝)).
    { iIntros (?). iPureIntro. rewrite <-Rplus_0_l. eapply pgl_dbind_adv; try done.
      by apply pgl_trivial.
    }
    by iIntros (?) "/=".
  Qed.

  Lemma wp_adequacy_val_fupd `{Countable sch_int_state} (ζ : sch_int_state) e es (sch: scheduler con_prob_lang_mdp sch_int_state) n v σ ε φ `{!TapeOblivious sch_int_state sch}:
    to_val e = Some v
    state_interp σ err_interp ε
    WP e {{ v, φ v }}
    |={, }=> pgl (sch_exec sch n (ζ, (Val v :: es, σ))) φ ε.
  Proof.
    iIntros (Hval) "(?&?&Hwp)".
    rewrite pgl_wp_unfold/pgl_wp_pre.
    iMod ("Hwp" with "[$]") as "H".
    simpl. rewrite Hval.
    iRevertε) "H".
    iApply state_step_coupl_ind.
    iModIntro.
    iIntros (??) "[%|[>(_&_&%)|[H|(%μ&%&%Herasable&%&%Hineq&H)]]]".
    - iPureIntro. by apply pgl_1.
    - iApply fupd_mask_intro; first done.
      iIntros "_".
      iPureIntro.
      erewrite sch_exec_is_final; last done.
      eapply pgl_mon_grading; last apply pgl_dret; auto.
    - iApply (fupd_mono _ _ (_)%I).
      { iPureIntro.
        intros H'.
        apply pgl_epsilon_limit; last exact.
        by apply Rle_ge.
      }
      iIntros (ε' ?).
      unshelve iDestruct ("H" $! (mknonnegreal ε' _) with "[]") as "[H _]"; [|done..].
      pose proof cond_nonneg ε. lra.
    - apply sch_erasable_sch_erasable_val in Herasable.
      erewrite <-Herasable; last done.
      iApply (fupd_mono _ _ (_)%I).
      { iPureIntro.
        intros H'.
        eapply pgl_mon_grading; first apply Hineq.
        rewrite <-Rplus_0_l.
        eapply pgl_dbind_adv; [| |exact H'|by apply pgl_trivial]; naive_solver.
      }
      iIntros (??).
      by iMod ("H" $! _) as "[? _]".
  Qed.

  Lemma state_step_coupl_erasure `{Countable sch_int_state} (ζ : sch_int_state) es σ ε Z n (sch: scheduler con_prob_lang_mdp sch_int_state) φ `{H0 : !TapeOblivious sch_int_state sch}:
    state_step_coupl σ ε Z -∗
    ( σ2 ε2, Z σ2 ε2 ={}=∗ |={}▷=>^(n)
                               pgl (sch_exec sch n (ζ, (es, σ2))) φ ε2) -∗
    |={}=> |={}▷=>^(n)
             pgl (sch_exec sch n (ζ, (es, σ))) φ ε.
  Proof.
    iRevertε).
    iApply state_step_coupl_ind.
    iIntros "!>" (σ ε) "[%|[?|[H|(%μ&%&%Herasable&%&%Hineq&H)]]] Hcont".
    - iApply step_fupdN_intro; first done.
      iPureIntro.
      by apply pgl_1.
    - by iMod ("Hcont" with "[$]").
    - iApply (step_fupdN_mono _ _ _ (⌜( ε', _)⌝)).
      { iPureIntro.
        intros.
        apply pgl_epsilon_limit; [simpl; by apply Rle_ge|done].
      }
      iIntros (ε' ?).
      unshelve iDestruct ("H" $! (mknonnegreal ε' _) with "[]") as "[H _]";
        [pose proof cond_nonneg ε; simpl in *; lra|done|simpl].
      iApply ("H" with "[$]").
    - apply sch_erasable_sch_erasable_val in Herasable.
      rewrite -Herasable.
      epose proof (step_fupdN_mono _ _ _ (pgl _ _ (_))%I) as h.
      iApply h.
      { iPureIntro.
        intros. eapply pgl_mon_grading; first (simpl; apply Hineq). exact.
      }
      iApply (pgl_dbind_adv' with "[][-]"); first (iPureIntro; naive_solver).
      iIntros (?).
      iMod ("H"$!_) as "[H _]".
      simpl.
      iApply ("H" with "[$]").
  Qed.

  Lemma state_step_coupl_erasure' `{Countable sch_int_state} (ζ : sch_int_state) e es e' σ ε Z n num (sch: scheduler con_prob_lang_mdp sch_int_state) φ `{!TapeOblivious sch_int_state sch}:
    ((e::es)!!num = Some e') ->
    to_val e = None ->
    to_val e' = None ->
    state_step_coupl σ ε Z -∗
    ( σ2 ε2, Z σ2 ε2 ={}=∗ |={}▷=>^(S n)
                               pgl (prim_step e' σ2 ≫= λ '(e3, σ3, efs), sch_exec sch n (ζ, (<[num:=e3]> (e :: es) ++ efs, σ3))) φ ε2) -∗
    |={}=> |={}▷=>^(S n)
             pgl (prim_step e' σ ≫= λ '(e3, σ3, efs), sch_exec sch n (ζ, (<[num:=e3]> (e :: es) ++ efs, σ3))) φ ε.
  Proof.
    intros ???.
    iRevertε).
    iApply state_step_coupl_ind.
    iIntros "!>" (σ ε) "[%|[?|[H|(%μ&%&%Herasable&%&%Hineq&H)]]] Hcont".
    - iApply step_fupdN_intro; first done.
      iPureIntro.
      by apply pgl_1.
    - by iMod ("Hcont" with "[$]").
    - iApply (step_fupdN_mono _ _ _ (⌜( ε', _)⌝)).
      { iPureIntro.
        intros.
        apply pgl_epsilon_limit; [simpl; by apply Rle_ge|done].
      }
      iIntros (ε' ?).
      unshelve iDestruct ("H" $! (mknonnegreal ε' _) with "[]") as "[H _]";
        [pose proof cond_nonneg ε; simpl in *; lra|done|simpl].
      iApply ("H" with "[$]").
    - unshelve erewrite (Rcoupl_eq_elim _ _ (prim_coupl_step_prim_sch_erasable _ _ _ _ _ _ _ μ _ _ _ _)); [done..|].
      epose proof (step_fupdN_mono _ _ _ (pgl _ _ _)%I) as h.
      iApply h.
      { iPureIntro.
        intros. eapply pgl_mon_grading; first (simpl; apply Hineq). exact.
      }
      iApply (pgl_dbind_adv'); first (iPureIntro; naive_solver).
      iIntros (?).
      iMod ("H" $! _) as "[H _]".
      simpl.
      iApply ("H" with "[$]").
  Qed.

  Lemma wp_refRcoupl_step_fupdN `{Countable sch_int_state} (ζ : sch_int_state) (ε : nonnegreal)
    (e : expr) es (σ : state) n φ (sch: scheduler con_prob_lang_mdp sch_int_state) `{!TapeOblivious sch_int_state sch}:
    state_interp σ err_interp (ε) WP e {{ v, φ v }} ([∗ list] e' es, WP e' {{ _, True%I }})
    |={,}=> |={}▷=>^n pgl (sch_exec sch n (ζ, (e::es, σ))) φ ε.
  Proof.
    iInduction n as [|n] "IH" forall (ζ ε e es σ); iIntros "((Hσh & Hσt) & Hε & Hwp & Hwps)".
    - Local Transparent sch_exec.
      destruct (to_val e) eqn:Heq.
      + apply of_to_val in Heq as <-.
        iApply wp_adequacy_val_fupd; last iFrame.
        done.
      + rewrite /sch_exec /=.
        iApply fupd_mask_intro; [set_solver|]; iIntros "_".
        iPureIntro. rewrite Heq.
        apply pgl_dzero,
        Rle_ge,
        cond_nonneg.
    - rewrite sch_exec_Sn /sch_step_or_final/=.
      destruct (to_val e) eqn:Heq.
      + apply of_to_val in Heq as <-.
        rewrite dret_id_left'/=.
        iMod (wp_adequacy_val_fupd with "[$]") as "%"; first done.
        repeat iModIntro.
        by iApply step_fupdN_intro.
      + rewrite {1}/sch_step. rewrite <-dbind_assoc.
        (* replace (ε) with (0+ε)*)
        iApply fupd_mask_intro; first done.
        iIntros "Hclose".
        rewrite -fupd_idemp.
        iApply (pgl_dbind' _ _ _ _ (S _)); first done.
        iIntros ([sch_σ sch_a]).
        iMod "Hclose" as "_".
        simpl. rewrite Heq.
        destruct ((e::es)!!sch_a) as [chosen_e|] eqn:Hlookup; rewrite Hlookup; last first.
        { (* step a thread thats out of bounds *)
          rewrite dmap_dret dret_id_left'.
          iApply fupd_mask_intro; first done.
          iIntros "Hclose".
          do 2 iModIntro. iMod "Hclose" as "_".
          iApply "IH". iFrame.
        }
        case_match eqn:Hcheckval.
        { (* step a thread thats a value *)
          rewrite dmap_dret dret_id_left'.
          iApply fupd_mask_intro; first done.
          iIntros "Hclose".
          do 2 iModIntro. iMod "Hclose" as "_".
          iApply "IH". iFrame.
        }
        rewrite dmap_comp/dmap-dbind_assoc.
        erewrite (distr_ext _ _); last first.
        { intros x. erewrite (dbind_ext_right _ _ (λ '(_,_,_), _)); last first.
          - intros [[??]?].
            by rewrite dret_id_left/=.
          - done.
        }
        destruct sch_a as [|sch_a].
        * (* step main thread *)
          rewrite pgl_wp_unfold /pgl_wp_pre. iSimpl in "Hwp". rewrite Heq.
          iMod ("Hwp" with "[$]") as "Hlift".
          (* replace (0+ε)*)
          iApply (state_step_coupl_erasure' with "[$Hlift]"); [done..|].
          iIntros (σ2 ε2) "(%&%&%&%Hineq&H)".
          epose proof (step_fupdN_mono _ _ _ (pgl _ _ _)%I) as h.
          iApply h.
          { iPureIntro.
            intros. eapply pgl_mon_grading; first (simpl; apply Hineq). exact.
          }
          replace (chosen_e) with e; last by (simpl in Hlookup; simplify_eq).
          iApply (pgl_dbind_adv' with "[][-]"); first (iPureIntro; naive_solver).
          iIntros ([[??]?]).
          iMod ("H" $! _ _ _ ) as "H".
          simpl.
          do 3 iModIntro.
          iApply (state_step_coupl_erasure with "[$][-]").
          iIntros (??) ">(?&?&?&?)".
          iApply ("IH" with "[-]"). iFrame.
        * (* step other threads*)
          simpl in Hlookup.
          apply list_elem_of_split_length in Hlookup as (l1 & l2 & -> & ->).
          iDestruct "Hwps" as "[Hl1 [Hwp' Hl2]]".
          rewrite (pgl_wp_unfold _ _ chosen_e)/pgl_wp_pre.
          iSimpl in "Hwp'".
          rewrite Hcheckval.
          iMod ("Hwp'" with "[$]") as "Hlift".
          (* replace (0+ε)*)
          iApply (state_step_coupl_erasure' _ e _ chosen_e with "[$]"); [|done..|].
          { simpl.
            rewrite lookup_app_r//.
            by replace (_-_)%nat with 0%nat by lia.
          }
          iIntros (σ2 ε2) "(%&%&%&%Hineq&H)".
          epose proof (step_fupdN_mono _ _ _ (pgl _ _ _)%I) as h.
          iApply h.
          { iPureIntro.
            intros. eapply pgl_mon_grading; first (simpl; apply Hineq). exact.
          }
          iApply (pgl_dbind_adv' with "[][-]"); first (iPureIntro; naive_solver).
          iIntros ([[??]?]).
          iMod ("H" $! _ _ _) as "H".
          simpl.
          do 3 iModIntro.
          iApply (state_step_coupl_erasure with "[$][-]").
          iIntros (??) ">(?&?&?&?)".
          iApply ("IH" with "[-]"). iFrame.
          rewrite insert_app_r_alt//.
          replace (_-_)%nat with 0%nat by lia.
          simpl.
          iFrame.
  Qed.

End adequacy.

Class conerisGpreS Σ := ConerisGpreS {
  conerisGpreS_iris :: invGpreS Σ;
  conerisGpreS_heap :: ghost_mapG Σ loc val;
  conerisGpreS_tapes :: ghost_mapG Σ loc tape;
  conerisGpreS_err :: ecGpreS Σ;
}.

Definition conerisΣ : gFunctors :=
  #[invΣ; ghost_mapΣ loc val;
    ghost_mapΣ loc tape;
    GFunctor (authR nonnegrealUR)].
Global Instance subG_conerisGPreS {Σ} : subG conerisΣ Σ conerisGpreS Σ.
Proof. solve_inG. Qed.

Theorem wp_pgl_multi Σ `{conerisGpreS Σ} `{Countable sch_int_state} (ζ : sch_int_state) n
  (e : expr) es (σ : state) (ε : R) (sch: scheduler con_prob_lang_mdp sch_int_state) φ `{!TapeOblivious sch_int_state sch}:
  0 <= ε
  ( `{conerisGS Σ}, ε -∗ (WP e {{ v, φ v }} [∗ list] e' es, WP e' {{ _, True%I }}))
  pgl (sch_exec sch n (ζ, (e::es, σ))) φ ε.
Proof.
  intros Hwp.
  eapply pure_soundness, (step_fupdN_soundness_no_lc _ n 0).
  iIntros (Hinv) "_".
  iMod (ghost_map_alloc σ.(heap)) as "[%γH [Hh _]]".
  iMod (ghost_map_alloc σ.(tapes)) as "[%γT [Ht _]]".
  (* Handle the trivial 1 <= ε case *)
  destruct (decide (ε < 1)) as [Hcr|Hcr]; last first.
  { iClear "Hh Ht".
    iApply (fupd_mask_intro); [eauto|].
    iIntros "_".
    iApply step_fupdN_intro; [eauto|].
    iApply laterN_intro; iPureIntro.
    apply not_Rlt, Rge_le in Hcr.
    rewrite /pgl; intros.
    eapply Rle_trans; [eapply prob_le_1|done]. }
  set ε' := mknonnegreal _ .
  iMod (ec_alloc ε') as (?) "[??]"; [done|].
  set (HclutchGS := HeapG Σ _ _ _ γH γT _).
  epose proof (wp_refRcoupl_step_fupdN _ ε') as h.
  iApply h.
  iFrame. by iApply Hwp.
  Unshelve. apply _.
Qed.

Theorem wp_pgl Σ `{conerisGpreS Σ} `{Countable sch_int_state} (ζ : sch_int_state) n
  (e : expr) (σ : state) (ε : R) (sch: scheduler con_prob_lang_mdp sch_int_state) φ `{!TapeOblivious sch_int_state sch}:
  0 <= ε
  ( `{conerisGS Σ}, ε -∗ WP e {{ v, φ v }})
  pgl (sch_exec sch n (ζ, ([e], σ))) φ ε.
Proof.
  intros ? Hwp.
  eapply wp_pgl_multi; [done..|].
  simpl.
  iIntros.
  iSplitL; last done.
  by iApply Hwp.
Qed.

Lemma pgl_closed_lim `{Countable sch_int_state} (ζ : sch_int_state) (e : expr) (σ : state) (ε : R)
  (sch: scheduler con_prob_lang_mdp sch_int_state) φ `{!TapeOblivious sch_int_state sch}:
  ( n, pgl (sch_exec sch n (ζ, ([e], σ))) φ ε)
  pgl (sch_lim_exec sch (ζ, ([e], σ))) φ ε .
Proof. intros Hn. by apply sch_lim_exec_continuous_prob. Qed.

Theorem wp_pgl_lim Σ `{conerisGpreS Σ} `{Countable sch_int_state} (ζ : sch_int_state)
  (e : expr) (σ : state) (ε : R) (sch: scheduler con_prob_lang_mdp sch_int_state) φ `{!TapeOblivious sch_int_state sch}:
  0 <= ε
  ( `{conerisGS Σ}, ε -∗ WP e {{ v, φ v }})
  pgl (sch_lim_exec sch (ζ, ([e], σ))) φ ε.
Proof.
  intros.
  apply pgl_closed_lim; first done.
  intros.
  by eapply wp_pgl.
Qed.

safety
Definition sch_mass_1 `{Countable sch_int_state} (sch: scheduler con_prob_lang_mdp sch_int_state) :=
   sch_cfg, SeriesC (sch sch_cfg)=1.

Section safety.
  Context `{!conerisGS Σ}.

  Lemma safety_dbind_adv `{Countable A, Countable A'}(h : A distr A')
    (μ : distr A) ε (ε2 : A R) R :
    SeriesC μ = 1 ->
    (exists r, forall a, 0 <= ε2(a) <= r)
    pgl μ R ε ->
    ( a, R a -> SeriesC (h a) >= 1 - ε2 a)
    SeriesC (μ ≫= h) >= 1 - (ε + Expval μ ε2).
  Proof.
    intros H' H1 Hpgl H2.
    rewrite dbind_mass/Expval.
    rewrite /pgl in Hpgl.
    setoid_rewrite <-SeriesC_scal_l.
    rewrite -H'.
    trans (SeriesC μ - ( prob μ (λ a : A, Datatypes.negb (bool_decide (R a))) + SeriesC (λ a : A, μ a * ε2 a))); last lra.
    replace (SeriesC μ) with (prob μ (λ a, bool_decide (R a))+prob μ (λ a, Datatypes.negb (bool_decide (R a)))); last first.
    { rewrite /prob. rewrite -SeriesC_plus.
      - apply SeriesC_ext. intros. case_bool_decide; simpl; lra.
      - apply ex_seriesC_filter_bool_pos; auto.
      - apply ex_seriesC_filter_bool_pos; auto.
    }
    rewrite Rminus_plus_r_l.
    rewrite /prob.
    rewrite Rminus_def.
    replace (-_) with (-1 * SeriesC (λ a, μ a * ε2 a)) by lra.
    rewrite -SeriesC_scal_l -SeriesC_plus; [|by apply ex_seriesC_filter_bool_pos|by apply ex_seriesC_scal_l, pmf_ex_seriesC_mult_fn].
    apply Rle_ge.
    rewrite Rcomplements.Rminus_le_0.
    rewrite Rminus_def.
    replace (-_) with (-1 * SeriesC (λ x : A, (if bool_decide (R x) then μ x else 0) + -1 * (μ x * ε2 x))) by lra.
    rewrite -SeriesC_scal_l -SeriesC_plus; last first.
    { apply ex_seriesC_scal_l, ex_seriesC_plus; first by apply ex_seriesC_filter_bool_pos.
      apply ex_seriesC_scal_l. by apply pmf_ex_seriesC_mult_fn.
    }
    { setoid_rewrite SeriesC_scal_l. apply pmf_ex_seriesC_mult_fn. naive_solver. }
    apply SeriesC_ge_0'.
    intros a.
    rewrite SeriesC_scal_l.
    case_bool_decide.
    - trans (μ a * (1 - ε2 a) - μ a + μ a * ε2 a); first lra.
      assert (μ a * (1 - ε2 a)<= μ a * SeriesC (h a)); last lra.
      apply Rmult_le_compat_l; first done.
      apply Rge_le.
      naive_solver.
    - apply Rplus_le_le_0_compat; first real_solver.
      assert (0<=μ a) by done.
      assert (0<=ε2 a) by naive_solver.
      rewrite Rmult_plus_distr_l -Rmult_assoc.
      replace (-1*-1) with 1 by lra.
      rewrite Rmult_0_r Rplus_0_l Rmult_1_l.
      real_solver.
  Qed.

  Lemma safety_dbind_adv' `{Countable A, Countable A'}(h : A distr A')
    (μ : distr A) ε (ε2 : A R) n R:
    SeriesC μ = 1 -∗
    exists r, forall a, 0 <= ε2(a) <= r -∗
               pgl μ R ε -∗
               ( a, R a ={}=∗ |={}▷=>^(n)SeriesC (h a) >= 1 - ε2 a) -∗
               |={}=> |={}▷=>^(n) SeriesC (μ ≫= h) >= 1 - (ε + Expval μ ε2) : iProp Σ.
  Proof.
    iIntros (???) "H".
    iApply (step_fupdN_mono _ _ _ (⌜( a, R a -> SeriesC (h a) >= 1 - ε2 a)⌝)).
    { iPureIntro. intros. by eapply safety_dbind_adv. }
    iIntros (??).
    by iApply "H".
  Qed.

  Lemma safety_dbind' `{Countable A, Countable A'}(h : A distr A')
    (μ : distr A) (ε : R) n:
    0 <= ε -∗
    SeriesC μ = 1 -∗
    ( a, |={}▷=>^(n)SeriesC (h a) >= 1 - ε) -∗
    |={}=> |={}▷=>^(n) SeriesC (μ ≫= h) >= 1 - ε : iProp Σ.
  Proof.
    iIntros (? Hmass) "H".
    pose (ε' (a:A):= ε).
    assert (ε=0+Expval μ ε') as K.
    { rewrite /ε'. rewrite Expval_const; last done. rewrite Hmass. simpl. lra. }
    rewrite {2}K.
    epose proof (safety_dbind_adv' h _ 0 ε' n) as hh.
    iPoseProof (hh with "[][][][H]") as "K".
    - done.
    - iPureIntro. exists ε. intros. naive_solver.
    - iPureIntro. by apply pgl_trivial.
    - iIntros. rewrite /ε'. iApply "H".
    - rewrite /ε'.
      rewrite Expval_const; last done.
      by rewrite Hmass Rmult_1_r.
  Qed.

  Lemma sch_erasable_mass `{Countable sch_int_state} (ζ : sch_int_state) (sch: scheduler con_prob_lang_mdp sch_int_state) `{!TapeOblivious sch_int_state sch} μ σ:
    sch_erasable (λ (t : Type) _ _
                   (sch : scheduler con_prob_lang_mdp t), TapeOblivious t sch) μ σ ->
    SeriesC μ = 1.
  Proof.
    rewrite /sch_erasable.
    intros H0.
    cut (SeriesC (μ≫=(λ σ', sch_exec sch 0 (ζ, ([Val (#())], σ')))) = SeriesC (sch_exec sch 0 (ζ, ([Val (#())], σ)))).
    - rewrite dbind_mass. simpl. rewrite dret_mass. rewrite SeriesC_scal_r.
      lra.
    - apply sch_erasable_sch_erasable_val in H0.
      by rewrite H0.
  Qed.


  Lemma state_step_coupl_erasure_safety `{Countable sch_int_state} (ζ : sch_int_state) es σ ε Z n (sch: scheduler con_prob_lang_mdp sch_int_state) `{H0 : !TapeOblivious sch_int_state sch}:
    state_step_coupl σ ε Z -∗
    ( σ2 ε2, Z σ2 ε2 ={}=∗ |={}▷=>^(n)
                               SeriesC (sch_pexec sch n (ζ, (es, σ2))) >= 1 - ε2) -∗
    |={}=> |={}▷=>^(n)
             SeriesC (sch_pexec sch n (ζ, (es, σ))) >= 1- nonneg ε.
  Proof.
    iRevertε).
    iApply state_step_coupl_ind.
    iIntros "!>" (σ ε) "[%|[?|[H|(%μ&%&%Herasable&%&%&H)]]] Hcont".
    - iApply step_fupdN_intro; first done.
      iPureIntro.
      trans 0; try lra.
      by apply Rle_ge.
    - by iMod ("Hcont" with "[$]").
    - iApply (step_fupdN_mono _ _ _ (⌜( ε', ε<ε'-> SeriesC (sch_pexec sch n (ζ, (es, σ))) >= 1 - ε')⌝)).
      { iPureIntro.
        intros H1.
        apply Rle_ge.
        apply Rle_plus_epsilon.
        intros eps Heps.
        assert (SeriesC (sch_pexec sch n (ζ, (es, σ))) >= 1 - (ε+eps)) as H2.
        - apply H1. lra.
        - apply Rge_le in H2. rewrite Rminus_plus_distr in H2.
          by rewrite Rcomplements.Rle_minus_l in H2.
      }
      iIntros (ε' ?).
      unshelve iDestruct ("H" $! (mknonnegreal ε' _) with "[]") as "[H _]";
        [pose proof cond_nonneg ε; simpl in *; lra|done|simpl].
      iApply ("H" with "[$]").
    - replace (SeriesC _) with (SeriesC (μ ≫= λ σ', sch_pexec sch n (ζ, (es, σ')))); last first.
      { rewrite /sch_erasable in Herasable.
        epose proof Herasable _ _ _ sch es ζ n _ as H4.
        assert (SeriesC (dmap
         (λ x, x.2.1)
         (μ ≫= λ σ', sch_pexec sch n (ζ, (es, σ')))) =
       SeriesC (dmap
         (λ x, x.2.1)
         (sch_pexec sch n (ζ, (es, σ))))) as H5.
        - f_equal. by rewrite H4.
        - by rewrite !dmap_mass in H5.
      }
      iApply (step_fupdN_mono _ _ _ (SeriesC (μ ≫= λ σ', sch_pexec sch n (ζ, (es, σ'))) >= 1 - (0 + Expval μ ε2))).
      { iPureIntro. intros. etrans; first exact.
        apply Rle_ge.
        apply Ropp_le_cancel.
        rewrite !Ropp_minus_distr. rewrite Rplus_0_l. by apply Rplus_le_compat_r.
      }
      iApply (safety_dbind_adv' _ _ _ (_)); [| |iPureIntro; by apply pgl_trivial|].
      { iPureIntro. eapply sch_erasable_mass; last first; done. }
      { iPureIntro. naive_solver. }
      iIntros (a HR).
      iMod ("H" $! _) as "[H _]".
      by iMod ("H" with "[$]").
  Qed.

  Lemma state_step_coupl_erasure_safety' `{Countable sch_int_state} (ζ : sch_int_state) e es e' σ ε Z n num (sch: scheduler con_prob_lang_mdp sch_int_state) `{!TapeOblivious sch_int_state sch}:
    ((e::es)!!num = Some e') ->
    to_val e = None ->
    to_val e' = None ->
    state_step_coupl σ ε Z -∗
    ( σ2 ε2, Z σ2 ε2 ={}=∗ |={}▷=>^(S n)
                               SeriesC (prim_step e' σ2 ≫= λ '(e3, σ3, efs), sch_pexec sch n (ζ, (<[num:=e3]> (e :: es) ++ efs, σ3))) >=1- ε2) -∗
    |={}=> |={}▷=>^(S n)
             SeriesC (prim_step e' σ ≫= λ '(e3, σ3, efs), sch_pexec sch n (ζ, (<[num:=e3]> (e :: es) ++ efs, σ3))) >= 1- nonneg ε.
  Proof.
    intros ???.
    iRevertε).
    iApply state_step_coupl_ind.
    iIntros "!>" (σ ε) "[%|[?|[H|(%μ&%&%Herasable&%&%&H)]]] Hcont".
    - iApply step_fupdN_intro; first done.
      iPureIntro.
      trans 0; try lra.
      by apply Rle_ge.
    - by iMod ("Hcont" with "[$]").
    - iApply (step_fupdN_mono _ _ _
                (⌜( ε', ε<ε'-> SeriesC (prim_step e' σ ≫= λ '(e3, σ3, efs), sch_pexec sch n (ζ, (<[num:=e3]> (e :: es) ++ efs, σ3))) >= 1 - ε')⌝)).
      { iPureIntro.
        intros K1.
        apply Rle_ge.
        apply Rle_plus_epsilon.
        intros eps Heps.
        assert (SeriesC (prim_step e' σ ≫= λ '(e3, σ3, efs), sch_pexec sch n (ζ, (<[num:=e3]> (e :: es) ++ efs, σ3))) >= 1 - (ε+eps)) as K2.
        - apply K1. lra.
        - apply Rge_le in K2. rewrite Rminus_plus_distr in K2.
          by rewrite Rcomplements.Rle_minus_l in K2.
      }
      iIntros (ε' ?).
      unshelve iDestruct ("H" $! (mknonnegreal ε' _) with "[]") as "[H _]";
        [pose proof cond_nonneg ε; simpl in *; lra|done|simpl].
      iApply ("H" with "[$]").
    - replace (SeriesC _) with (SeriesC (μ ≫= (λ σ', prim_step e' σ' ≫= λ '(e3, σ3, efs), sch_pexec sch n (ζ, (<[num:=e3]> (e :: es) ++ efs, σ3))))); last first.
      { (* rewrite /sch_erasable in Herasable. *)
        (* epose proof Herasable _ _ _ sch es ζ n _ as H4. *)
        assert (SeriesC (dmap
                           (λ x, x.2.1)
                           (μ
     ≫= λ σ', prim_step e' σ'
          ≫= λ '(e3, σ3, efs), sch_pexec sch n (ζ, (<[num:=e3]> (e :: es) ++ efs, σ3)))) =
                SeriesC (dmap
                           (λ x, x.2.1)
                           (prim_step e' σ ≫= λ '(e3, σ3, efs), sch_pexec sch n (ζ, (<[num:=e3]> (e :: es) ++ efs, σ3))))) as K.
        - unshelve epose proof prim_coupl_step_prim_pexec_sch_erasable e n es σ ζ e' num μ _ _ _ _ as K'; try done.
          apply Rcoupl_eq_elim in K'. by rewrite K'.
        - by rewrite !dmap_mass in K.
      }
      iApply (step_fupdN_mono _ _ _ (SeriesC (μ ≫= (λ σ', prim_step e' σ' ≫= λ '(e3, σ3, efs), sch_pexec sch n (ζ, (<[num:=e3]> (e :: es) ++ efs, σ3)))) >= 1 - (0 + Expval μ ε2))).
      { iPureIntro. intros. etrans; first exact.
        apply Rle_ge.
        apply Ropp_le_cancel.
        rewrite !Ropp_minus_distr. rewrite Rplus_0_l. by apply Rplus_le_compat_r.
      }
      iApply (safety_dbind_adv' _ _ _ (_)); [| |iPureIntro; by apply pgl_trivial|].
      { iPureIntro. eapply sch_erasable_mass; last first; done. }
      { iPureIntro. naive_solver. }
      iIntros (a HR).
      iMod ("H" $! _) as "[H _]".
      by iMod ("H" with "[$]").
  Qed.

  Lemma wp_safety_step_fupdN `{Countable sch_int_state} (ζ : sch_int_state) (ε : nonnegreal)
    (e : expr) es (σ : state) n φ (sch: scheduler con_prob_lang_mdp sch_int_state) `{!TapeOblivious sch_int_state sch}:
    sch_mass_1 sch->
    state_interp σ err_interp (ε) WP e {{ v, φ v }} ([∗ list] e' es, WP e' {{ _, True%I }})
    |={,}=> |={}▷=>^n SeriesC (sch_pexec sch n (ζ, (e::es, σ))) >= 1 - ε.
  Proof.
    intros Hmass.
    iInduction n as [|n] "IH" forall (ζ ε e es σ); iIntros "((Hσh & Hσt) & Hε & Hwp & Hwps)".
    - iApply (fupd_mask_intro); [eauto|].
      iIntros "_".
      iApply step_fupdN_intro; [eauto|].
      iPureIntro.
      rewrite sch_pexec_O dret_mass.
      pose proof cond_nonneg ε. simpl. lra.
    - rewrite sch_pexec_Sn /sch_step_or_final/=.
      case_match eqn:Hval.
      { erewrite SeriesC_ext; last (intros; by rewrite dret_id_left').
        iApply fupd_mask_intro; first done.
        iIntros "Hclose".
        do 2 iModIntro.
        iMod "Hclose". iApply "IH".
        iFrame.
      }
      rewrite /sch_step. rewrite <-dbind_assoc.
      iApply fupd_mask_intro; first done.
      iIntros "Hclose".
      rewrite -fupd_idemp.
      iApply (safety_dbind' _ _ _ (S _)); [done..|].
      iIntros ([ζ' thid]).
      simpl. rewrite Hval.
      destruct ((e::es)!!thid) as [chosen_e|] eqn:Hlookup; rewrite Hlookup; last first.
      { erewrite SeriesC_ext; last first.
        - intros.
          by rewrite dmap_dret dret_id_left'.
        - do 2 iModIntro. iMod "Hclose".
          iApply "IH". iFrame. }
      case_match eqn:Hcheckval.
      { erewrite SeriesC_ext; last first.
        - intros.
          by rewrite dmap_dret dret_id_left'.
        - do 2 iModIntro. iMod "Hclose".
          iApply "IH". iFrame. }
      rewrite dmap_comp/dmap-dbind_assoc.
      erewrite (distr_ext _ _); last first.
      { intros x. erewrite (dbind_ext_right _ _ (λ '(_,_,_), _)); last first.
        - intros [[??]?].
          by rewrite dret_id_left/=.
        - done. }
      destruct thid as [|thid].
      + iMod "Hclose". rewrite pgl_wp_unfold/pgl_wp_pre.
        iMod ("Hwp" with "[$]") as "Hwp".
        rewrite -fupd_idemp.
        iApply (state_step_coupl_erasure_safety' with "[$Hwp]"); try done.
        iIntros (σ2 ε2). simpl. rewrite Hval.
        rewrite /prog_coupl.
        iIntros "( %εb & %Hred &%Hbound & %Hineq & Hlift)".
        assert (e = chosen_e) as ?; last simplify_eq.
        { simpl in Hlookup. by simplify_eq. }
        iApply (step_fupdN_mono _ _ _ (SeriesC
           (prim_step chosen_e σ2 ≫= λ '(e3, σ3, efs), sch_pexec sch n (ζ', (e3 :: es ++ efs, σ3))) >=
                                         1 - (0 + Expval (prim_step chosen_e σ2) εb))).
        { iPureIntro. simpl in *.
          intros. etrans; first exact.
          apply Rle_ge.
          apply Ropp_le_cancel.
          rewrite !Ropp_minus_distr. rewrite Rplus_0_l. by apply Rplus_le_compat_r.
        }
        iApply (safety_dbind_adv' _ _ _ _ (S n) with "[][][]").
        * iPureIntro. eapply mixin_prim_step_mass; [apply con_language_mixin|done].
        * iPureIntro. naive_solver.
        * iPureIntro; by apply pgl_trivial.
        * iIntros (((e' & σ') & efs) _).
          iMod ("Hlift"$! _ _ _).
          simpl.
          do 3 iModIntro.
          iApply (state_step_coupl_erasure_safety with "[$]").
          iIntros (??) ">(?&?&?&?)". iApply "IH". iFrame.
      + simpl in Hlookup.
        apply list_elem_of_split_length in Hlookup as (l1 & l2 & -> & ->).
        iDestruct "Hwps" as "[Hl1 [Hwp' Hl2]]".
        rewrite (pgl_wp_unfold _ _ chosen_e)/pgl_wp_pre.
        iSimpl in "Hwp'".
        rewrite Hcheckval.
        iMod "Hclose".
        iMod ("Hwp'" with "[$]") as "Hlift".
        rewrite -fupd_idemp.
        iApply (state_step_coupl_erasure_safety' with "[$]"); try done.
        { simpl. by apply list_lookup_middle. }
        iIntros (σ2 ε2). simpl. rewrite /prog_coupl.
        iIntros "( %εb & %Hred &%Hbound & %Hineq & Hlift)".
        iApply (step_fupdN_mono _ _ _ (SeriesC
           (prim_step chosen_e σ2 ≫= λ '(e3, σ3, efs),
                 sch_pexec sch n (ζ', (e :: <[length l1:=e3]> (l1 ++ chosen_e :: l2) ++ efs, σ3))) >=
                                         1 - (0 + Expval (prim_step chosen_e σ2) εb))).
        { iPureIntro. simpl in *.
          intros. etrans; first exact.
          apply Rle_ge.
          apply Ropp_le_cancel.
          rewrite !Ropp_minus_distr.
          rewrite Rplus_0_l. by apply Rplus_le_compat_r. }
        iApply (safety_dbind_adv' _ _ _ _ (S n) with "[][][]").
        * iPureIntro. eapply mixin_prim_step_mass; [apply con_language_mixin|done].
        * iPureIntro. naive_solver.
        * iPureIntro. by apply pgl_trivial.
        * iIntros (((e' & σ') & efs) HR).
          iMod ("Hlift" $! _ _ _).
          simpl.
          do 3 iModIntro.
          iApply (state_step_coupl_erasure_safety with "[$]").
          iIntros (??) ">(?&?&?&?)". iApply "IH". iFrame.
          rewrite insert_app_r_alt//.
          replace (_-_)%nat with 0%nat by lia.
          simpl. iFrame.
  Qed.

End safety.

Theorem wp_safety_multi Σ `{conerisGpreS Σ} `{Countable sch_int_state} (ζ : sch_int_state) n
  (e : expr) es (σ : state) (ε : R) (sch: scheduler con_prob_lang_mdp sch_int_state) φ `{!TapeOblivious sch_int_state sch}:
  0 <= ε
  sch_mass_1 sch ->
  ( `{conerisGS Σ}, ε -∗ (WP e {{ v, φ v }} [∗ list] e' es, WP e' {{ _, True%I }}))
  SeriesC (sch_pexec sch n (ζ, (e::es, σ))) >= 1 - ε.
Proof.
  intros Hmass Hwp.
  eapply pure_soundness, (step_fupdN_soundness_no_lc _ n 0).
  iIntros (Hinv) "_".
  iMod (ghost_map_alloc σ.(heap)) as "[%γH [Hh _]]".
  iMod (ghost_map_alloc σ.(tapes)) as "[%γT [Ht _]]".
  (* Handle the trivial 1 <= ε case *)
  destruct (decide (ε < 1)) as [Hcr|Hcr]; last first.
  { iClear "Hh Ht".
    iApply (fupd_mask_intro); [eauto|].
    iIntros "_".
    iApply step_fupdN_intro; [eauto|].
    iApply laterN_intro; iPureIntro.
    trans 0; [by apply Rle_ge, SeriesC_ge_0|lra].
  }
  set ε' := mknonnegreal _ .
  iMod (ec_alloc ε') as (?) "[??]"; [done|].
  set (HclutchGS := HeapG Σ _ _ _ γH γT _).
  epose proof (wp_safety_step_fupdN _ ε') as h.
  iApply h; first done.
  iFrame. by iApply Hwp.
  Unshelve. apply _.
Qed.

Theorem wp_safety Σ `{conerisGpreS Σ} `{Countable sch_int_state} (ζ : sch_int_state)
  (e : expr) (σ : state) (ε : R) n (sch: scheduler con_prob_lang_mdp sch_int_state) φ `{!TapeOblivious sch_int_state sch}:
  0 <= ε
  sch_mass_1 sch ->
  ( `{conerisGS Σ}, ε -∗ WP e {{ v, φ v }})
  SeriesC (sch_pexec sch n (ζ, ([e], σ))) >= 1 - ε.
Proof.
  iIntros ( Hmass Hwp).
  eapply wp_safety_multi; [done..|].
  iIntros.
  iSplitL; last done.
  by iApply Hwp.
Qed.