From 5398cfde0d0b5c515ea10e058b7aa9e72257cd6c Mon Sep 17 00:00:00 2001
From: Markus de Medeiros <markusdemedeiros@outlook.com>
Date: Thu, 15 Feb 2024 16:12:43 -0500
Subject: [PATCH] update sampler examples to use the twp

---
 .../approx_higherorder_incremental.v          |  98 +++++
 .../approx_higherorder_rejection_sampler.v    | 376 +++++++-----------
 .../approximate_samplers/approx_random_walk.v |  24 +-
 .../approx_rejection_sampler_presampled.v     |  19 +-
 .../approximate_samplers/approx_sampler_lib.v | 151 ++-----
 .../approximate_samplers/approx_walkSAT.v     | 100 +++--
 6 files changed, 341 insertions(+), 427 deletions(-)
 create mode 100644 theories/examples/approximate_samplers/approx_higherorder_incremental.v

diff --git a/theories/examples/approximate_samplers/approx_higherorder_incremental.v b/theories/examples/approximate_samplers/approx_higherorder_incremental.v
new file mode 100644
index 000000000..d7b1a7699
--- /dev/null
+++ b/theories/examples/approximate_samplers/approx_higherorder_incremental.v
@@ -0,0 +1,98 @@
+(** * Higher order specification for incremental sampling algorithms *)
+From clutch.ub_logic Require Export ub_clutch ub_rules.
+From clutch Require Export examples.approximate_samplers.approx_sampler_lib.
+From Coquelicot Require Import Series.
+Require Import Lra.
+
+Section incremental_spec.
+  Local Open Scope R.
+  Context `{!ub_clutchGS Σ}.
+
+  (* Ψ : state
+     ξ : error
+     L : progress bound
+   *)
+
+  Definition incr_sampling_scheme_spec (sampler checker : val) (Ψ : nat -> iProp Σ) (ξ : nat -> nonnegreal) L iL E Θ : iProp Σ :=
+    ( (* Either 0 credit or 0 progress => we will sample a value which the checker accepts
+         Allowed to consume (or invalidate Ψ) in this process *)
+      ((€ (ξ 0%nat) ∨ Ψ 0%nat) -∗ WP sampler #() @ E [{ fun s => WP checker (Val s) @ E [{fun v => ⌜v = #true⌝ ∗ Θ s }] }]) ∗
+      (* Given any amount of credit and progress, we can get a sample such that... *)
+     □ (∀ i p, (⌜((S p) <= L)%nat ⌝ ∗ ⌜((S i) < iL)%nat ⌝ ∗ € (ξ (S i)) ∗ Ψ (S p)) -∗
+            WP sampler #() @ E [{ fun s =>
+                   (*...we're done by chance alone, or... *)
+                  (WP checker (Val s) @ E [{fun v => ⌜v = #true⌝ ∗ (Θ s)}]) ∨
+                   (*... we make prgress, and can run the checker on the sample without losing progress, or *)
+                  (€ (ξ (S i)) ∗ Ψ p ∗ (Ψ p -∗ WP checker (Val s) @ E [{fun v => Ψ p ∗ (⌜v = #false⌝ ∨ (⌜v = #true⌝ ∗ Θ s))}])) ∨
+                   (*... we lose progress & amplify error, and can run the checker on the sample without losing progress. *)
+                  (∃ p', ⌜(p' <= L)%nat ⌝ ∗ € (ξ i) ∗ Ψ p' ∗ (Ψ p' -∗ WP checker (Val s) @ E [{fun v => Ψ p' ∗ (⌜v = #false⌝ ∨ (⌜v = #true⌝ ∗ Θ s))}]))}]))%I.
+
+
+  Lemma ho_incremental_ubdd_approx_safe (sampler checker : val) Ψ ξ L E i iL p Θ :
+    ⊢ ⌜(p <= L)%nat ⌝ ∗
+      ⌜(i < iL)%nat ⌝ ∗
+    incr_sampling_scheme_spec sampler checker Ψ ξ L iL E Θ ∗
+    € (ξ i) ∗
+    Ψ p -∗
+    WP (gen_rejection_sampler sampler checker) @ E [{ fun v => Θ v }].
+  Proof.
+    rewrite /incr_sampling_scheme_spec.
+    iIntros "(%Hl&%Hil&(Hfinal&#Hamp)&Hcr&HΨ)".
+    rewrite /gen_rejection_sampler.
+    do 7 wp_pure.
+    iRevert (Hl).
+    iInduction i as [|i'] "IHerror" forall (p).
+    - (* base case for error credits *)
+      iIntros "%Hl".
+      wp_pures.
+      wp_bind (sampler _).
+      wp_apply (ub_twp_wand with "[Hfinal Hcr]"); first (iApply "Hfinal"; iFrame).
+      iIntros (s) "Hcheck"; wp_pures.
+      wp_apply (ub_twp_wand with "Hcheck").
+      iIntros (v) "(-> & HΘ)"; wp_pures.
+      eauto.
+    - (* inductive case for error credits *)
+      iIntros "%Hl".
+      iInduction p as [|p'] "IHp".
+      + (* base case for progress measure *)
+        wp_pures.
+        wp_bind (sampler _).
+        wp_apply (ub_twp_wand with "[Hfinal HΨ]"); first (iApply "Hfinal"; iFrame).
+        iIntros (s) "Hcheck"; wp_pures.
+        wp_apply (ub_twp_wand with "Hcheck").
+        iIntros (v) "(-> & HΘ)"; wp_pures.
+        eauto.
+      + (* Inductive case for progress measure *)
+        wp_pures.
+        wp_bind (sampler _).
+        wp_apply (ub_twp_wand with "[Hamp Hcr HΨ]"); first (iApply "Hamp"; iFrame; eauto).
+        iIntros (s) "[Hluck | [(Hcr&HΨ&Hcheck)|[%p'' (%Hp''&Hcr&HΨ&Hcheck)]]]".
+        * (* luck *)
+          wp_pures.
+          wp_bind (checker _).
+          wp_apply (ub_twp_wand with "Hluck").
+          iIntros (?) "(-> & HΘ)".
+          wp_pures.
+          eauto.
+        * (* progress *)
+          wp_pures.
+          wp_bind (checker _).
+          wp_apply (ub_twp_wand with "[Hcheck HΨ]"); first (iApply ("Hcheck" with "[$]")).
+          iIntros (r) "(HΨ&[-> | (-> & A)])".
+          -- (* not lucky: checker rejects *)
+             wp_pure. iApply ("IHp" with "[] Hfinal Hcr HΨ").
+             iPureIntro. lia.
+          -- (* lucky: checker accepts *)
+             wp_pures. eauto.
+        * (* amplification *)
+          wp_pures.
+          wp_bind (checker _).
+          wp_apply (ub_twp_wand with "[Hcheck HΨ]"); first (iApply ("Hcheck" with "[$]")).
+          iIntros (r) "(HΨ&[-> | (-> & A)])".
+          -- (* not lucky: checker rejects *)
+             assert (HiL' : (i' < iL)%nat) by lia.
+             wp_pure. iApply ("IHerror" $! HiL' with "Hfinal Hcr HΨ"). eauto.
+          -- (* lucky: checker accepts *)
+             wp_pures. eauto.
+    Qed.
+End incremental_spec.
diff --git a/theories/examples/approximate_samplers/approx_higherorder_rejection_sampler.v b/theories/examples/approximate_samplers/approx_higherorder_rejection_sampler.v
index e82461f5c..61f54ac07 100644
--- a/theories/examples/approximate_samplers/approx_higherorder_rejection_sampler.v
+++ b/theories/examples/approximate_samplers/approx_higherorder_rejection_sampler.v
@@ -9,47 +9,42 @@ Section spec.
   Local Open Scope R.
   Context `{!ub_clutchGS Σ}.
 
-  Definition sampling_scheme_spec (sampler checker : val) 𝜀factor 𝜀final E : iProp Σ
+  (* Spec for samplers which need to amplify some number of times, and then spend.
+     In particular, this is a spec for bounded samplers. *)
+  Definition bounded_sampling_scheme_spec (sampler checker : val) 𝜀factor 𝜀final E Θ : iProp Σ
     := ((∀ 𝜀,
-          {{{ € 𝜀 }}}
+          [[{ € 𝜀 }]]
             ((Val sampler) #())%E @ E
-          {{{ (v : val), RET v;
-               ((WP ((Val checker) v) @ E {{ λ v', ⌜v' = #true ⌝ }}) ∨
-               (∃ 𝜀', € 𝜀' ∗ ⌜𝜀 <= 𝜀' * 𝜀factor ⌝ ∗ (WP ((Val checker) v) @ E {{ λ v', ⌜v' = #true \/ v' = #false⌝ }})))}}}) ∗
+          [[{ (v : val), RET v;
+               ((WP ((Val checker) v) @ E [{ λ v', ⌜v' = #true ⌝ ∗ Θ v }]) ∨
+               (∃ 𝜀', € 𝜀' ∗ ⌜𝜀 <= 𝜀' * 𝜀factor ⌝ ∗ (WP ((Val checker) v) @ E [{ λ v', ⌜v' = #false⌝ ∨ (⌜v' = #true ⌝ ∗ Θ v)}])))}]]) ∗
         (∀ v : val,
-          {{{ € 𝜀final }}}
+          [[{ € 𝜀final }]]
             ((Val sampler) v) @ E
-          {{{ (v' : val), RET v'; (WP ((Val checker) v') @ E {{ λ v', ⌜v' = #true ⌝ }})}}}))%I.
-
-End spec.
-
-
-Section accounting.
-  (** Specification for general (stateful) bounded rejection samplers which makes use of
-      Iris' higher order Hoare triples *)
-  Local Open Scope R.
-  Context `{!ub_clutchGS Σ}.
-
-
-  Program Definition generic_geometric_error (r 𝜀final : nonnegreal) (depth : nat) : nonnegreal
-    := (𝜀final * (mknonnegreal (r ^ depth) _))%NNR.
-  Next Obligation. intros. apply pow_le. by destruct r. Qed.
-
-  Lemma final_generic_geometric_error (r 𝜀final : nonnegreal) : (generic_geometric_error r 𝜀final 0%nat) = 𝜀final.
-  Proof. apply nnreal_ext; by rewrite /generic_geometric_error /= Rmult_1_r. Qed.
+          [[{ (v' : val), RET v'; (WP ((Val checker) v') @ E [{ λ v'', ⌜v'' = #true ⌝ ∗ Θ v' }])}]]))%I.
+
+  (* Easier to explain spec.
+     It is allowed to wait arbitrarily long before spending,
+     and thus can wait to spend credit €1. *)
+  Definition sampling_scheme_spec (sampler checker : val) εfactor E Θ : iProp Σ
+    := (∀ ε,
+       [[{ € ε }]]
+            ((Val sampler) #())%E @ E
+       [[{ (v : val), RET v;
+               ((WP ((Val checker) v) @ E [{ λ v', ⌜v' = #true ⌝ ∗ Θ v }]) ∨
+               (∃ ε', € ε' ∗ ⌜ε <= ε' * εfactor ⌝ ∗ (WP ((Val checker) v) @ E [{ λ v', ⌜v' = #false⌝ ∨ (⌜v' = #true ⌝ ∗ Θ v)}])))}]])%I.
 
-  Lemma simpl_generic_geometric_error (r 𝜀final : nonnegreal) (depth : nat) :
-    (not (eq (nonneg r) 0)) ->
-    (nnreal_div (generic_geometric_error r 𝜀final (S depth)) r)%NNR = (generic_geometric_error r 𝜀final  depth).
+  Definition unbounded_scheme_implies_bounded sampler checker (εfactor : nonnegreal) E Θ :
+    ⊢ sampling_scheme_spec sampler checker εfactor E Θ -∗
+      bounded_sampling_scheme_spec sampler checker εfactor nnreal_one E Θ.
   Proof.
-    intros.
-    rewrite /generic_geometric_error /nnreal_div /nnreal_mult.
-    apply  nnreal_ext; simpl.
-    rewrite Rmult_assoc;  apply Rmult_eq_compat_l.
-    rewrite -Rmult_comm -Rmult_assoc Rinv_l; [lra|auto].
+    iIntros "?".
+    iSplitL; iFrame.
+    iIntros (? ?) "!> ?".
+    iExFalso; iApply (ec_spend with "[$]").
+    simpl; lra.
   Qed.
-
-End accounting.
+End spec.
 
 
 Section safety.
@@ -57,12 +52,12 @@ Section safety.
   Context `{!ub_clutchGS Σ}.
 
   (* safety for higher-order bounded rejection samplers *)
-  Lemma ho_bounded_approx_safe (sampler checker : val) (r εfinal : nonnegreal) (depth : nat) E :
+  Lemma ho_bounded_approx_safe (sampler checker : val) (r εfinal : nonnegreal) (depth : nat) E Θ :
     (0 < r < 1) ->
     (0 < εfinal < 1) ->
-    sampling_scheme_spec sampler checker r εfinal E -∗
+    bounded_sampling_scheme_spec sampler checker r εfinal E Θ -∗
     € (generic_geometric_error r εfinal depth) -∗
-    (WP (gen_bounded_rejection_sampler #(S depth) sampler checker) @ E {{ fun v => ∃ v', ⌜ v = SOMEV v' ⌝}})%I.
+    (WP (gen_bounded_rejection_sampler #(S depth) sampler checker) @ E [{ fun v => Θ v }])%I.
   Proof.
     (* initial setup *)
     rewrite /sampling_scheme_spec.
@@ -75,110 +70,93 @@ Section safety.
       { iApply (ec_weaken with "Hcr"); rewrite /generic_geometric_error /=; lra. }
       iIntros (next_sample) "Hcheck_accept".
       wp_pures; wp_bind (checker next_sample)%E.
-      iApply (ub_wp_wand with "Hcheck_accept").
-      iIntros (?) "#->"; wp_pures.
-      iModIntro; iExists next_sample; iFrame; auto.
+      iApply (ub_twp_wand with "Hcheck_accept").
+      iIntros (?) "(#-> & ?)"; wp_pures.
+      iModIntro; iFrame.
     - wp_pures.
       replace (bool_decide _) with false; last (symmetry; apply bool_decide_eq_false; lia).
       wp_pures; wp_bind (sampler #())%E.
       iApply ("Hamplify" $! (generic_geometric_error r εfinal (S depth')) with "Hcr").
-      iIntros (next_sample) "!> [Hcheck_accept|[%ε'(Hcr&%Hε'&Hcheck_reject)]]"; wp_pures.
+      iIntros (next_sample) "[Hcheck_accept|[%ε'(Hcr&%Hε'&Hcheck_reject)]]"; wp_pures.
       + wp_bind (checker next_sample)%V.
-        iApply (ub_wp_wand with "Hcheck_accept").
-        iIntros (?) "#->"; wp_pures.
-        iModIntro; iExists next_sample; iFrame; auto.
+        iApply (ub_twp_wand with "Hcheck_accept").
+        iIntros (?) "(#-> & ?)"; wp_pures.
+        iModIntro; iFrame.
       + wp_bind (checker next_sample)%V.
-        iApply (ub_wp_wand with "Hcheck_reject").
-        iIntros (?) "%Hresult".
+        iApply (ub_twp_wand with "Hcheck_reject").
+        iIntros (?) "Hresult".
         iSpecialize ("IH" with "[Hcr]").
         * iApply (ec_spend_le_irrel with "Hcr").
           rewrite /generic_geometric_error /=.
           apply (Rmult_le_reg_r r); auto.
           by rewrite /generic_geometric_error /=
                      (Rmult_comm r _) -Rmult_assoc in Hε'.
-        * destruct Hresult as [-> | ->].
-          { wp_pures; eauto. }
+        * iDestruct "Hresult" as "[-> | (-> & ?)]".
           { do 2 wp_pure.
             iClear "#".
             replace #((S (S depth')) - 1) with #(S depth'); [| do 2 f_equal; lia].
             iApply "IH". }
+          { wp_pures; eauto. }
   Qed.
 
 
-  (* safety for higher-order unbounded rejection samplers (almost the same proof) *)
-  Lemma ho_ubdd_approx_safe (sampler checker : val) (r εfinal : nonnegreal) (depth : nat) E :
+  (* safety for higher-order unbounded rejection samplers *)
+  Lemma ho_ubdd_approx_safe (sampler checker : val) (r : nonnegreal) (depth : nat) E Θ :
     (0 < r < 1) ->
-    (0 < εfinal < 1) ->
-    sampling_scheme_spec sampler checker r εfinal E -∗
-    ▷ € (generic_geometric_error r εfinal depth) -∗
-    (WP (gen_rejection_sampler sampler checker) @ E {{ fun v => ∃ v', ⌜ v = SOMEV v' ⌝}})%I.
+    sampling_scheme_spec sampler checker r E Θ -∗
+    € (generic_geometric_error r nnreal_one depth) -∗
+    (WP (gen_rejection_sampler sampler checker) @ E [{ fun v => Θ v }])%I.
   Proof.
     rewrite /sampling_scheme_spec.
-    iIntros ([Hr_pos Hr] [Hεfinal_pos Hεfinal]) "(#Hamplify&#Haccept) Hcr".
+    iIntros ([Hr_pos Hr]) "#Hamplify Hcr".
     rewrite /gen_rejection_sampler.
     do 7 wp_pure.
     iInduction depth as [|depth' Hdepth'] "IH".
     - wp_pures; wp_bind (sampler #())%E.
-      wp_apply ("Haccept" with "[Hcr]").
-      { iApply (ec_weaken with "Hcr"); rewrite /generic_geometric_error /=; lra. }
-      iIntros (next_sample) "Hcheck_accept".
-      wp_pures; wp_bind (checker next_sample)%E.
-      iApply (ub_wp_wand with "Hcheck_accept").
-      iIntros (?) "#->"; wp_pures.
-      iModIntro; iExists next_sample; iFrame; auto.
+      iExFalso; iApply (ec_spend with "[$]").
+      rewrite /generic_geometric_error /=; lra.
     - wp_pures.
       wp_pures; wp_bind (sampler #())%E.
-      iApply ("Hamplify" $! (generic_geometric_error r εfinal (S depth')) with "Hcr").
-      iIntros (next_sample) "!> [Hcheck_accept|[%ε'(Hcr&%Hε'&Hcheck_reject)]]"; wp_pures.
+      iApply ("Hamplify" $! (generic_geometric_error r nnreal_one (S depth')) with "Hcr").
+      iIntros (next_sample) "[Hcheck_accept|[%ε'(Hcr&%Hε'&Hcheck_reject)]]"; wp_pures.
       + wp_bind (checker next_sample)%V.
-        iApply (ub_wp_wand with "Hcheck_accept").
-        iIntros (?) "#->"; wp_pures.
-        iModIntro; iExists next_sample; iFrame; auto.
+        iApply (ub_twp_wand with "Hcheck_accept").
+        iIntros (?) "(#-> & ?)"; wp_pures.
+        iModIntro; iFrame.
       + wp_bind (checker next_sample)%V.
-        iApply (ub_wp_wand with "Hcheck_reject").
-        iIntros (?) "%Hresult".
+        iApply (ub_twp_wand with "Hcheck_reject").
+        iIntros (?) "Hresult".
         iSpecialize ("IH" with "[Hcr]").
         * iApply (ec_spend_le_irrel with "Hcr").
           rewrite /generic_geometric_error /=.
           apply (Rmult_le_reg_r r); auto.
           by rewrite /generic_geometric_error /=
                      (Rmult_comm r _) -Rmult_assoc in Hε'.
-        * destruct Hresult as [-> | ->].
-          { wp_pures; eauto. }
+        * iDestruct "Hresult" as "[-> | (-> & ?)]".
           { wp_pure.
             iClear "#".
             replace #((S (S depth')) - 1) with #(S depth'); [| do 2 f_equal; lia].
             iApply "IH". }
+          { wp_pures; eauto. }
   Qed.
 
 
   (** Limiting argument: any amount of credit suffices to show the unbounded sampler is safe *)
-  Lemma ho_ubdd_safe (sampler checker : val) (r εfinal ε : nonnegreal) E :
+  Lemma ho_ubdd_safe (sampler checker : val) (r ε : nonnegreal) E Θ :
     (0 < r < 1) ->
-    (0 < εfinal < 1) ->
     0 < ε ->
-    ⊢ sampling_scheme_spec sampler checker r εfinal E -∗
+    ⊢ sampling_scheme_spec sampler checker r E Θ -∗
       €ε -∗
-      WP gen_rejection_sampler sampler checker @ E {{ v, ∃ v', ⌜ v = SOMEV v' ⌝ }}.
+      WP gen_rejection_sampler sampler checker @ E [{ v, Θ v }].
   Proof.
-    iIntros (? ? ?) "#Hspec Hcr".
-    remember (/ NNRbar_to_real (NNRbar.Finite εfinal) * nonneg ε) as p.
-    assert (Hp : (0 < p)).
-    { rewrite Heqp.
-      apply Rmult_lt_0_compat; try done.
-      apply Rinv_0_lt_compat.
-      by destruct H0. }
-    assert (H' : r < 1); first by destruct H.
-    destruct (error_limit' r H' (mkposreal p Hp)) as [depth Hlim].
-    iApply (ho_ubdd_approx_safe _ _ _ _ depth);[|done|done|];[done|]. (* weird unification order *)
-    iApply (ec_spend_le_irrel with "Hcr").
+    iIntros ([? Hr] Hε) "#Hspec Hcr".
     rewrite /generic_geometric_error /=.
-    apply (Rmult_le_reg_l (/ εfinal)).
-    { apply Rinv_0_lt_compat; by destruct H0. }
-    rewrite /= Heqp in Hlim.
-    rewrite -Rmult_assoc Rinv_l; last lra.
-    rewrite Rmult_1_l.
-    by apply Rlt_le.
+    destruct (error_limit' r Hr (mkposreal (nonneg ε) Hε)) as [d Hlim].
+    iApply (ho_ubdd_approx_safe with "[]"); [eauto| iApply "Hspec" |].
+    iApply (ec_spend_le_irrel with "Hcr").
+    rewrite /generic_geometric_error /= Rmult_1_l.
+    apply Rlt_le.
+    simpl in Hlim; eauto.
   Qed.
 End safety.
 
@@ -216,79 +194,46 @@ Section higherorder_rand.
     - apply not_0_INR; lia.
   Qed.
 
+  Definition rand_sampling_scheme_goal (n' : nat) (v : val) : iProp Σ := ∃ (n : nat), ⌜v = #n /\ (n <= n')%nat⌝.
+
   Lemma rand_sampling_scheme_spec (n' m' : nat) (Hnm : (n' < m')%nat) E :
     ⊢ sampling_scheme_spec
           (λ: "_", rand #m')%V
           (λ: "sample", "sample" ≤ #n')%V
           (err_factor (S n') (S m'))
-          (err_factor (S n') (S m'))
-          E.
+          E
+          (rand_sampling_scheme_goal n').
   Proof.
     Opaque INR.
     rewrite /sampling_scheme_spec.
-    iStartProof; iSplit.
-    - (* sampling rule *)
-      iIntros (ε Φ) "!> Hcr HΦ"; wp_pures.
-      iApply (wp_couple_rand_adv_comp  m' _ _ ε (rand_ε2 n' m' ε) _ with "Hcr").
-      { (* uniform bound *)
-        eexists (nnreal_div ε (err_factor (S n') (S m'))); intros s.
-        rewrite /rand_ε2.
-        case_bool_decide; simpl; [|lra].
-        apply Rmult_le_pos.
-        - by apply cond_nonneg.
-        - apply Rlt_le, Rinv_0_lt_compat, Rmult_lt_0_compat.
-          + apply lt_0_INR. lia.
-          + apply Rinv_0_lt_compat. apply lt_0_INR. lia.
-      }
-
-      { (* series convergence *)
-        by apply sample_err_mean_higherorder. }
-
-      iNext; iIntros (s) "Hcr".
-      iApply "HΦ".
-      destruct (le_gt_dec s n') as [|] eqn:Hdec; [iLeft | iRight].
-      + (* the sample is inbounds, the checker should accept *)
-        wp_pures; iModIntro; iPureIntro.
-        do 2 f_equal; apply bool_decide_true; lia.
-      + (* the sample is out of bounds *)
-        iExists _; iSplitL; first iFrame.
-        iSplit.
-        * (* credit is amplified *)
-          iPureIntro.
-          rewrite /rand_ε2 bool_decide_eq_false_2; last first.
-          { rewrite /not; intros; lia. }
-          rewrite /nnreal_div Rmult_assoc /nnreal_inv; simpl.
-          rewrite Rinv_l; [lra|].
-          apply Rmult_integral_contrapositive; split.
-          -- apply Rgt_not_eq, Rlt_gt, lt_0_INR; lia.
-          -- apply Rinv_neq_0_compat.
-             apply Rgt_not_eq, Rlt_gt, lt_0_INR; lia.
-        * (* sampler rejects *)
-          wp_pures; iModIntro; iPureIntro. right.
-          do 2 f_equal; apply bool_decide_false; lia.
-
-    - (* checking rule *)
-      iIntros (s Φ) "!> Hcr HΦ"; wp_pures.
-      wp_apply (wp_rand_err_list_nat _ m' (seq (S n') ((S m') - (S n')))).
-      iSplitL "Hcr".
-      + (* credit accounting *)
-        iApply (ec_spend_irrel with "Hcr").
-        rewrite /err_factor seq_length /=.
-        apply Rmult_eq_compat_l.
-        do 2 f_equal; lia.
-      + iIntros (s') "%Hs'".
-        iApply "HΦ"; wp_pures.
-        iModIntro; iPureIntro.
-        do 2 f_equal; apply bool_decide_true.
-        rewrite List.Forall_forall in Hs'.
-        specialize Hs' with s'.
-        apply Znot_gt_le.
-        intros Hcont; apply Hs'; last reflexivity.
-        rewrite in_seq.
-        split; first lia.
-        replace (S n' + (S m' - S n'))%nat with (S m') by lia.
-        specialize (fin_to_nat_lt s'); by lia.
-
+    iStartProof.
+    iIntros (ε Φ) "!> Hcr HΦ"; wp_pures.
+    iApply (twp_couple_rand_adv_comp1 m' _ _ ε (rand_ε2 n' m' ε) _ with "Hcr").
+    { by apply sample_err_mean_higherorder. }
+    iIntros (s) "Hcr".
+    iApply "HΦ".
+    destruct (le_gt_dec s n') as [|] eqn:Hdec; [iLeft | iRight].
+    + (* the sample is inbounds, the checker should accept *)
+      wp_pures; iModIntro; iPureIntro.
+      split.
+      { do 2 f_equal; apply bool_decide_true; lia. }
+      { eexists _. split; [by f_equal|done]. }
+    + (* the sample is out of bounds *)
+      iExists _; iSplitL; first iFrame.
+      iSplit.
+      * (* credit is amplified *)
+        iPureIntro.
+        rewrite /rand_ε2 bool_decide_eq_false_2; last first.
+        { rewrite /not; intros; lia. }
+        rewrite /nnreal_div Rmult_assoc /nnreal_inv; simpl.
+        rewrite Rinv_l; [lra|].
+        apply Rmult_integral_contrapositive; split.
+        -- apply Rgt_not_eq, Rlt_gt, lt_0_INR; lia.
+        -- apply Rinv_neq_0_compat.
+           apply Rgt_not_eq, Rlt_gt, lt_0_INR; lia.
+      * (* sampler rejects *)
+        wp_pures; iModIntro; iPureIntro. left.
+        { do 2 f_equal; apply bool_decide_false; lia. }
     Unshelve.
     { rewrite Nat2Z.id; apply TCEq_refl. }
   Qed.
@@ -296,8 +241,6 @@ Section higherorder_rand.
 End higherorder_rand.
 
 
-
-
 Section higherorder_flip2.
   (** Instantiation of the higher-order spec for a pair of coin flips *)
   Local Open Scope R.
@@ -321,22 +264,18 @@ Section higherorder_flip2.
     := (fun z => if (flip_is_1 z) then εh else εt).
 
   Lemma flip_amplification (ε1 εh εt : nonnegreal) (Hmean : (εh + εt) = 2 * ε1 ) E :
-    {{{ € ε1 }}}
+    [[{ € ε1 }]]
       rand #1 @ E
-    {{{ v, RET #v; ⌜(v = 0%nat) \/ (v = 1%nat) ⌝ ∗ € (scale_flip ε1 εh εt #v) }}}.
+    [[{ v, RET #v; ⌜(v = 0%nat) \/ (v = 1%nat) ⌝ ∗ € (scale_flip ε1 εh εt #v) }]].
   Proof.
     iIntros (Φ) "Hcr HΦ".
-    iApply (wp_couple_rand_adv_comp 1%nat  _ _ ε1 (ε2_flip1 ε1 εh εt) _ with "Hcr").
-    - (* uniform bound *)
-      exists (εh + εt)%NNR; intros n.
-      rewrite /ε2_flip1.
-      destruct (fin_to_bool n); destruct εt, εh; rewrite /bound /=; lra.
+    iApply (twp_couple_rand_adv_comp1 1%nat  _ _ ε1 (ε2_flip1 ε1 εh εt) _ with "Hcr").
     - (* series mean *)
       rewrite SeriesC_finite_foldr /enum /fin_finite /fin_enum /ε2_flip1 /=.
       rewrite Rplus_0_r -Rmult_plus_distr_l Rplus_comm Hmean /=.
       lra.
     - (* continutation *)
-      iNext. iIntros (n) "Hcr".
+      iIntros (n) "Hcr".
       iApply ("HΦ" $! _); iSplitR.
       + iPureIntro.
         inv_fin n; first (left; done).
@@ -350,85 +289,56 @@ Section higherorder_flip2.
       { apply TCEq_refl. }
   Qed.
 
+  Definition flip2_goal (v : val) : iProp Σ := ⌜v = (#1%nat, #1%nat)%V⌝.
 
   Lemma flip2_sampling_scheme_spec E :
     ⊢ sampling_scheme_spec
           (λ: "_", Pair (rand #1) (rand #1))
           (λ: "sample", (((Fst "sample") = #1) && ((Snd "sample") = #1)))
           (nnreal_div (nnreal_nat 3%nat) (nnreal_nat 4%nat))
-          (nnreal_div (nnreal_nat 3%nat) (nnreal_nat 4%nat))
-          E.
+          E
+          flip2_goal.
   Proof.
     rewrite /sampling_scheme_spec.
-    iStartProof; iSplit.
-    - (* amplification rule *)
-      iIntros (𝜀 Φ) "!> Hcr HΦ"; wp_pures.
-      wp_apply (flip_amplification 𝜀
-                  (nnreal_mult 𝜀 (nnreal_div (nnreal_nat 2) (nnreal_nat 3)))
-                  (nnreal_mult 𝜀 (nnreal_div (nnreal_nat 4) (nnreal_nat 3)))
-                   with "Hcr"); [simpl; lra| ].
+    iStartProof.
+    iIntros (𝜀 Φ) "!> Hcr HΦ"; wp_pures.
+    wp_apply (flip_amplification 𝜀
+                (nnreal_mult 𝜀 (nnreal_div (nnreal_nat 2) (nnreal_nat 3)))
+                (nnreal_mult 𝜀 (nnreal_div (nnreal_nat 4) (nnreal_nat 3)))
+                 with "Hcr"); [simpl; lra| ].
+    iIntros (v) "(%Hv&Hcr)".
+    destruct Hv as [-> | ->].
+    + (* first flip was zero, check is going to false and the second flip doesn't matter. *)
+      wp_bind (rand _)%E; iApply twp_rand; auto.
+      iIntros (v') "_"; wp_pures; iModIntro; iApply "HΦ".
+      iRight; iExists _.
+      iSplitL "Hcr"; [iFrame|].
+      iSplit.
+      * (* credit accounting *)
+        iPureIntro; simpl; lra.
+      * (* step the checker *)
+        wp_pures; case_bool_decide; wp_pures; auto.
+    + (* first flip was 1, we only have 2/3 error so we need to amplify up to 4/3
+          in the case that the second flip is not 1 *)
+      replace (scale_flip 𝜀 _ _ _) with (𝜀 * nnreal_div (nnreal_nat 2) (nnreal_nat 3))%NNR; last first.
+      { rewrite /scale_flip /flip_is_1 /=. by apply nnreal_ext. }
+      remember (𝜀 * nnreal_div (nnreal_nat 2) (nnreal_nat 3))%NNR as 𝜀'.
+      wp_bind (rand #1 )%E.
+      wp_apply (flip_amplification 𝜀' nnreal_zero (nnreal_mult 𝜀' (nnreal_nat 2)) with "Hcr").
+      { simpl. lra. }
       iIntros (v) "(%Hv&Hcr)".
       destruct Hv as [-> | ->].
-      + (* first flip was zero, check is going to false and the second flip doesn't matter. *)
-        wp_bind (rand _)%E; iApply wp_rand; auto.
-        iNext; iIntros (v') "_"; wp_pures; iModIntro; iApply "HΦ".
+      * (* second flip was zero *)
+        wp_pures; iModIntro; iApply "HΦ".
         iRight; iExists _.
         iSplitL "Hcr"; [iFrame|].
         iSplit.
-        * (* credit accounting *)
-          iPureIntro; simpl; lra.
-        * (* step the checker *)
-          wp_pures; case_bool_decide; wp_pures; auto.
-      + (* first flip was 1, we only have 2/3 error so we need to amplify up to 4/3
-            in the case that the second flip is not 1 *)
-        replace (scale_flip 𝜀 _ _ _) with (𝜀 * nnreal_div (nnreal_nat 2) (nnreal_nat 3))%NNR; last first.
-        { rewrite /scale_flip /flip_is_1 /=. by apply nnreal_ext. }
-        remember (𝜀 * nnreal_div (nnreal_nat 2) (nnreal_nat 3))%NNR as 𝜀'.
-        wp_bind (rand #1 )%E.
-        wp_apply (flip_amplification 𝜀' nnreal_zero (nnreal_mult 𝜀' (nnreal_nat 2)) with "Hcr").
-        { simpl. lra. }
-        iIntros (v) "(%Hv&Hcr)".
-        destruct Hv as [-> | ->].
-        * (* second flip was zero *)
-          wp_pures; iModIntro; iApply "HΦ".
-          iRight; iExists _.
-          iSplitL "Hcr"; [iFrame|].
-          iSplit.
-          -- (* credit accounting *)
-            iPureIntro; rewrite Heq𝜀' /=; lra.
-          -- (* step the checker *)
-            wp_pures; auto.
-        * (* both flips are 1, checker should accept *)
-          wp_pures; iModIntro; iApply "HΦ".
-          iLeft; wp_pures; auto.
-
-    - (* credit spending rule *)
-      iIntros (s Φ) "!> Hcr HΦ"; wp_pures.
-      wp_bind (rand #1)%E.
-
-      (* give € 1 to the 0 flip, and € 1/2 to the 1 flip *)
-      wp_apply (flip_amplification
-                  (nnreal_div (nnreal_nat 3) (nnreal_nat 4))
-                  (nnreal_div (nnreal_nat 1) (nnreal_nat 2))
-                  nnreal_one with "Hcr"); [simpl; lra|].
-      iIntros (v') "(%Hv'&Hcr)".
-      destruct Hv' as [-> | ->].
-      + (* first flip is zero: but we can spend € 1 to conclude *)
-        iApply (wp_ec_spend with "Hcr").
-        * rewrite /scale_flip /flip_is_1 /=; lra.
-        * rewrite /to_val; done.
-      +  (* we have € 1/2 so we can make the second flip be 1 too *)
-        wp_bind (rand #1)%E.
-        iApply (wp_rand_err _ _ 0%fin with "[Hcr HΦ]").
-        iSplitL "Hcr". { iApply (ec_spend_irrel with "Hcr"). rewrite /=; lra. }
-        iIntros (v') "%Hv'".
+        -- (* credit accounting *)
+          iPureIntro; rewrite Heq𝜀' /=; lra.
+        -- (* step the checker *)
+          wp_pures; auto.
+      * (* both flips are 1, checker should accept *)
         wp_pures; iModIntro; iApply "HΦ".
-        wp_pures; case_bool_decide; wp_pures; auto.
-        (* we have a contradiction in Hv' and H *)
-        exfalso.
-        rewrite /not in Hv', H.
-        inv_fin v'; first auto.
-        intros v'; inv_fin v'; first auto.
-        intros v'; by apply Fin.case0.
+        iLeft; wp_pures; auto.
   Qed.
 End higherorder_flip2.
diff --git a/theories/examples/approximate_samplers/approx_random_walk.v b/theories/examples/approximate_samplers/approx_random_walk.v
index 30b4f9e07..65ebc34e1 100644
--- a/theories/examples/approximate_samplers/approx_random_walk.v
+++ b/theories/examples/approximate_samplers/approx_random_walk.v
@@ -1,6 +1,7 @@
 (** * Error credit proof that an unbounded integer walk returns to the origin *)
 From clutch.ub_logic Require Export ub_clutch ub_rules.
 From clutch Require Export examples.approximate_samplers.approx_sampler_lib.
+From clutch Require Export examples.approximate_samplers.approx_higherorder_incremental.
 From Coquelicot Require Import Series.
 Require Import Lra.
 
@@ -105,8 +106,8 @@ Section integer_walk.
   (* sampler either gives us progress or amplifies our error *)
   Lemma wp_sampler_amp εᵢ εₐ l p i kwf E :
     ⊢ I εᵢ εₐ l kwf (S p) ∗ € (AX εᵢ εₐ kwf (S i)) -∗
-      WP (int_walk_sampler #l) @ E {{ fun _ => ((I εᵢ εₐ l kwf p ∗ € (AX εᵢ εₐ kwf (S i))) ∨
-                                     (I εᵢ εₐ l kwf (S (S p)) ∗ € (AX εᵢ εₐ kwf i)))}}.
+      WP (int_walk_sampler #l) @ E [{ fun _ => ((I εᵢ εₐ l kwf p ∗ € (AX εᵢ εₐ kwf (S i))) ∨
+                                     (I εᵢ εₐ l kwf (S (S p)) ∗ € (AX εᵢ εₐ kwf i)))}].
   Proof.
     iIntros "([%z (Hl & HcrIC & %Hz & HcrAC)] & HcrAX)".
     rewrite /int_walk_sampler.
@@ -131,7 +132,7 @@ Section integer_walk.
       excluded by virtue of the (S p) in the amp spec.
 
      *)
-    wp_apply (wp_couple_rand_adv_comp1 _ _ _
+    wp_apply (twp_couple_rand_adv_comp1 _ _ _
                 ((IC εᵢ (S p)) + (AC εᵢ εₐ (L εᵢ - S p) (I_obligation_1 εᵢ (S p)) kwf))%NNR
                 (integer_walk_distr εᵢ εₐ (S p) kwf) with "[HcrAC HcrIC]").
     { (* Series mean *)
@@ -220,33 +221,38 @@ Section integer_walk.
         (AX εᵢ εₐ kwf)
         (L εᵢ)
         B
-        E.
+        E
+        (fun _ => ⌜True ⌝).
   Proof.
     iSplit.
     - (* Spending rules *)
       iIntros "[Hcr | [%z (Hl & _ & %Hz & _)]]".
       + (* Credit spending rule *)
-        wp_apply (wp_ec_spend _ _ _ nnreal_one); simpl; [lra|eauto|].
+        wp_apply (twp_ec_spend _ _ _ nnreal_one); simpl; [lra|eauto|].
         iApply (ec_spend_le_irrel with "Hcr"); simpl.
         rewrite Rmult_0_l Rminus_0_r.
         apply Rmax_r.
       + (* Progress spending rule *)
         rewrite /int_walk_sampler; wp_pures.
-        wp_bind (rand _)%E; wp_apply wp_rand; eauto.
+        wp_bind (rand _)%E; wp_apply twp_rand; eauto.
         iIntros (n) "_"; wp_pures.
         rewrite /int_walk_checker.
         (* the rest of the symbolic execution doesn't change depeding on the value.  *)
         case_bool_decide; repeat (try wp_pures; try wp_load; try wp_store).
         * (* l ↦ -3 *)
-          iModIntro. iPureIntro. f_equal. f_equal.
+          iModIntro. iPureIntro.
+          split; eauto.
+          f_equal. f_equal.
           apply bool_decide_eq_true_2. lia.
         * (* l ↦ -1 *)
-          iModIntro. iPureIntro. f_equal. f_equal.
+          iModIntro. iPureIntro.
+          split; eauto.
+          f_equal. f_equal.
           apply bool_decide_eq_true_2. lia.
     - iModIntro.
       iIntros (i p) "(%Hpwf&%Hb&HcrAX&HI)".
       wp_pure.
-      wp_apply (ub_wp_wand with "[HcrAX HI]"); first iApply (wp_sampler_amp with "[$]").
+      wp_apply (ub_twp_wand with "[HcrAX HI]"); first iApply (wp_sampler_amp with "[$]").
       iIntros (s) "[(HI&HAX)|(HI&HAX)]".
       + (* progress *)
         iLeft; iFrame.
diff --git a/theories/examples/approximate_samplers/approx_rejection_sampler_presampled.v b/theories/examples/approximate_samplers/approx_rejection_sampler_presampled.v
index 4bd086a1c..1c52da729 100644
--- a/theories/examples/approximate_samplers/approx_rejection_sampler_presampled.v
+++ b/theories/examples/approximate_samplers/approx_rejection_sampler_presampled.v
@@ -20,14 +20,14 @@ Section presampled_flip2.
         gen_rejection_sampler
         (λ: "_", Pair (rand("α")#1) (rand("α") #1))
         (λ: "sample", (((Fst "sample") = #1) && ((Snd "sample") = #1)))
-        @ E {{ v, ∃ v', ⌜ v = SOMEV v' ⌝ }}.
+        @ E [{ v, ⌜ v = (#1%Z, #1%Z)%V ⌝ }].
   Proof.
     iIntros (?) "Hcr".
     wp_bind (alloc _)%I.
-    wp_apply (wp_alloc_tape); auto.
+    wp_apply (twp_alloc_tape); auto.
     iIntros (α) "Hα".
     rewrite Z2Nat.inj_succ; try lia.
-    wp_apply (presample_planner_aligned _ _ _ _ _ _ _ _ [1%fin; 1%fin]); auto; [apply H|].
+    wp_apply (twp_presample_planner_aligned _ _ _ _ _ _ _ _ [1%fin; 1%fin]); auto; [apply H|].
     iFrame.
     iIntros "[%junk Hα] /=".
     pose flip2_junk_inv k s : iProp Σ := (∃ j, α ↪ (S (Z.to_nat 0); j ++ s) ∗ ⌜length j = (2 * k)%nat ⌝)%I.
@@ -45,23 +45,24 @@ Section presampled_flip2.
       iDestruct "Hinv" as "[%j (Hα & %Hl)] /=".
       rewrite (nil_length_inv _ Hl) /=.
       wp_pures.
-      wp_bind (Rand _ _); wp_apply (wp_rand_tape with "Hα"); iIntros "Hα".
-      wp_bind (Rand _ _); wp_apply (wp_rand_tape with "Hα"); iIntros "Hα".
+      wp_bind (Rand _ _); wp_apply (twp_rand_tape with "Hα"); iIntros "Hα".
+      wp_bind (Rand _ _); wp_apply (twp_rand_tape with "Hα"); iIntros "Hα".
       wp_pures.
-      iModIntro; iExists _; iPureIntro. done.
+      iModIntro; eauto.
     - rewrite /flip2_junk_inv.
       iDestruct "Hinv" as "[%j (Hα & %Hl)] /=".
       rewrite Nat.mul_succ_r Nat.add_comm /= in Hl.
       destruct j as [| s0 j0]. { simpl in Hl. exfalso; lia. }
       destruct j0 as [| s1 j']. { simpl in Hl. exfalso; lia. }
       wp_pures.
-      wp_bind (Rand _ _); wp_apply (wp_rand_tape with "Hα"); iIntros "Hα".
-      wp_bind (Rand _ _); wp_apply (wp_rand_tape with "Hα"); iIntros "Hα".
+      wp_bind (Rand _ _); wp_apply (twp_rand_tape with "Hα"); iIntros "Hα".
+      wp_bind (Rand _ _); wp_apply (twp_rand_tape with "Hα"); iIntros "Hα".
       iSpecialize ("IH" with "[Hα]").
       { iExists _; iFrame; iPureIntro. do 2 rewrite cons_length in Hl. congruence. }
       wp_pures.
       case_bool_decide; [wp_pures; case_bool_decide|].
-      + wp_pures; iModIntro; iExists _; auto.
+      + wp_pures. iModIntro; iPureIntro.
+        rewrite H0 H1. done.
       + wp_pure. iApply "IH".
       + do 2 wp_pure; iApply "IH".
   Qed.
diff --git a/theories/examples/approximate_samplers/approx_sampler_lib.v b/theories/examples/approximate_samplers/approx_sampler_lib.v
index 7e2480511..3325caa91 100644
--- a/theories/examples/approximate_samplers/approx_sampler_lib.v
+++ b/theories/examples/approximate_samplers/approx_sampler_lib.v
@@ -6,16 +6,19 @@ Require Import Lra.
 Set Default Proof Using "Type*".
 
 Section samplers.
+  (* Placeholder stuck expression *)
+  Definition error : expr := (#42 #42)%E.
+
   (* higher order boundeded rejection sampler *)
   Definition gen_bounded_rejection_sampler :=
     (λ: "depth" "sampler" "checker",
         let: "do_sample" :=
           (rec: "f" "tries_left" :=
               if: ("tries_left" - #1) < #0
-                then NONE
+                then error
                 else let: "next_sample" := ("sampler" #()) in
                      if: ("checker" "next_sample")
-                        then SOME "next_sample"
+                        then "next_sample"
                         else "f" ("tries_left" - #1))
         in "do_sample" "depth")%E.
 
@@ -27,7 +30,7 @@ Section samplers.
           (rec: "f" "_" :=
              let: "next_sample" := ("sampler" #()) in
               if: ("checker" "next_sample")
-                  then SOME "next_sample"
+                  then "next_sample"
                   else "f" #())
         in "do_sample" #())%E.
 End samplers.
@@ -84,16 +87,12 @@ Section finSeries.
   Qed.
 End finSeries.
 
-
-
-
-
 Section accounting.
   Local Open Scope R.
 
+  (* Amplification factor for rejecting n of m possible outcomes *)
   Definition err_factor n m : nonnegreal := (nnreal_div (nnreal_nat (m - n)%nat) (nnreal_nat m%nat)).
 
-
   Lemma err_factor_lt1 n m : (0 < n)%nat -> (n < m)%nat -> err_factor n m < 1.
   Proof.
     intros ? ?.
@@ -117,7 +116,7 @@ Section accounting.
   Qed.
 
 
-  (* closed form for the error in the bounded sampler, with a given number of tries remaining *)
+  (* Combined factor after a series of``depth`` amplification *)
   Program Definition bdd_cf_error n m depth (Hnm : (n < m)%nat) := mknonnegreal ((err_factor n m) ^ depth) _.
   Next Obligation.
     intros.
@@ -128,7 +127,7 @@ Section accounting.
   Qed.
 
 
-  (* distribution of error mass ε₁ for a given sample:
+  (* Error credit distribution for a basic ``rand`` rejection sampler.
       - zero error given to cases which are inbounds
       - uniform error to the recursive case *)
   Definition bdd_cf_sampling_error n m ε₁ : (fin m) -> nonnegreal
@@ -138,7 +137,6 @@ Section accounting.
             else (nnreal_div ε₁ (err_factor n m)).
 
 
-  (* simplify amplified errors *)
   Lemma simplify_amp_err (n m d': nat) (s : fin m) Hnm :
     (s <? n)%nat = false ->
     bdd_cf_sampling_error n m (bdd_cf_error n m (S d') Hnm) s = (bdd_cf_error n m d' Hnm).
@@ -167,7 +165,7 @@ Section accounting.
   Qed.
 
 
-  (* error distribution is bounded above for each possible sample *)
+  (* Error distribution upper bound *)
   Lemma sample_err_wf n m d (s : fin m) Hnm :
     (0 < n)%nat ->
     bdd_cf_sampling_error n m (bdd_cf_error n m (S d) Hnm) s <= 1.
@@ -189,7 +187,7 @@ Section accounting.
   Qed.
 
 
-  (* mean of error distribution is preserved *)
+  (* Mean of error distribution is preserved *)
   Lemma sample_err_mean n' m' (Hnm : (n' < m')%nat) 𝜀₁ :
     SeriesC (λ s : fin (S m'), (1 / S m') * bdd_cf_sampling_error (S n') (S m') 𝜀₁ s) = 𝜀₁.
   Proof.
@@ -238,118 +236,21 @@ Section accounting.
   Qed.
 
 
-End accounting.
-
-
-
-
-Section incremental_spec.
-  Local Open Scope R.
-  Context `{!ub_clutchGS Σ}.
-
-  (* Ψ : state
-     ξ : error
-     L : progress bound
-   *)
-  Definition incr_sampling_scheme_spec (sampler checker : val) (Ψ : nat -> iProp Σ) (ξ : nat -> nonnegreal) L iL E : iProp Σ :=
-    ( (* Either 0 credit or 0 progress => we will sample a value which the checker accepts
-         Allowed to consume (or invalidate Ψ) in this process *)
-      ((€ (ξ 0%nat) ∨ Ψ 0%nat) -∗ WP sampler #() @ E {{ fun s => WP checker (Val s) @ E {{fun v => ⌜v = #true⌝}}}}) ∗
-      (* Given any amount of credit and progress, we can get a sample such that... *)
-     □ (∀ i p, (⌜((S p) <= L)%nat ⌝ ∗ ⌜((S i) < iL)%nat ⌝ ∗ € (ξ (S i)) ∗ Ψ (S p)) -∗
-            WP sampler #() @ E {{ fun s =>
-                   (*...we're done by chance alone, or... *)
-                  (WP checker (Val s) @ E {{fun v => ⌜v = #true⌝}}) ∨
-                   (*... we make prgress, and can run the checker on the sample without losing progress, or *)
-                  (€ (ξ (S i)) ∗ Ψ p ∗ (Ψ p -∗ WP checker (Val s) @ E {{fun v => Ψ p ∗ ∃ b: bool, ⌜v = #b⌝}})) ∨
-                   (*... we lose progress & amplify error, and can run the checker on the sample without losing progress. *)
-                  (∃ p', ⌜(p' <= L)%nat ⌝ ∗ € (ξ i) ∗ Ψ p' ∗ (Ψ p' -∗ WP checker (Val s) @ E {{fun v => Ψ p' ∗ ∃ b: bool, ⌜v = #b⌝}}))}}))%I.
+  Program Definition generic_geometric_error (r 𝜀final : nonnegreal) (depth : nat) : nonnegreal
+    := (𝜀final * (mknonnegreal (r ^ depth) _))%NNR.
+  Next Obligation. intros. apply pow_le. by destruct r. Qed.
 
+  Lemma final_generic_geometric_error (r 𝜀final : nonnegreal) : (generic_geometric_error r 𝜀final 0%nat) = 𝜀final.
+  Proof. apply nnreal_ext; by rewrite /generic_geometric_error /= Rmult_1_r. Qed.
 
-  Lemma ho_incremental_ubdd_approx_safe (sampler checker : val) Ψ ξ L E i iL p :
-    ⊢ ⌜(p <= L)%nat ⌝ ∗
-      ⌜(i < iL)%nat ⌝ ∗
-    incr_sampling_scheme_spec sampler checker Ψ ξ L iL E ∗
-    € (ξ i) ∗
-    Ψ p -∗
-    WP (gen_rejection_sampler sampler checker) @ E {{ fun v => ∃ v', ⌜ v = SOMEV v' ⌝}}.
+  Lemma simpl_generic_geometric_error (r 𝜀final : nonnegreal) (depth : nat) :
+    (not (eq (nonneg r) 0)) ->
+    (nnreal_div (generic_geometric_error r 𝜀final (S depth)) r)%NNR = (generic_geometric_error r 𝜀final  depth).
   Proof.
-    rewrite /incr_sampling_scheme_spec.
-    iIntros "(%Hl&%Hil&(Hfinal&#Hamp)&Hcr&HΨ)".
-    rewrite /gen_rejection_sampler.
-    do 7 wp_pure.
-    iRevert (Hl).
-    iInduction i as [|i'] "IHerror" forall (p).
-    - (* base case for error credits *)
-      iIntros "%Hl".
-      wp_pures.
-      wp_bind (sampler _).
-      wp_apply (ub_wp_wand with "[Hfinal Hcr]"); first (iApply "Hfinal"; iFrame).
-      iIntros (s) "Hcheck"; wp_pures.
-      wp_apply (ub_wp_wand with "Hcheck").
-      iIntros (v) "->"; wp_pures.
-      eauto.
-    - (* inductive case for error credits *)
-      iIntros "%Hl".
-      iInduction p as [|p'] "IHp".
-      + (* base case for progress measure *)
-        wp_pures.
-        wp_bind (sampler _).
-        wp_apply (ub_wp_wand with "[Hfinal HΨ]"); first (iApply "Hfinal"; iFrame).
-        iIntros (s) "Hcheck"; wp_pures.
-        wp_apply (ub_wp_wand with "Hcheck").
-        iIntros (v) "->"; wp_pures.
-        eauto.
-      + (* Inductive case for progress measure *)
-        wp_pures.
-        wp_bind (sampler _).
-        wp_apply (ub_wp_wand with "[Hamp Hcr HΨ]"); first (iApply "Hamp"; iFrame; eauto).
-        iIntros (s) "[Hluck | [(Hcr&HΨ&Hcheck)|[%p'' (%Hp''&Hcr&HΨ&Hcheck)]]]".
-        * (* luck *)
-          wp_pures.
-          wp_bind (checker _).
-          wp_apply (ub_wp_wand with "Hluck").
-          iIntros (?) "->".
-          wp_pures.
-          eauto.
-        * (* progress *)
-          wp_pures.
-          wp_bind (checker _).
-          wp_apply (ub_wp_wand with "[Hcheck HΨ]"); first (iApply ("Hcheck" with "[$]")).
-          iIntros (r) "(HΨ&[%b ->])".
-          destruct b as [|].
-          -- (* lucky: checker accepts *)
-             wp_pures. eauto.
-          -- (* not lucky: checker rejects *)
-             wp_pure. iApply ("IHp" with "[] Hfinal Hcr HΨ").
-             iPureIntro. lia.
-        * (* amplification *)
-          wp_pures.
-          wp_bind (checker _).
-          wp_apply (ub_wp_wand with "[Hcheck HΨ]"); first (iApply ("Hcheck" with "[$]")).
-          iIntros (r) "(HΨ&[%b ->])".
-          destruct b as [|].
-          -- (* lucky: checker accepts *)
-             wp_pures. eauto.
-          -- (* not lucky: checker rejects *)
-             assert (HiL' : (i' < iL)%nat) by lia.
-             wp_pure. iApply ("IHerror" $! HiL' with "Hfinal Hcr HΨ"). eauto.
-    Qed.
-
-
-End incremental_spec.
-
-
-Section remove_me.
-
-  Local Open Scope R.
-  Context `{!ub_clutchGS Σ}.
-
-  Lemma wp_ec_spend e E Φ ε :
-    (1 <= ε.(nonneg))%R →
-    (to_val e = None) ->
-    € ε -∗ WP e @ E {{ Φ }}.
-  Proof. Admitted.
-
-
-End remove_me.
+    intros.
+    rewrite /generic_geometric_error /nnreal_div /nnreal_mult.
+    apply  nnreal_ext; simpl.
+    rewrite Rmult_assoc;  apply Rmult_eq_compat_l.
+    rewrite -Rmult_comm -Rmult_assoc Rinv_l; [lra|auto].
+  Qed.
+End accounting.
diff --git a/theories/examples/approximate_samplers/approx_walkSAT.v b/theories/examples/approximate_samplers/approx_walkSAT.v
index 1f8cab789..5af2b449b 100644
--- a/theories/examples/approximate_samplers/approx_walkSAT.v
+++ b/theories/examples/approximate_samplers/approx_walkSAT.v
@@ -1,6 +1,7 @@
 (** * Termination with probability 1 of the WalkSAT algorithm *)
 From clutch.ub_logic Require Export ub_clutch ub_rules.
 From clutch Require Export examples.approximate_samplers.approx_sampler_lib.
+From clutch Require Export examples.approximate_samplers.approx_higherorder_incremental.
 From Coquelicot Require Import Series.
 Require Import Lra.
 
@@ -27,8 +28,6 @@ Section higherorder_walkSAT.
     | inv_cons (b : bool) (m' : list bool) (asn' : val) : (inv_asn' m' asn') -> inv_asn' (b :: m') (SOMEV (#b, asn')).
   Definition inv_asn m asn : Prop := length m = N /\ inv_asn' m asn.
 
-  (* Placeholder stuck expression *)
-  Definition error : expr := (#42 #42)%E.
 
   (* Set up a random assignment of n boolean variables *)
   Definition mk_init_asn': val :=
@@ -77,7 +76,7 @@ Section higherorder_walkSAT.
   (* eval_asn spec *)
   Definition wp_eval_asn m asn E (n : nat) :
     (⊢ ⌜ (n < (length m))%nat ⌝ -∗ ⌜ inv_asn' m asn ⌝ -∗
-       WP (eval_asn asn #n)%E @ E {{ fun v => ⌜#(m !!! n : bool) = v⌝}})%I .
+       WP (eval_asn asn #n)%E @ E [{ fun v => ⌜#(m !!! n : bool) = v⌝}])%I .
   Proof using N HN.
     iIntros "%Hn %Hinv".
     iRevert (Hn).
@@ -99,7 +98,7 @@ Section higherorder_walkSAT.
           destruct Hc; try lia.
           by rewrite -H0 /= Nat2Z.inj_0 in H. }
         destruct n' as [|n''] eqn:Hn'; [by rewrite Nat2Z.inj_0 in H |].
-        wp_apply (ub_wp_wand with "[IH]").
+        wp_apply (ub_twp_wand with "[IH]").
         { iApply "IH".
           iPureIntro.
           rewrite cons_length in Hlen.
@@ -128,7 +127,7 @@ Section higherorder_walkSAT.
   (* update_asn spec *)
   Definition wp_update_asn m asn E n (b: bool) :
     (⊢ ⌜(n < length m)%nat ⌝ -∗ ⌜inv_asn' m asn ⌝ -∗
-        WP (update_asn asn #n #b)%E @ E {{ fun asn' => ⌜inv_asn' (<[n := b]> m) asn' ⌝}})%I.
+        WP (update_asn asn #n #b)%E @ E [{ fun asn' => ⌜inv_asn' (<[n := b]> m) asn' ⌝}])%I.
   Proof using N HN.
     iIntros "%Hn %Hinv".
     iRevert (Hn).
@@ -152,7 +151,7 @@ Section higherorder_walkSAT.
           pose Hc := Nat.le_0_l; apply (Nat.lt_eq_cases 0%nat n') in Hc.
           destruct Hc; try lia.
           by rewrite -H0 /= Nat2Z.inj_0 in H. }
-        wp_apply (ub_wp_wand with "[IH]").
+        wp_apply (ub_twp_wand with "[IH]").
         { iApply "IH".
           iPureIntro.
           rewrite cons_length in Hlen.
@@ -395,14 +394,14 @@ Section higherorder_walkSAT.
   (* evaluate_fvar spec*)
   Lemma wp_evaluate_fvar l asn m v E :
     (⊢ ⌜ inv_asn m asn ⌝ -∗ l ↦ asn -∗
-       WP (evaluate_fvar v) asn @ E {{ fun v' => l ↦ asn ∗ ⌜v' = #(fvar_SAT m v)⌝ }} )%I.
+       WP (evaluate_fvar v) asn @ E [{ fun v' => l ↦ asn ∗ ⌜v' = #(fvar_SAT m v)⌝ }] )%I.
   Proof.
     iIntros "%Hinv Hl".
     destruct v as [p v vwf].
     rewrite /evaluate_fvar.
     wp_pures.
     wp_bind (eval_asn _ _)%E.
-    wp_apply (ub_wp_wand with "[]").
+    wp_apply (ub_twp_wand with "[]").
     { iApply wp_eval_asn; iPureIntro; last first.
       - rewrite /inv_asn in Hinv. by destruct Hinv.
       - destruct Hinv; lia. }
@@ -420,25 +419,25 @@ Section higherorder_walkSAT.
   (* evaluate_clause spec *)
   Lemma wp_evaluate_clause l asn m c E :
     (⊢ ⌜ inv_asn m asn ⌝ -∗ l ↦ asn -∗
-     WP (evaluate_clause c) asn @ E {{ fun v => l ↦ asn ∗ ⌜v = #(clause_SAT m c)⌝ }} )%I.
+     WP (evaluate_clause c) asn @ E [{ fun v => l ↦ asn ∗ ⌜v = #(clause_SAT m c)⌝ }])%I.
   Proof.
     iIntros "%Hinv Hl".
     destruct c as [e1 e2 e3].
     rewrite /evaluate_clause.
     wp_pures.
     wp_bind (evaluate_fvar _ _).
-    wp_apply (ub_wp_wand with "[Hl]").
+    wp_apply (ub_twp_wand with "[Hl]").
     { iApply wp_evaluate_fvar; [eauto|iFrame]. }
     iIntros (s1) "(Hl&%Hs1)".
     destruct (fvar_SAT m e1) as [|] eqn:HeqS1; rewrite Hs1; wp_pures.
     { iModIntro; iFrame; iPureIntro; f_equal. simpl; by rewrite HeqS1. }
     wp_bind (evaluate_fvar _ _).
-    wp_apply (ub_wp_wand with "[Hl]").
+    wp_apply (ub_twp_wand with "[Hl]").
     { iApply wp_evaluate_fvar; [eauto|iFrame]. }
     iIntros (s2) "(Hl&%Hs2)".
     destruct (fvar_SAT m e2) as [|] eqn:HeqS2; rewrite Hs2; wp_pures.
     { iModIntro; iFrame; iPureIntro; f_equal. simpl; by rewrite HeqS2 orb_true_r. }
-    wp_apply (ub_wp_wand with "[Hl]").
+    wp_apply (ub_twp_wand with "[Hl]").
     { iApply wp_evaluate_fvar; [eauto|iFrame]. }
     iIntros (s3) "(Hl&%Hs3)".
     destruct (fvar_SAT m e3) as [|] eqn:HeqS3; rewrite Hs3.
@@ -490,7 +489,7 @@ Section higherorder_walkSAT.
 
   (* spec for helper lemma *)
   Lemma wp_clause_to_index (c: clause) (target : fin 3) E :
-    ⊢ (WP (clause_to_index c #target)%E @ E {{ fun i => ⌜ i = #(fVar_index (proj_clause_value c target))⌝ }})%I.
+    ⊢ (WP (clause_to_index c #target)%E @ E [{ fun i => ⌜ i = #(fVar_index (proj_clause_value c target))⌝ }])%I.
   Proof.
     iStartProof. rewrite /proj_clause_value/clause_to_index /proj_clause_value /fVar_index.
     destruct c.
@@ -577,25 +576,25 @@ Section higherorder_walkSAT.
     ⊢ (l ↦ asn -∗
        € (εInv ε0 (S p)) -∗
        WP (resample_clause c #l)%E @ E
-         {{ fun _ =>
+         [{ fun _ =>
               ∃ asn' m', (l ↦ asn') ∗
                          ⌜inv_asn m' asn' ⌝ ∗
                          ( (* Flips the target variable and loses some credit, or... *)
                            ( € (εInv ε0 p) ∗
                             ⌜m' = (<[(fVar_index (proj_clause_value c target)) := negb (m !!! (fVar_index (proj_clause_value c target)))]> m)⌝) ∨
                             (* ...obtains the amplified credit *)
-                            (€ (εAmplified ε0)))}})%I.
+                            (€ (εAmplified ε0)))}])%I.
   Proof.
     iIntros (Hinv) "Hl Hε".
     Opaque update_asn.
     rewrite /resample_clause.
     wp_pures.
-    wp_apply (wp_load with "Hl").
+    wp_apply (twp_load with "Hl").
     iIntros "Hl".
     wp_pures.
     wp_bind (rand _)%E.
     replace (length m) with N in Hp; [|by destruct Hinv].
-    wp_apply (wp_couple_rand_adv_comp1 _ _ _ _ (εDistr_resampler _ _ _ target) with "Hε").
+    wp_apply (twp_couple_rand_adv_comp1 _ _ _ _ (εDistr_resampler _ _ _ target) with "Hε").
     {
       rewrite εDistr_mean.
       rewrite /εInv.
@@ -606,11 +605,11 @@ Section higherorder_walkSAT.
 
     - (* sampler chooses the target index and flips it *)
       wp_bind (clause_to_index c _)%E.
-      wp_apply (ub_wp_wand); first iApply (wp_clause_to_index c i).
+      wp_apply (ub_twp_wand); first iApply (wp_clause_to_index c i).
       iIntros (i') "->".
       wp_pures.
       wp_bind (eval_asn _ _)%E.
-      wp_apply (ub_wp_wand with "[]").
+      wp_apply (ub_twp_wand with "[]").
       { iApply wp_eval_asn; iPureIntro; last first.
         - rewrite /inv_asn in Hinv. by destruct Hinv.
         - destruct (proj_clause_value c i) as [? ? ?].
@@ -619,7 +618,7 @@ Section higherorder_walkSAT.
       iIntros (v) "<-".
       wp_pures.
       wp_bind (update_asn _ _ _).
-      wp_apply (ub_wp_wand with "[]").
+      wp_apply (ub_twp_wand with "[]").
       { iApply wp_update_asn; iPureIntro; last first.
         - rewrite /inv_asn in Hinv. by destruct Hinv.
         - destruct (proj_clause_value c i) as [? ? ?].
@@ -627,7 +626,6 @@ Section higherorder_walkSAT.
           simpl; lia. }
       iIntros (v) "%Hinv'".
       wp_pures.
-
       wp_store.
       iModIntro.
       iExists _, _.
@@ -657,11 +655,11 @@ Section higherorder_walkSAT.
       done.
     - (* sampler chooses the wrong index, step through and conclude by the amplification  *)
       wp_bind (clause_to_index c _)%E.
-      wp_apply (ub_wp_wand); first iApply (wp_clause_to_index c i).
+      wp_apply (ub_twp_wand); first iApply (wp_clause_to_index c i).
       iIntros (i') "->".
       wp_pures.
       wp_bind (eval_asn _ _)%E.
-      wp_apply (ub_wp_wand with "[]").
+      wp_apply (ub_twp_wand with "[]").
       { iApply wp_eval_asn; iPureIntro; last first.
         - rewrite /inv_asn in Hinv. by destruct Hinv.
         - destruct (proj_clause_value c i) as [? ? ?].
@@ -670,7 +668,7 @@ Section higherorder_walkSAT.
       iIntros (v) "<-".
       wp_pures.
       wp_bind (update_asn _ _ _).
-      wp_apply (ub_wp_wand with "[]").
+      wp_apply (ub_twp_wand with "[]").
       { iApply wp_update_asn; iPureIntro; last first.
         - rewrite /inv_asn in Hinv. by destruct Hinv.
         - destruct (proj_clause_value c i) as [? ? ?].
@@ -703,7 +701,7 @@ Section higherorder_walkSAT.
 
   (* running the checker *)
   Lemma wp_check l asn m f E :
-    (⊢ ⌜ inv_asn m asn ⌝ -∗ l ↦ asn -∗ ((WP ((Val (checker f)) #l) @ E {{ λ v', l ↦ asn ∗ ⌜v' = #(formula_SAT m f)⌝ }})))%I.
+    (⊢ ⌜ inv_asn m asn ⌝ -∗ l ↦ asn -∗ ((WP ((Val (checker f)) #l) @ E [{ λ v', l ↦ asn ∗ ⌜v' = #(formula_SAT m f)⌝ }])))%I.
   Proof.
     iInduction f as [|c f'] "IH".
     - iIntros "%Hinv Hl".
@@ -715,12 +713,12 @@ Section higherorder_walkSAT.
       wp_bind (! _)%E.
       wp_load.
       wp_bind (evaluate_clause _ _)%E.
-      wp_apply (ub_wp_wand with "[Hl]").
+      wp_apply (ub_twp_wand with "[Hl]").
       { wp_apply wp_evaluate_clause; [|iFrame]. iPureIntro. eapply Hinv.  }
       iIntros (ev) "(Hl&->)".
       destruct (clause_SAT m c) as [|] eqn:Hcsat.
       + wp_pure.
-        wp_apply (ub_wp_wand with "[Hl]").
+        wp_apply (ub_twp_wand with "[Hl]").
         { iApply "IH"; [eauto|iFrame]. }
         iIntros (v) "(Hl&%Hf')".
         iFrame; iPureIntro.
@@ -738,7 +736,7 @@ Section higherorder_walkSAT.
     (⊢ ⌜formula_SAT m f = true ⌝ -∗
        ⌜ inv_asn m asn ⌝ -∗
        l ↦ asn -∗
-       (WP ((Val (sampler f)) #l) @ E {{ λ v', l ↦ asn }}))%I.
+       (WP ((Val (sampler f)) #l) @ E [{ λ v', l ↦ asn }]))%I.
   Proof.
     iInduction f as [|c f'] "IHf".
     - iIntros "? ? ?".
@@ -752,7 +750,7 @@ Section higherorder_walkSAT.
       wp_load.
       wp_pures.
       wp_bind (evaluate_clause _ _)%E.
-      wp_apply (ub_wp_wand with "[Hl]").
+      wp_apply (ub_twp_wand with "[Hl]").
       { wp_apply wp_evaluate_clause; [|iFrame].
         iPureIntro. eapply Hinv.  }
       iIntros (v) "(Hl&->)".
@@ -780,9 +778,9 @@ Section higherorder_walkSAT.
        l ↦ asn -∗
        € (εInv ε (S p)) -∗
        (WP ((Val (sampler f)) #l) @ E
-          {{ λ v', ∃ asn' m', l ↦ asn' ∗ ⌜ inv_asn m' asn' ⌝ ∗
+          [{ λ v', ∃ asn' m', l ↦ asn' ∗ ⌜ inv_asn m' asn' ⌝ ∗
                       ((⌜(progress_measure f m' solution < progress_measure f m solution)%nat ⌝ ∗ €(εInv ε p)) ∨
-                       (€ (εAmplified ε)) )}}))%I.
+                       (€ (εAmplified ε)) )}]))%I.
     Proof.
       iIntros "%Hp %Hsol_len %Hsol %Hm %Hinv Hl Hε".
       destruct (find_progress _ _ Hm) as [f1 [f2 [c (-> & Hf1 & Hc)]]].
@@ -794,14 +792,14 @@ Section higherorder_walkSAT.
         wp_pures.
         wp_load.
         wp_bind (evaluate_clause _ _)%E.
-        wp_apply (ub_wp_wand with "[Hl]").
+        wp_apply (ub_twp_wand with "[Hl]").
         { wp_apply (wp_evaluate_clause with "[] Hl").
           iPureIntro; eauto. }
         iIntros (r) "(Hl&->)".
         rewrite Hc'.
         wp_pure.
         replace (f1' ++ [c] ++ f2) with (f1' ++ c :: f2) by auto.
-        wp_apply (ub_wp_wand with "[Hl Hε]").
+        wp_apply (ub_twp_wand with "[Hl Hε]").
         { iApply ("IH" with "[] [] [] Hl Hε").
           - iPureIntro. rewrite -Hm /formula_SAT /= /fmap. f_equal. auto.
           - iPureIntro. rewrite -Hsol /formula_SAT /= /fmap. f_equal.
@@ -822,7 +820,7 @@ Section higherorder_walkSAT.
       wp_pures.
       wp_load.
       wp_bind (evaluate_clause _ _)%E.
-      wp_apply (ub_wp_wand with "[Hl]").
+      wp_apply (ub_twp_wand with "[Hl]").
       { wp_apply (wp_evaluate_clause with "[] Hl"). iPureIntro; eapply Hinv. }
       iIntros (r) "(Hl&->)".
       rewrite Hc; wp_pures.
@@ -834,7 +832,7 @@ Section higherorder_walkSAT.
       apply andb_prop in Hsol; destruct Hsol as [Hc_solution Hf2_solution].
       destruct (progress_is_possible_clause _ _ _ Hc_solution Hc) as [targetFV [HtargetClause HtargetFV]].
       destruct (reflect_progress_to_target targetFV _ HtargetClause) as [target Htarget].
-      wp_apply (ub_wp_wand with "[Hε Hl]").
+      wp_apply (ub_twp_wand with "[Hε Hl]").
       { wp_apply ((resample_amplify _ target) with "Hl Hε"); last first.
         - eapply Hinv.
         - destruct Hinv; lia. }
@@ -843,7 +841,6 @@ Section higherorder_walkSAT.
       iFrame.
       iSplit; [iPureIntro; eapply Hasn'|].
       iDestruct "Hs" as "[[Hε %H]|Hε]".
-
       - (* Flip is lucky and makes progress *)
         iLeft; iFrame.
         iPureIntro.
@@ -897,14 +894,15 @@ Section higherorder_walkSAT.
           (εProgress ε)
           N
           HiL
-          E)%I.
+          E
+          (fun _ => ⌜True⌝ ))%I.
   Proof.
     iIntros "%HHil %Hsolution %Hsl %Hnontrivial".
     rewrite /incr_sampling_scheme_spec.
     iSplit.
     - iIntros "[Hcr | [%asn [%m (Hl & Hcr & %Hinv & %Hp)]]]".
       + (* € 1 case: spend *)
-        iApply (wp_ec_spend with "Hcr"); [|auto].
+        iApply (twp_ec_spend with "Hcr"); [|auto].
         apply final_progress.
       + (* Ψ 0 case *)
         apply Nat.le_0_r in Hp.
@@ -912,35 +910,37 @@ Section higherorder_walkSAT.
         simplify_eq.
         (* using Ψ, asn now equals the solution. step the sampler... *)
         wp_pures.
-        wp_apply (ub_wp_wand with "[Hl]").
+        wp_apply (ub_twp_wand with "[Hl]").
         { wp_apply wp_sampler_done; iFrame; iPureIntro; eauto. }
         iIntros (v) "Hl".
         (* then step the checker... *)
         wp_pures.
-        wp_apply (ub_wp_wand with "[Hl]").
+        wp_apply (ub_twp_wand with "[Hl]").
         { iApply wp_check; [|iFrame].
           iPureIntro; apply Hinv. }
-        iIntros (r) "(Hl&->)"; iPureIntro; do 2 f_equal.
-        simplify_eq; done.
+        iIntros (r) "(Hl&->)"; iPureIntro.
+        split; eauto.
+        do 2 f_equal. done.
     - iIntros (i p) "!> (%Hp_bound & %Hi_bound & Hε & [%asn [%m (Hl & Hcr & %Hinv & %Hp)]])".
       wp_pures.
       (* step the sampler differently depending on if it is SAT or not *)
       destruct (formula_SAT m f) as [|] eqn:Hsat.
       + (* SAT: we can't make progress or amplify, but that is be okay, since we can pass the check *)
-        wp_apply (ub_wp_wand with "[Hl]").
+        wp_apply (ub_twp_wand with "[Hl]").
         { wp_apply wp_sampler_done; try by iPureIntro. iFrame. }
         iIntros (?) "Hl".
         iLeft.
         wp_pures.
-        iApply (ub_wp_wand with "[Hε Hcr Hl]").
+        iApply (ub_twp_wand with "[Hε Hcr Hl]").
         { iApply wp_check; iFrame. iPureIntro. eapply Hinv. }
         iIntros (?) "[? ->]"; iPureIntro.
+        split; eauto.
         by do 2 f_equal.
       + (* UNSAT *)
         (* Step to the resampling step, and amplify *)
         rewrite /sampler.
         wp_pures.
-        wp_apply (ub_wp_wand with "[Hl Hcr]").
+        wp_apply (ub_twp_wand with "[Hl Hcr]").
         { wp_apply (wp_sampler_amplify with "[] [] [] [] [] Hl [Hcr]"); last iFrame; try eauto. }
         iIntros (s) "[%asn' [%m' (Hl & %Hinv' & [(%Hp' & A)|Hamp])]]".
         * (* makes progress *)
@@ -950,19 +950,18 @@ Section higherorder_walkSAT.
           { iExists _, _. iFrame. iSplit; iPureIntro; eauto. lia. }
           iIntros "[%asn'' [%m'' (Hl & Hcr & %Hasn'' & %Ap'')]]".
           wp_pures.
-          iApply (ub_wp_wand with "[Hl]").
+          iApply (ub_twp_wand with "[Hl]").
           { iApply wp_check; iFrame. iPureIntro. eauto. }
           iIntros (?) "(Hl & ->)".
           iSplitL.
           {  iFrame. iExists _, _. iFrame. iSplit; iPureIntro; eauto. }
-          iExists _; iPureIntro; eauto.
+          destruct (formula_SAT m'' f); iPureIntro; [right|left]; eauto.
         * (* amplifies *)
           iRight; iRight.
           iFrame.
           (* Revert back to iProgress N *)
           iExists N.
           iSplitR; eauto.
-
           (* Transfer the amplfied credits between the invariants *)
           iAssert (€ (εInv ε N) ∗ € (pos_to_nn (εExcess ε)) )%I with "[Hamp]" as "[Hinv Hexcess]".
           { iApply ec_split.
@@ -978,7 +977,6 @@ Section higherorder_walkSAT.
             do 2 (rewrite Rmax_right; [|apply HHil; lia]).
             rewrite S_INR. lra.
           }
-
           iFrame.
           iSplitL.
           -- iExists _, _; iFrame; iSplit; iPureIntro; eauto.
@@ -995,12 +993,12 @@ Section higherorder_walkSAT.
                   apply IH. }
           -- iIntros "[%asn'' [%m'' (Hl & ? & % & %)]]".
              wp_pures.
-             wp_apply (ub_wp_wand with "[Hl]").
+             wp_apply (ub_twp_wand with "[Hl]").
              { iApply wp_check; [|iFrame]. iPureIntro; eauto. }
              iIntros (?) "[? ->]".
              iFrame.
              iSplitL.
              { iExists _, _. iFrame.  iSplit; eauto. }
-             eauto.
+             destruct (formula_SAT m'' f); iPureIntro; [right|left]; eauto.
     Qed.
 End higherorder_walkSAT.