diff --git a/CHANGELOG.md b/CHANGELOG.md
index 0f4d979450f8b2e1c8e21994903d77596fe69853..518d58d5d46c76f3f6e8517934b74c3fb3f0c648 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -27,6 +27,14 @@ lemma.
 * Make the `inG` instances for `libG` fields local, so they are only used inside
   the library that defines the `libG`.
 
+**Changes in `iris_heap_lang`:**
+
+* Changed the `num_laters_per_step` of `heap_lang` to `λ n, n`, signifying that
+  each step of the weakest precondition strips `n` laters, where `n` is the
+  number of steps taken so far. This number is tied to ghost state in the state
+  interpretation, which is exposed, updated, and used with new lemmas
+  `wp_lb_init`, `wp_lb_update`, and `wp_step_fupdN_lb`. (by Jonas Kastberg Hinrichsen)
+
 ## Iris 3.6.0 (2022-01-22)
 
 The highlights and most notable changes of this release are:
diff --git a/iris/program_logic/total_adequacy.v b/iris/program_logic/total_adequacy.v
index a8d0f5f09f6ae1a2c011c54bf6072c58b40dbee6..0e6dd10e749d30247a128ab429b7bb9c442ff6a8 100644
--- a/iris/program_logic/total_adequacy.v
+++ b/iris/program_logic/total_adequacy.v
@@ -116,22 +116,27 @@ Proof.
 Qed.
 End adequacy.
 
-Theorem twp_total Σ Λ `{!invGpreS Σ} s e σ Φ :
+Theorem twp_total Σ Λ `{!invGpreS Σ} s e σ Φ n :
   (∀ `{Hinv : !invGS Σ},
      ⊢ |={⊤}=> ∃
-         (stateI : state Λ → list (observation Λ) → nat → iProp Σ)
-         (fork_post : val Λ → iProp Σ),
+         (stateI : state Λ → nat → list (observation Λ) → nat → iProp Σ)
+         (** We abstract over any instance of [irisG], and thus any value of
+             the field [num_laters_per_step]. This is needed because instances
+             of [irisG] (e.g., the one of HeapLang) are shared between WP and
+             TWP, where TWP simply ignores [num_laters_per_step]. *)
+         (num_laters_per_step : nat → nat)
+         (fork_post : val Λ → iProp Σ)
+         state_interp_mono,
        let _ : irisGS Λ Σ :=
-           IrisG _ _ Hinv (λ σ _, stateI σ) fork_post (λ _, 0)
-                 (λ _ _ _ _, fupd_intro _ _)
+           IrisG _ _ Hinv stateI fork_post num_laters_per_step state_interp_mono
        in
-       stateI σ [] 0 ∗ WP e @ s; ⊤ [{ Φ }]) →
+       stateI σ n [] 0 ∗ WP e @ s; ⊤ [{ Φ }]) →
   sn erased_step ([e], σ). (* i.e. ([e], σ) is strongly normalizing *)
 Proof.
   intros Hwp. apply (soundness (M:=iResUR Σ) _  1); simpl.
   apply (fupd_plain_soundness ⊤ ⊤ _)=> Hinv.
-  iMod (Hwp) as (stateI fork_post) "[Hσ H]".
-  iApply (@twptp_total _ _ (IrisG _ _ Hinv (λ σ _, stateI σ) fork_post _ _)
-                       _ 0 with "Hσ").
+  iMod (Hwp) as (stateI num_laters_per_step fork_post stateI_mono) "[Hσ H]".
+  set (iG := IrisG _ _ Hinv stateI fork_post num_laters_per_step stateI_mono).
+  iApply (@twptp_total _ _ iG _ n with "Hσ").
   by iApply (@twp_twptp _ _ (IrisG _ _ Hinv _ fork_post _ _)).
 Qed.
diff --git a/iris_heap_lang/adequacy.v b/iris_heap_lang/adequacy.v
index 42814e06e6ed009146c21a9cd4faf7187639c751..9d4ad1e7ff8b381d819fcff35361d453167d142e 100644
--- a/iris_heap_lang/adequacy.v
+++ b/iris_heap_lang/adequacy.v
@@ -1,4 +1,5 @@
 From iris.algebra Require Import auth.
+From iris.base_logic.lib Require Import mono_nat.
 From iris.proofmode Require Import proofmode.
 From iris.program_logic Require Export weakestpre adequacy.
 From iris.heap_lang Require Import proofmode notation.
@@ -9,23 +10,42 @@ Class heapGpreS Σ := HeapGpreS {
   heapGpreS_heap :> gen_heapGpreS loc (option val) Σ;
   heapGpreS_inv_heap :> inv_heapGpreS loc (option val) Σ;
   heapGpreS_proph :> proph_mapGpreS proph_id (val * val) Σ;
+  heapGS_step_cnt :> mono_natG Σ;
 }.
 
 Definition heapΣ : gFunctors :=
-  #[invΣ; gen_heapΣ loc (option val); inv_heapΣ loc (option val); proph_mapΣ proph_id (val * val)].
+  #[invΣ; gen_heapΣ loc (option val); inv_heapΣ loc (option val);
+    proph_mapΣ proph_id (val * val); mono_natΣ].
 Global Instance subG_heapGpreS {Σ} : subG heapΣ Σ → heapGpreS Σ.
 Proof. solve_inG. Qed.
 
+(* TODO: The [wp_adequacy] lemma is insufficient for a state interpretation
+with a non-constant step index function. We thus use the more general
+[wp_strong_adequacy] lemma. The proof below replicates part of the proof of
+[wp_adequacy], and it hence would make sense to see if we can prove a version
+of [wp_adequacy] for a non-constant step version. *)
 Definition heap_adequacy Σ `{!heapGpreS Σ} s e σ φ :
   (∀ `{!heapGS Σ}, ⊢ inv_heap_inv -∗ WP e @ s; ⊤ {{ v, ⌜φ v⌝ }}) →
   adequate s e σ (λ v _, φ v).
 Proof.
-  intros Hwp; eapply (wp_adequacy _ _); iIntros (? κs) "".
+  intros Hwp.
+  apply adequate_alt; intros t2 σ2 [n [κs ?]]%erased_steps_nsteps.
+  eapply (wp_strong_adequacy Σ _); [|done].
+  iIntros (Hinv).
   iMod (gen_heap_init σ.(heap)) as (?) "[Hh _]".
   iMod (inv_heap_init loc (option val)) as (?) ">Hi".
   iMod (proph_map_init κs σ.(used_proph_id)) as (?) "Hp".
-  iModIntro. iExists
-    (λ σ κs, (gen_heap_interp σ.(heap) ∗ proph_map_interp κs σ.(used_proph_id))%I),
-    (λ _, True%I).
-  iFrame. iApply (Hwp (HeapGS _ _ _ _ _)). done.
+  iMod (mono_nat_own_alloc) as (γ) "[Hsteps _]".
+  iDestruct (Hwp (HeapGS _ _ _ _ _ _ _) with "Hi") as "Hwp".
+  iModIntro. iExists s.
+  iExists (λ σ ns κs nt, (gen_heap_interp σ.(heap) ∗
+                          proph_map_interp κs σ.(used_proph_id) ∗
+                          mono_nat_auth_own γ 1 ns))%I.
+  iExists [(λ v, ⌜φ v⌝%I)], (λ _, True)%I, _ => /=.
+  iFrame.
+  iIntros (es' t2' -> ? ?) " _ H _".
+  iApply fupd_mask_intro_discard; [done|]. iSplit; [|done].
+  iDestruct (big_sepL2_cons_inv_r with "H") as (e' ? ->) "[Hwp H]".
+  iDestruct (big_sepL2_nil_inv_r with "H") as %->.
+  iIntros (v2 t2'' [= -> <-]). by rewrite to_of_val.
 Qed.
diff --git a/iris_heap_lang/primitive_laws.v b/iris_heap_lang/primitive_laws.v
index 0af58af239b8cc83a7b465365026f62d84e440fb..e51974e0260253e43b896516f27d21a55571d937 100644
--- a/iris_heap_lang/primitive_laws.v
+++ b/iris_heap_lang/primitive_laws.v
@@ -3,6 +3,7 @@ the Iris lifting lemmas. *)
 
 From iris.proofmode Require Import proofmode.
 From iris.bi.lib Require Import fractional.
+From iris.base_logic.lib Require Import mono_nat.
 From iris.base_logic.lib Require Export gen_heap proph_map gen_inv_heap.
 From iris.program_logic Require Export weakestpre total_weakestpre.
 From iris.program_logic Require Import ectx_lifting total_ectx_lifting.
@@ -15,16 +16,60 @@ Class heapGS Σ := HeapGS {
   heapGS_gen_heapGS :> gen_heapGS loc (option val) Σ;
   heapGS_inv_heapGS :> inv_heapGS loc (option val) Σ;
   heapGS_proph_mapGS :> proph_mapGS proph_id (val * val) Σ;
+  heapGS_step_name : gname;
+  heapGS_step_cnt : mono_natG Σ;
 }.
+Local Existing Instance heapGS_step_cnt.
 
-Global Instance heapGS_irisGS `{!heapGS Σ} : irisGS heap_lang Σ := {
+Section steps.
+  Context `{!heapGS Σ}.
+
+  Local Definition steps_auth (n : nat) : iProp Σ :=
+    mono_nat_auth_own heapGS_step_name 1 n.
+
+  Definition steps_lb (n : nat) : iProp Σ :=
+    mono_nat_lb_own heapGS_step_name n.
+
+  Local Lemma steps_lb_valid n m :
+    steps_auth n -∗ steps_lb m -∗ ⌜m ≤ n⌝.
+  Proof.
+    iIntros "Hauth Hlb".
+    by iDestruct (mono_nat_lb_own_valid with "Hauth Hlb") as %[_ Hle].
+  Qed.
+
+  Local Lemma steps_lb_get n :
+    steps_auth n -∗ steps_lb n.
+  Proof. apply mono_nat_lb_own_get. Qed.
+
+  Local Lemma steps_auth_update n n' :
+    n ≤ n' → steps_auth n ==∗ steps_auth n' ∗ steps_lb n'.
+  Proof. intros Hle. by apply mono_nat_own_update. Qed.
+
+  Local Lemma steps_auth_update_S n :
+    steps_auth n ==∗ steps_auth (S n).
+  Proof.
+    iIntros "Hauth".
+    iMod (mono_nat_own_update with "Hauth") as "[$ _]"; [lia|done].
+  Qed.
+
+  Lemma steps_lb_le n n' :
+    n' ≤ n → steps_lb n -∗ steps_lb n'.
+  Proof. intros Hle. by apply mono_nat_lb_own_le. Qed.
+
+End steps.
+
+Global Program Instance heapGS_irisGS `{heapGS Σ} : irisGS heap_lang Σ := {
   iris_invGS := heapGS_invGS;
-  state_interp σ _ κs _ :=
-    (gen_heap_interp σ.(heap) ∗ proph_map_interp κs σ.(used_proph_id))%I;
+  state_interp σ step_cnt κs _ :=
+    (gen_heap_interp σ.(heap) ∗ proph_map_interp κs σ.(used_proph_id) ∗
+     steps_auth step_cnt)%I;
   fork_post _ := True%I;
-  num_laters_per_step _ := 0;
-  state_interp_mono _ _ _ _ := fupd_intro _ _
+  num_laters_per_step n := n;
 }.
+Next Obligation.
+  iIntros (?? σ ns κs nt)  "/= ($ & $ & H)".
+  by iMod (steps_auth_update_S with "H") as "$".
+Qed.
 
 (** Since we use an [option val] instance of [gen_heap], we need to overwrite
 the notations.  That also helps for scopes and coercions. *)
@@ -68,6 +113,61 @@ Implicit Types σ : state.
 Implicit Types v : val.
 Implicit Types l : loc.
 
+Lemma wp_lb_init s E e Φ :
+  TCEq (to_val e) None →
+  (steps_lb 0 -∗ WP e @ s; E {{ v, Φ v }}) -∗
+  WP e @ s; E {{ Φ }}.
+Proof.
+  (** TODO: We should try to use a generic lifting lemma (and avoid [wp_unfold])
+  here, since this breaks the WP abstraction. *)
+  rewrite !wp_unfold /wp_pre /=. iIntros (->) "Hwp".
+  iIntros (σ1 ns κ κs m) "(Hσ & Hκ & Hsteps)".
+  iDestruct (steps_lb_get with "Hsteps") as "#Hlb".
+  iDestruct (steps_lb_le _ 0 with "Hlb") as "Hlb0"; [lia|].
+  iSpecialize ("Hwp" with "Hlb0"). iApply ("Hwp" $! σ1 ns κ κs m). iFrame.
+Qed.
+
+Lemma wp_lb_update s n E e Φ :
+  TCEq (to_val e) None →
+  steps_lb n -∗
+  WP e @ s; E {{ v, steps_lb (S n) -∗ Φ v }} -∗
+  WP e @ s; E {{ Φ }}.
+Proof.
+  (** TODO: We should try to use a generic lifting lemma (and avoid [wp_unfold])
+  here, since this breaks the WP abstraction. *)
+  rewrite !wp_unfold /wp_pre /=. iIntros (->) "Hlb Hwp".
+  iIntros (σ1 ns κ κs m) "(Hσ & Hκ & Hsteps)".
+  iDestruct (steps_lb_valid with "Hsteps Hlb") as %?.
+  iMod ("Hwp" $! σ1 ns κ κs m with "[$Hσ $Hκ $Hsteps]") as "[%Hs Hwp]".
+  iModIntro. iSplit; [done|].
+  iIntros (e2 σ2 efs Hstep).
+  iMod ("Hwp" with "[//]") as "Hwp".
+  iIntros "!> !>". iMod "Hwp" as "Hwp". iIntros "!>".
+  iApply (step_fupdN_wand with "Hwp").
+  iIntros "Hwp". iMod "Hwp" as "(($ & $ & Hsteps)& Hwp & $)".
+  iDestruct (steps_lb_get with "Hsteps") as "#HlbS".
+  iDestruct (steps_lb_le _ (S n) with "HlbS") as "#HlbS'"; [lia|].
+  iModIntro. iFrame "Hsteps".
+  iApply (wp_wand with "Hwp"). iIntros (v) "HΦ". by iApply "HΦ".
+Qed.
+
+Lemma wp_step_fupdN_lb s n E1 E2 e P Φ :
+  TCEq (to_val e) None →
+  E2 ⊆ E1 →
+  steps_lb n -∗
+  (|={E1∖E2,∅}=> |={∅}▷=>^(S n) |={∅,E1∖E2}=> P) -∗
+  WP e @ s; E2 {{ v, P ={E1}=∗ Φ v }} -∗
+  WP e @ s; E1 {{ Φ }}.
+Proof.
+  iIntros (He HE) "Hlb HP Hwp".
+  iApply wp_step_fupdN; [done|].
+  iSplit; [|by iFrame].
+  iIntros (σ ns κs nt) "(? & ? & Hsteps)".
+  iDestruct (steps_lb_valid with "Hsteps Hlb") as %Hle.
+  iApply fupd_mask_intro; [set_solver|].
+  iIntros "_". iPureIntro. rewrite /num_laters_per_step /=. lia.
+Qed.
+
 (** Recursive functions: we do not use this lemmas as it is easier to use Löb
 induction directly, but this demonstrates that we can state the expected
 reasoning principle for recursive functions, without any visible ▷. *)
@@ -87,16 +187,20 @@ Lemma wp_fork s E e Φ :
   ▷ WP e @ s; ⊤ {{ _, True }} -∗ ▷ Φ (LitV LitUnit) -∗ WP Fork e @ s; E {{ Φ }}.
 Proof.
   iIntros "He HΦ". iApply wp_lift_atomic_head_step; [done|].
-  iIntros (σ1 ns κ κs nt) "Hσ !>"; iSplit; first by eauto with head_step.
-  iIntros "!>" (v2 σ2 efs Hstep); inv_head_step. by iFrame.
+  iIntros (σ1 ns κ κs nt) "(?&?&Hsteps) !>"; iSplit; first by eauto with head_step.
+  iIntros "!>" (v2 σ2 efs Hstep); inv_head_step.
+  iMod (steps_auth_update_S with "Hsteps") as "Hsteps".
+  by iFrame.
 Qed.
 
 Lemma twp_fork s E e Φ :
   WP e @ s; ⊤ [{ _, True }] -∗ Φ (LitV LitUnit) -∗ WP Fork e @ s; E [{ Φ }].
 Proof.
   iIntros "He HΦ". iApply twp_lift_atomic_head_step; [done|].
-  iIntros (σ1 ns κs nt) "Hσ !>"; iSplit; first by eauto with head_step.
-  iIntros (κ v2 σ2 efs Hstep); inv_head_step. by iFrame.
+  iIntros (σ1 ns κs nt) "(?&?&Hsteps) !>"; iSplit; first by eauto with head_step.
+  iIntros (κ v2 σ2 efs Hstep); inv_head_step.
+  iMod (steps_auth_update_S with "Hsteps") as "Hsteps".
+  by iFrame.
 Qed.
 
 (** Heap *)
@@ -224,13 +328,14 @@ Lemma twp_allocN_seq s E v n :
       (l +ₗ (i : nat)) ↦ v ∗ meta_token (l +ₗ (i : nat)) ⊤ }]].
 Proof.
   iIntros (Hn Φ) "_ HΦ". iApply twp_lift_atomic_head_step_no_fork; first done.
-  iIntros (σ1 ns κs nt) "[Hσ Hκs] !>"; iSplit; first by destruct n; auto with lia head_step.
+  iIntros (σ1 ns κs nt) "(Hσ & Hκs & Hsteps) !>"; iSplit; first by destruct n; auto with lia head_step.
   iIntros (κ v2 σ2 efs Hstep); inv_head_step.
   iMod (gen_heap_alloc_big _ (heap_array _ (replicate (Z.to_nat n) v)) with "Hσ")
     as "(Hσ & Hl & Hm)".
   { apply heap_array_map_disjoint.
     rewrite replicate_length Z2Nat.id; auto with lia. }
-  iModIntro; do 2 (iSplit; first done). iFrame "Hσ Hκs". iApply "HΦ".
+  iMod (steps_auth_update_S with "Hsteps") as "Hsteps".
+  iModIntro; do 2 (iSplit; first done). iFrame "Hσ Hκs Hsteps". iApply "HΦ".
   iApply big_sepL_sep. iSplitL "Hl".
   - by iApply heap_array_to_seq_mapsto.
   - iApply (heap_array_to_seq_meta with "Hm"). by rewrite replicate_length.
@@ -263,10 +368,11 @@ Lemma twp_free s E l v :
   [[{ RET LitV LitUnit; True }]].
 Proof.
   iIntros (Φ) "Hl HΦ". iApply twp_lift_atomic_head_step_no_fork; first done.
-  iIntros (σ1 ns κs nt) "[Hσ Hκs] !>". iDestruct (gen_heap_valid with "Hσ Hl") as %?.
+  iIntros (σ1 ns κs nt) "(Hσ & Hκs & Hsteps) !>". iDestruct (gen_heap_valid with "Hσ Hl") as %?.
   iSplit; first by eauto with head_step.
   iIntros (κ v2 σ2 efs Hstep); inv_head_step.
   iMod (gen_heap_update with "Hσ Hl") as "[$ Hl]".
+  iMod (steps_auth_update_S with "Hsteps") as "Hsteps".
   iModIntro. iSplit; first done. iSplit; first done. iFrame. by iApply "HΦ".
 Qed.
 Lemma wp_free s E l v :
@@ -281,10 +387,12 @@ Lemma twp_load s E l dq v :
   [[{ l ↦{dq} v }]] Load (Val $ LitV $ LitLoc l) @ s; E [[{ RET v; l ↦{dq} v }]].
 Proof.
   iIntros (Φ) "Hl HΦ". iApply twp_lift_atomic_head_step_no_fork; first done.
-  iIntros (σ1 ns κs nt) "[Hσ Hκs] !>". iDestruct (gen_heap_valid with "Hσ Hl") as %?.
+  iIntros (σ1 ns κs nt) "(Hσ & Hκs & Hsteps) !>". iDestruct (gen_heap_valid with "Hσ Hl") as %?.
   iSplit; first by eauto with head_step.
   iIntros (κ v2 σ2 efs Hstep); inv_head_step.
-  iModIntro; iSplit=> //. iSplit; first done. iFrame. by iApply "HΦ".
+  iMod (steps_auth_update_S with "Hsteps") as "Hsteps".
+  iModIntro. iSplit; [done|]. iSplit; [done|]. iFrame.
+  by iApply "HΦ".
 Qed.
 Lemma wp_load s E l dq v :
   {{{ ▷ l ↦{dq} v }}} Load (Val $ LitV $ LitLoc l) @ s; E {{{ RET v; l ↦{dq} v }}}.
@@ -298,9 +406,10 @@ Lemma twp_store s E l v' v :
   [[{ RET LitV LitUnit; l ↦ v }]].
 Proof.
   iIntros (Φ) "Hl HΦ". iApply twp_lift_atomic_head_step_no_fork; first done.
-  iIntros (σ1 ns κs nt) "[Hσ Hκs] !>". iDestruct (gen_heap_valid with "Hσ Hl") as %?.
+  iIntros (σ1 ns κs nt) "(Hσ & Hκs & Hsteps) !>". iDestruct (gen_heap_valid with "Hσ Hl") as %?.
   iSplit; first by eauto with head_step.
   iIntros (κ v2 σ2 efs Hstep); inv_head_step.
+  iMod (steps_auth_update_S with "Hsteps") as "Hsteps".
   iMod (gen_heap_update with "Hσ Hl") as "[$ Hl]".
   iModIntro. iSplit; first done. iSplit; first done. iFrame. by iApply "HΦ".
 Qed.
@@ -317,9 +426,10 @@ Lemma twp_xchg s E l v' v :
   [[{ RET v'; l ↦ v }]].
 Proof.
   iIntros (Φ) "Hl HΦ". iApply twp_lift_atomic_head_step_no_fork; first done.
-  iIntros (σ1 ns κs nt) "[Hσ Hκs] !>". iDestruct (gen_heap_valid with "Hσ Hl") as %?.
+  iIntros (σ1 ns κs nt) "(Hσ & Hκs & Hsteps) !>". iDestruct (gen_heap_valid with "Hσ Hl") as %?.
   iSplit; first by eauto with head_step.
   iIntros (κ v2 σ2 efs Hstep); inv_head_step.
+  iMod (steps_auth_update_S with "Hsteps") as "Hsteps".
   iMod (gen_heap_update with "Hσ Hl") as "[$ Hl]".
   iModIntro. iSplit; first done. iSplit; first done. iFrame. by iApply "HΦ".
 Qed.
@@ -337,10 +447,11 @@ Lemma twp_cmpxchg_fail s E l dq v' v1 v2 :
   [[{ RET PairV v' (LitV $ LitBool false); l ↦{dq} v' }]].
 Proof.
   iIntros (?? Φ) "Hl HΦ". iApply twp_lift_atomic_head_step_no_fork; first done.
-  iIntros (σ1 ns κs nt) "[Hσ Hκs] !>". iDestruct (gen_heap_valid with "Hσ Hl") as %?.
+  iIntros (σ1 ns κs nt) "(Hσ & Hκs & Hsteps) !>". iDestruct (gen_heap_valid with "Hσ Hl") as %?.
   iSplit; first by eauto with head_step.
   iIntros (κ v2' σ2 efs Hstep); inv_head_step.
   rewrite bool_decide_false //.
+  iMod (steps_auth_update_S with "Hsteps") as "Hsteps".
   iModIntro; iSplit; first done. iSplit; first done. iFrame. by iApply "HΦ".
 Qed.
 Lemma wp_cmpxchg_fail s E l dq v' v1 v2 :
@@ -358,11 +469,12 @@ Lemma twp_cmpxchg_suc s E l v1 v2 v' :
   [[{ RET PairV v' (LitV $ LitBool true); l ↦ v2 }]].
 Proof.
   iIntros (?? Φ) "Hl HΦ". iApply twp_lift_atomic_head_step_no_fork; first done.
-  iIntros (σ1 ns κs nt) "[Hσ Hκs] !>". iDestruct (gen_heap_valid with "Hσ Hl") as %?.
+  iIntros (σ1 ns κs nt) "(Hσ & Hκs & Hsteps) !>". iDestruct (gen_heap_valid with "Hσ Hl") as %?.
   iSplit; first by eauto with head_step.
   iIntros (κ v2' σ2 efs Hstep); inv_head_step.
   rewrite bool_decide_true //.
   iMod (gen_heap_update with "Hσ Hl") as "[$ Hl]".
+  iMod (steps_auth_update_S with "Hsteps") as "Hsteps".
   iModIntro. iSplit; first done. iSplit; first done. iFrame. by iApply "HΦ".
 Qed.
 Lemma wp_cmpxchg_suc s E l v1 v2 v' :
@@ -379,10 +491,11 @@ Lemma twp_faa s E l i1 i2 :
   [[{ RET LitV (LitInt i1); l ↦ LitV (LitInt (i1 + i2)) }]].
 Proof.
   iIntros (Φ) "Hl HΦ". iApply twp_lift_atomic_head_step_no_fork; first done.
-  iIntros (σ1 ns κs nt) "[Hσ Hκs] !>". iDestruct (gen_heap_valid with "Hσ Hl") as %?.
+  iIntros (σ1 ns κs nt) "(Hσ & Hκs & Hsteps) !>". iDestruct (gen_heap_valid with "Hσ Hl") as %?.
   iSplit; first by eauto with head_step.
   iIntros (κ e2 σ2 efs Hstep); inv_head_step.
   iMod (gen_heap_update with "Hσ Hl") as "[$ Hl]".
+  iMod (steps_auth_update_S with "Hsteps") as "Hsteps".
   iModIntro. do 2 (iSplit; first done). iFrame. by iApply "HΦ".
 Qed.
 Lemma wp_faa s E l i1 i2 :
@@ -399,9 +512,10 @@ Lemma wp_new_proph s E :
   {{{ pvs p, RET (LitV (LitProphecy p)); proph p pvs }}}.
 Proof.
   iIntros (Φ) "_ HΦ". iApply wp_lift_atomic_head_step_no_fork; first done.
-  iIntros (σ1 ns κ κs nt) "[Hσ HR] !>". iSplit; first by eauto with head_step.
+  iIntros (σ1 ns κ κs nt) "(Hσ & HR & Hsteps) !>". iSplit; first by eauto with head_step.
   iIntros "!>" (v2 σ2 efs Hstep). inv_head_step.
   rename select proph_id into p.
+  iMod (steps_auth_update_S with "Hsteps") as "Hsteps".
   iMod (proph_map_new_proph p with "HR") as "[HR Hp]"; first done.
   iModIntro; iSplit; first done. iFrame. by iApply "HΦ".
 Qed.
@@ -458,20 +572,22 @@ Proof.
   (* TODO we should try to use a generic lifting lemma (and avoid [wp_unfold])
      here, since this breaks the WP abstraction. *)
   iIntros (A He) "Hp WPe". rewrite !wp_unfold /wp_pre /= He. simpl in *.
-  iIntros (σ1 ns κ κs nt) "[Hσ Hκ]".
+  iIntros (σ1 ns κ κs nt) "(Hσ & Hκ & Hsteps)".
   destruct κ as [|[p' [w' v']] κ' _] using rev_ind.
-  - iMod ("WPe" $! σ1 ns [] κs nt with "[$Hσ $Hκ]") as "[Hs WPe]". iModIntro. iSplit.
+  - iMod ("WPe" $! σ1 ns [] κs nt with "[$Hσ $Hκ $Hsteps]") as "[Hs WPe]".
+    iModIntro. iSplit.
     { iDestruct "Hs" as "%". iPureIntro. destruct s; [ by apply resolve_reducible | done]. }
     iIntros (e2 σ2 efs step). exfalso. apply step_resolve in step; last done.
     inv_head_step. match goal with H: ?κs ++ [_] = [] |- _ => by destruct κs end.
   - rewrite -assoc.
-    iMod ("WPe" $! σ1 0 _ _ nt with "[$Hσ $Hκ]") as "[Hs WPe]". iModIntro. iSplit.
+    iMod ("WPe" $! σ1 ns _ _ nt with "[$Hσ $Hκ $Hsteps]") as "[Hs WPe]". iModIntro. iSplit.
     { iDestruct "Hs" as %?. iPureIntro. destruct s; [ by apply resolve_reducible | done]. }
     iIntros (e2 σ2 efs step). apply step_resolve in step; last done.
     inv_head_step; simplify_list_eq.
     iMod ("WPe" $! (Val w') σ2 efs with "[%]") as "WPe".
     { by eexists [] _ _. }
-    iModIntro. iNext. iModIntro. iMod "WPe" as ">[[$ Hκ] WPe]".
+    iModIntro. iNext. iMod "WPe" as "WPe". iModIntro.
+    iApply (step_fupdN_wand with "WPe"); iIntros "> [($ & Hκ & $) WPe]".
     iMod (proph_map_resolve_proph p' (w',v') κs with "[$Hκ $Hp]") as (vs' ->) "[$ HPost]".
     iModIntro. rewrite !wp_unfold /wp_pre /=. iDestruct "WPe" as "[HΦ $]".
     iMod "HΦ". iModIntro. by iApply "HΦ".
diff --git a/iris_heap_lang/total_adequacy.v b/iris_heap_lang/total_adequacy.v
index a0389c87f65d1b4b0dc8a9a6605f81f89f981617..f947da923faec7dc30a5c7e9bc61e88cada29131 100644
--- a/iris_heap_lang/total_adequacy.v
+++ b/iris_heap_lang/total_adequacy.v
@@ -1,4 +1,5 @@
 From iris.proofmode Require Import proofmode.
+From iris.base_logic.lib Require Import mono_nat.
 From iris.program_logic Require Export total_adequacy.
 From iris.heap_lang Require Export adequacy.
 From iris.heap_lang Require Import proofmode notation.
@@ -12,9 +13,12 @@ Proof.
   iMod (gen_heap_init σ.(heap)) as (?) "[Hh _]".
   iMod (inv_heap_init loc (option val)) as (?) ">Hi".
   iMod (proph_map_init [] σ.(used_proph_id)) as (?) "Hp".
+  iMod (mono_nat_own_alloc 0) as (γ) "[Hsteps _]".
   iModIntro.
   iExists
-    (λ σ κs _, (gen_heap_interp σ.(heap) ∗ proph_map_interp κs σ.(used_proph_id))%I),
-    (λ _, True%I); iFrame.
-  iApply (Hwp (HeapGS _ _ _ _ _)). done.
+    (λ σ ns κs _, (gen_heap_interp σ.(heap) ∗
+                   proph_map_interp κs σ.(used_proph_id) ∗
+                   mono_nat_auth_own γ 1 ns)%I),
+    id, (λ _, True%I), _; iFrame.
+  by iApply (Hwp (HeapGS _ _ _ _ _ _ _)).
 Qed.