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.