diff --git a/theories/heap_lang/array.v b/theories/heap_lang/array.v
index 43c64c41f916da7da52dca6e623a4895cbb4d4ca..deacf72e8f933fea6c5b047d5c0dd1dfee487f56 100644
--- a/theories/heap_lang/array.v
+++ b/theories/heap_lang/array.v
@@ -1,7 +1,6 @@
 From stdpp Require Import fin_maps.
 From iris.bi Require Import lib.fractional.
 From iris.proofmode Require Import tactics.
-From iris.program_logic Require Export weakestpre.
 From iris.heap_lang Require Export lifting.
 From iris.heap_lang Require Import tactics notation.
 Set Default Proof Using "Type".
@@ -80,7 +79,7 @@ Proof.
   rewrite drop_insert; last by lia. done.
 Qed.
 
-(** Allocation *)
+(** * Rules for allocation *)
 Lemma mapsto_seq_array l q v n :
   ([∗ list] i ∈ seq 0 n, (l +ₗ (i : nat)) ↦{q} v) -∗
   l ↦∗{q} replicate n v.
@@ -92,17 +91,6 @@ Proof.
   setoid_rewrite <-loc_add_assoc. iApply "IH". done.
 Qed.
 
-Lemma wp_allocN s E v n :
-  0 < n →
-  {{{ True }}} AllocN (Val $ LitV $ LitInt $ n) (Val v) @ s; E
-  {{{ l, RET LitV (LitLoc l); l ↦∗ replicate (Z.to_nat n) v ∗
-         [∗ list] i ∈ seq 0 (Z.to_nat n), meta_token (l +ₗ (i : nat)) ⊤ }}}.
-Proof.
-  iIntros (Hzs Φ) "_ HΦ". iApply wp_allocN_seq; [done..|]. iNext.
-  iIntros (l) "Hlm". iApply "HΦ".
-  iDestruct (big_sepL_sep with "Hlm") as "[Hl $]".
-  by iApply mapsto_seq_array.
-Qed.
 Lemma twp_allocN s E v n :
   0 < n →
   [[{ True }]] AllocN (Val $ LitV $ LitInt $ n) (Val v) @ s; E
@@ -114,17 +102,16 @@ Proof.
   iDestruct (big_sepL_sep with "Hlm") as "[Hl $]".
   by iApply mapsto_seq_array.
 Qed.
-
-Lemma wp_allocN_vec s E v n :
+Lemma wp_allocN s E v n :
   0 < n →
-  {{{ True }}}
-    AllocN #n v @ s ; E
-  {{{ l, RET #l; l ↦∗ vreplicate (Z.to_nat n) v ∗
+  {{{ True }}} AllocN (Val $ LitV $ LitInt $ n) (Val v) @ s; E
+  {{{ l, RET LitV (LitLoc l); l ↦∗ replicate (Z.to_nat n) v ∗
          [∗ list] i ∈ seq 0 (Z.to_nat n), meta_token (l +ₗ (i : nat)) ⊤ }}}.
 Proof.
-  iIntros (Hzs Φ) "_ HΦ". iApply wp_allocN; [ lia | done | .. ]. iNext.
-  iIntros (l) "[Hl Hm]". iApply "HΦ". rewrite vec_to_list_replicate. iFrame.
+  iIntros (? Φ) "_ HΦ". iApply (twp_wp_step with "HΦ").
+  iApply twp_allocN; [auto..|]; iIntros (l) "H HΦ". by iApply "HΦ".
 Qed.
+
 Lemma twp_allocN_vec s E v n :
   0 < n →
   [[{ True }]]
@@ -135,19 +122,18 @@ Proof.
   iIntros (Hzs Φ) "_ HΦ". iApply twp_allocN; [ lia | done | .. ].
   iIntros (l) "[Hl Hm]". iApply "HΦ". rewrite vec_to_list_replicate. iFrame.
 Qed.
-
-(** Access to array elements *)
-
-Lemma wp_load_offset s E l q off vs v :
-  vs !! off = Some v →
-  {{{ ▷ l ↦∗{q} vs }}} ! #(l +ₗ off) @ s; E {{{ RET v; l ↦∗{q} vs }}}.
+Lemma wp_allocN_vec s E v n :
+  0 < n →
+  {{{ True }}}
+    AllocN #n v @ s ; E
+  {{{ l, RET #l; l ↦∗ vreplicate (Z.to_nat n) v ∗
+         [∗ list] i ∈ seq 0 (Z.to_nat n), meta_token (l +ₗ (i : nat)) ⊤ }}}.
 Proof.
-  iIntros (Hlookup Φ) "Hl HΦ".
-  iDestruct (update_array l _ _ _ _ Hlookup with "Hl") as "[Hl1 Hl2]".
-  iApply (wp_load with "Hl1"). iIntros "!> Hl1". iApply "HΦ".
-  iDestruct ("Hl2" $! v) as "Hl2". rewrite list_insert_id; last done.
-  iApply "Hl2". iApply "Hl1".
+  iIntros (? Φ) "_ HΦ". iApply (twp_wp_step with "HΦ").
+  iApply twp_allocN_vec; [auto..|]; iIntros (l) "H HΦ". by iApply "HΦ".
 Qed.
+
+(** * Rules for accessing array elements *)
 Lemma twp_load_offset s E l q off vs v :
   vs !! off = Some v →
   [[{ l ↦∗{q} vs }]] ! #(l +ₗ off) @ s; E [[{ RET v; l ↦∗{q} vs }]].
@@ -158,25 +144,21 @@ Proof.
   iDestruct ("Hl2" $! v) as "Hl2". rewrite list_insert_id; last done.
   iApply "Hl2". iApply "Hl1".
 Qed.
+Lemma wp_load_offset s E l q off vs v :
+  vs !! off = Some v →
+  {{{ ▷ l ↦∗{q} vs }}} ! #(l +ₗ off) @ s; E {{{ RET v; l ↦∗{q} vs }}}.
+Proof.
+  iIntros (? Φ) ">H HΦ". iApply (twp_wp_step with "HΦ").
+  iApply (twp_load_offset with "H"); [eauto..|]; iIntros "H HΦ". by iApply "HΦ".
+Qed.
 
-
-Lemma wp_load_offset_vec s E l q sz (off : fin sz) (vs : vec val sz) :
-  {{{ ▷ l ↦∗{q} vs }}} ! #(l +ₗ off) @ s; E {{{ RET vs !!! off; l ↦∗{q} vs }}}.
-Proof. apply wp_load_offset. by apply vlookup_lookup. Qed.
 Lemma twp_load_offset_vec s E l q sz (off : fin sz) (vs : vec val sz) :
   [[{ l ↦∗{q} vs }]] ! #(l +ₗ off) @ s; E [[{ RET vs !!! off; l ↦∗{q} vs }]].
 Proof. apply twp_load_offset. by apply vlookup_lookup. Qed.
+Lemma wp_load_offset_vec s E l q sz (off : fin sz) (vs : vec val sz) :
+  {{{ ▷ l ↦∗{q} vs }}} ! #(l +ₗ off) @ s; E {{{ RET vs !!! off; l ↦∗{q} vs }}}.
+Proof. apply wp_load_offset. by apply vlookup_lookup. Qed.
 
-
-Lemma wp_store_offset s E l off vs v :
-  is_Some (vs !! off) →
-  {{{ ▷ l ↦∗ vs }}} #(l +ₗ off) <- v @ s; E {{{ RET #(); l ↦∗ <[off:=v]> vs }}}.
-Proof.
-  iIntros ([w Hlookup] Φ) ">Hl HΦ".
-  iDestruct (update_array l _ _ _ _ Hlookup with "Hl") as "[Hl1 Hl2]".
-  iApply (wp_store with "Hl1"). iNext. iIntros "Hl1".
-  iApply "HΦ". iApply "Hl2". iApply "Hl1".
-Qed.
 Lemma twp_store_offset s E l off vs v :
   is_Some (vs !! off) →
   [[{ l ↦∗ vs }]] #(l +ₗ off) <- v @ s; E [[{ RET #(); l ↦∗ <[off:=v]> vs }]].
@@ -186,33 +168,27 @@ Proof.
   iApply (twp_store with "Hl1"). iIntros "Hl1".
   iApply "HΦ". iApply "Hl2". iApply "Hl1".
 Qed.
-
-Lemma wp_store_offset_vec s E l sz (off : fin sz) (vs : vec val sz) v :
-  {{{ ▷ l ↦∗ vs }}} #(l +ₗ off) <- v @ s; E {{{ RET #(); l ↦∗ vinsert off v vs }}}.
+Lemma wp_store_offset s E l off vs v :
+  is_Some (vs !! off) →
+  {{{ ▷ l ↦∗ vs }}} #(l +ₗ off) <- v @ s; E {{{ RET #(); l ↦∗ <[off:=v]> vs }}}.
 Proof.
-  setoid_rewrite vec_to_list_insert. apply wp_store_offset.
-  eexists. by apply vlookup_lookup.
+  iIntros (? Φ) ">H HΦ". iApply (twp_wp_step with "HΦ").
+  iApply (twp_store_offset with "H"); [eauto..|]; iIntros "H HΦ". by iApply "HΦ".
 Qed.
+
 Lemma twp_store_offset_vec s E l sz (off : fin sz) (vs : vec val sz) v :
   [[{ l ↦∗ vs }]] #(l +ₗ off) <- v @ s; E [[{ RET #(); l ↦∗ vinsert off v vs }]].
 Proof.
   setoid_rewrite vec_to_list_insert. apply twp_store_offset.
   eexists. by apply vlookup_lookup.
 Qed.
-
-Lemma wp_cmpxchg_suc_offset s E l off vs v' v1 v2 :
-  vs !! off = Some v' →
-  v' = v1 →
-  vals_compare_safe v' v1 →
-  {{{ ▷ l ↦∗ vs }}}
-    CmpXchg #(l +â‚— off) v1 v2 @ s; E
-  {{{ RET (v', #true); l ↦∗ <[off:=v2]> vs }}}.
+Lemma wp_store_offset_vec s E l sz (off : fin sz) (vs : vec val sz) v :
+  {{{ ▷ l ↦∗ vs }}} #(l +ₗ off) <- v @ s; E {{{ RET #(); l ↦∗ vinsert off v vs }}}.
 Proof.
-  iIntros (Hlookup ?? Φ) "Hl HΦ".
-  iDestruct (update_array l _ _ _ _ Hlookup with "Hl") as "[Hl1 Hl2]".
-  iApply (wp_cmpxchg_suc with "Hl1"); [done..|].
-  iNext. iIntros "Hl1". iApply "HΦ". iApply "Hl2". iApply "Hl1".
+  iIntros (Φ) ">H HΦ". iApply (twp_wp_step with "HΦ").
+  iApply (twp_store_offset_vec with "H"); [eauto..|]; iIntros "H HΦ". by iApply "HΦ".
 Qed.
+
 Lemma twp_cmpxchg_suc_offset s E l off vs v' v1 v2 :
   vs !! off = Some v' →
   v' = v1 →
@@ -226,17 +202,18 @@ Proof.
   iApply (twp_cmpxchg_suc with "Hl1"); [done..|].
   iIntros "Hl1". iApply "HΦ". iApply "Hl2". iApply "Hl1".
 Qed.
-
-Lemma wp_cmpxchg_suc_offset_vec s E l sz (off : fin sz) (vs : vec val sz) v1 v2 :
-  vs !!! off = v1 →
-  vals_compare_safe (vs !!! off) v1 →
+Lemma wp_cmpxchg_suc_offset s E l off vs v' v1 v2 :
+  vs !! off = Some v' →
+  v' = v1 →
+  vals_compare_safe v' v1 →
   {{{ ▷ l ↦∗ vs }}}
     CmpXchg #(l +â‚— off) v1 v2 @ s; E
-  {{{ RET (vs !!! off, #true); l ↦∗ vinsert off v2 vs }}}.
+  {{{ RET (v', #true); l ↦∗ <[off:=v2]> vs }}}.
 Proof.
-  intros. setoid_rewrite vec_to_list_insert. eapply wp_cmpxchg_suc_offset=> //.
-  by apply vlookup_lookup.
+  iIntros (??? Φ) ">H HΦ". iApply (twp_wp_step with "HΦ").
+  iApply (twp_cmpxchg_suc_offset with "H"); [eauto..|]; iIntros "H HΦ". by iApply "HΦ".
 Qed.
+
 Lemma twp_cmpxchg_suc_offset_vec s E l sz (off : fin sz) (vs : vec val sz) v1 v2 :
   vs !!! off = v1 →
   vals_compare_safe (vs !!! off) v1 →
@@ -247,22 +224,17 @@ Proof.
   intros. setoid_rewrite vec_to_list_insert. eapply twp_cmpxchg_suc_offset=> //.
   by apply vlookup_lookup.
 Qed.
-
-Lemma wp_cmpxchg_fail_offset s E l q off vs v0 v1 v2 :
-  vs !! off = Some v0 →
-  v0 ≠ v1 →
-  vals_compare_safe v0 v1 →
-  {{{ ▷ l ↦∗{q} vs }}}
+Lemma wp_cmpxchg_suc_offset_vec s E l sz (off : fin sz) (vs : vec val sz) v1 v2 :
+  vs !!! off = v1 →
+  vals_compare_safe (vs !!! off) v1 →
+  {{{ ▷ l ↦∗ vs }}}
     CmpXchg #(l +â‚— off) v1 v2 @ s; E
-  {{{ RET (v0, #false); l ↦∗{q} vs }}}.
+  {{{ RET (vs !!! off, #true); l ↦∗ vinsert off v2 vs }}}.
 Proof.
-  iIntros (Hlookup HNEq Hcmp Φ) ">Hl HΦ".
-  iDestruct (update_array l _ _ _ _ Hlookup with "Hl") as "[Hl1 Hl2]".
-  iApply (wp_cmpxchg_fail with "Hl1"); first done.
-  { destruct Hcmp; by [ left | right ]. }
-  iIntros "!> Hl1". iApply "HΦ". iDestruct ("Hl2" $! v0) as "Hl2".
-  rewrite list_insert_id; last done. iApply "Hl2". iApply "Hl1".
+  iIntros (?? Φ) ">H HΦ". iApply (twp_wp_step with "HΦ").
+  iApply (twp_cmpxchg_suc_offset_vec with "H"); [eauto..|]; iIntros "H HΦ". by iApply "HΦ".
 Qed.
+
 Lemma twp_cmpxchg_fail_offset s E l q off vs v0 v1 v2 :
   vs !! off = Some v0 →
   v0 ≠ v1 →
@@ -278,14 +250,18 @@ Proof.
   iIntros "Hl1". iApply "HΦ". iDestruct ("Hl2" $! v0) as "Hl2".
   rewrite list_insert_id; last done. iApply "Hl2". iApply "Hl1".
 Qed.
-
-Lemma wp_cmpxchg_fail_offset_vec s E l q sz (off : fin sz) (vs : vec val sz) v1 v2 :
-  vs !!! off ≠ v1 →
-  vals_compare_safe (vs !!! off) v1 →
+Lemma wp_cmpxchg_fail_offset s E l q off vs v0 v1 v2 :
+  vs !! off = Some v0 →
+  v0 ≠ v1 →
+  vals_compare_safe v0 v1 →
   {{{ ▷ l ↦∗{q} vs }}}
     CmpXchg #(l +â‚— off) v1 v2 @ s; E
-  {{{ RET (vs !!! off, #false); l ↦∗{q} vs }}}.
-Proof. intros. eapply wp_cmpxchg_fail_offset=> //. by apply vlookup_lookup. Qed.
+  {{{ RET (v0, #false); l ↦∗{q} vs }}}.
+Proof.
+  iIntros (??? Φ) ">H HΦ". iApply (twp_wp_step with "HΦ").
+  iApply (twp_cmpxchg_fail_offset with "H"); [eauto..|]; iIntros "H HΦ". by iApply "HΦ".
+Qed.
+
 Lemma twp_cmpxchg_fail_offset_vec s E l q sz (off : fin sz) (vs : vec val sz) v1 v2 :
   vs !!! off ≠ v1 →
   vals_compare_safe (vs !!! off) v1 →
@@ -293,17 +269,14 @@ Lemma twp_cmpxchg_fail_offset_vec s E l q sz (off : fin sz) (vs : vec val sz) v1
     CmpXchg #(l +â‚— off) v1 v2 @ s; E
   [[{ RET (vs !!! off, #false); l ↦∗{q} vs }]].
 Proof. intros. eapply twp_cmpxchg_fail_offset=> //. by apply vlookup_lookup. Qed.
+Lemma wp_cmpxchg_fail_offset_vec s E l q sz (off : fin sz) (vs : vec val sz) v1 v2 :
+  vs !!! off ≠ v1 →
+  vals_compare_safe (vs !!! off) v1 →
+  {{{ ▷ l ↦∗{q} vs }}}
+    CmpXchg #(l +â‚— off) v1 v2 @ s; E
+  {{{ RET (vs !!! off, #false); l ↦∗{q} vs }}}.
+Proof. intros. eapply wp_cmpxchg_fail_offset=> //. by apply vlookup_lookup. Qed.
 
-Lemma wp_faa_offset s E l off vs (i1 i2 : Z) :
-  vs !! off = Some #i1 →
-  {{{ ▷ l ↦∗ vs }}} FAA #(l +ₗ off) #i2 @ s; E
-  {{{ RET LitV (LitInt i1); l ↦∗ <[off:=#(i1 + i2)]> vs }}}.
-Proof.
-  iIntros (Hlookup Φ) "Hl HΦ".
-  iDestruct (update_array l _ _ _ _ Hlookup with "Hl") as "[Hl1 Hl2]".
-  iApply (wp_faa with "Hl1").
-  iNext. iIntros "Hl1". iApply "HΦ". iApply "Hl2". iApply "Hl1".
-Qed.
 Lemma twp_faa_offset s E l off vs (i1 i2 : Z) :
   vs !! off = Some #i1 →
   [[{ l ↦∗ vs }]] FAA #(l +ₗ off) #i2 @ s; E
@@ -314,15 +287,15 @@ Proof.
   iApply (twp_faa with "Hl1").
   iIntros "Hl1". iApply "HΦ". iApply "Hl2". iApply "Hl1".
 Qed.
-
-Lemma wp_faa_offset_vec s E l sz (off : fin sz) (vs : vec val sz) (i1 i2 : Z) :
-  vs !!! off = #i1 →
+Lemma wp_faa_offset s E l off vs (i1 i2 : Z) :
+  vs !! off = Some #i1 →
   {{{ ▷ l ↦∗ vs }}} FAA #(l +ₗ off) #i2 @ s; E
-  {{{ RET LitV (LitInt i1); l ↦∗ vinsert off #(i1 + i2) vs }}}.
+  {{{ RET LitV (LitInt i1); l ↦∗ <[off:=#(i1 + i2)]> vs }}}.
 Proof.
-  intros. setoid_rewrite vec_to_list_insert. apply wp_faa_offset=> //.
-  by apply vlookup_lookup.
+  iIntros (? Φ) ">H HΦ". iApply (twp_wp_step with "HΦ").
+  iApply (twp_faa_offset with "H"); [eauto..|]; iIntros "H HΦ". by iApply "HΦ".
 Qed.
+
 Lemma twp_faa_offset_vec s E l sz (off : fin sz) (vs : vec val sz) (i1 i2 : Z) :
   vs !!! off = #i1 →
   [[{ l ↦∗ vs }]] FAA #(l +ₗ off) #i2 @ s; E
@@ -331,7 +304,14 @@ Proof.
   intros. setoid_rewrite vec_to_list_insert. apply twp_faa_offset=> //.
   by apply vlookup_lookup.
 Qed.
-
+Lemma wp_faa_offset_vec s E l sz (off : fin sz) (vs : vec val sz) (i1 i2 : Z) :
+  vs !!! off = #i1 →
+  {{{ ▷ l ↦∗ vs }}} FAA #(l +ₗ off) #i2 @ s; E
+  {{{ RET LitV (LitInt i1); l ↦∗ vinsert off #(i1 + i2) vs }}}.
+Proof.
+  iIntros (? Φ) ">H HΦ". iApply (twp_wp_step with "HΦ").
+  iApply (twp_faa_offset_vec with "H"); [eauto..|]; iIntros "H HΦ". by iApply "HΦ".
+Qed.
 End lifting.
 
 Typeclasses Opaque array.
diff --git a/theories/heap_lang/lifting.v b/theories/heap_lang/lifting.v
index d8261d93360891807f1456820050984c089b0fb7..9f403fd1a81c6296a7484b44cef8045725870fe6 100644
--- a/theories/heap_lang/lifting.v
+++ b/theories/heap_lang/lifting.v
@@ -3,7 +3,7 @@ From iris.proofmode Require Import tactics.
 From iris.algebra Require Import auth gmap.
 From iris.base_logic Require Export gen_heap.
 From iris.base_logic.lib Require Export proph_map.
-From iris.program_logic Require Export weakestpre.
+From iris.program_logic Require Export weakestpre total_weakestpre.
 From iris.program_logic Require Import ectx_lifting total_ectx_lifting.
 From iris.heap_lang Require Export lang.
 From iris.heap_lang Require Import tactics notation.
@@ -271,24 +271,6 @@ Proof.
   rewrite big_opM_singleton; iDestruct "Hvs" as "[$ Hvs]". by iApply "IH".
 Qed.
 
-Lemma wp_allocN_seq s E v n :
-  0 < n →
-  {{{ True }}} AllocN (Val $ LitV $ LitInt $ n) (Val v) @ s; E
-  {{{ l, RET LitV (LitLoc l); [∗ list] i ∈ seq 0 (Z.to_nat n),
-      (l +ₗ (i : nat)) ↦ v ∗ meta_token (l +ₗ (i : nat)) ⊤ }}}.
-Proof.
-  iIntros (Hn Φ) "_ HΦ". iApply wp_lift_atomic_head_step_no_fork; first done.
-  iIntros (σ1 κ κs k) "[Hσ Hκs] !>"; iSplit; first by auto with lia.
-  iNext; iIntros (v2 σ2 efs Hstep); inv_head_step.
-  iMod (gen_heap_alloc_gen _ (heap_array l (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; iSplit; first done. iFrame "Hσ Hκs". 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.
-Qed.
 Lemma twp_allocN_seq s E v n :
   0 < n →
   [[{ True }]] AllocN (Val $ LitV $ LitInt $ n) (Val v) @ s; E
@@ -307,28 +289,29 @@ Proof.
   - by iApply heap_array_to_seq_mapsto.
   - iApply (heap_array_to_seq_meta with "Hm"). by rewrite replicate_length.
 Qed.
-
-Lemma wp_alloc s E v :
-  {{{ True }}} Alloc (Val v) @ s; E {{{ l, RET LitV (LitLoc l); l ↦ v ∗ meta_token l ⊤ }}}.
+Lemma wp_allocN_seq s E v n :
+  0 < n →
+  {{{ True }}} AllocN (Val $ LitV $ LitInt $ n) (Val v) @ s; E
+  {{{ l, RET LitV (LitLoc l); [∗ list] i ∈ seq 0 (Z.to_nat n),
+      (l +ₗ (i : nat)) ↦ v ∗ meta_token (l +ₗ (i : nat)) ⊤ }}}.
 Proof.
-  iIntros (Φ) "_ HΦ". iApply wp_allocN_seq; [auto with lia..|].
-  iIntros "!>" (l) "/= (? & _)". rewrite loc_add_0. iApply "HΦ"; iFrame.
+  iIntros (Hn Φ) "_ HΦ". iApply (twp_wp_step with "HΦ").
+  iApply twp_allocN_seq; [auto..|]; iIntros (l) "H HΦ". by iApply "HΦ".
 Qed.
+
 Lemma twp_alloc s E v :
   [[{ True }]] Alloc (Val v) @ s; E [[{ l, RET LitV (LitLoc l); l ↦ v ∗ meta_token l ⊤ }]].
 Proof.
   iIntros (Φ) "_ HΦ". iApply twp_allocN_seq; [auto with lia..|].
   iIntros (l) "/= (? & _)". rewrite loc_add_0. iApply "HΦ"; iFrame.
 Qed.
-
-Lemma wp_load s E l q v :
-  {{{ ▷ l ↦{q} v }}} Load (Val $ LitV $ LitLoc l) @ s; E {{{ RET v; l ↦{q} v }}}.
+Lemma wp_alloc s E v :
+  {{{ True }}} Alloc (Val v) @ s; E {{{ l, RET LitV (LitLoc l); l ↦ v ∗ meta_token l ⊤ }}}.
 Proof.
-  iIntros (Φ) ">Hl HΦ". iApply wp_lift_atomic_head_step_no_fork; first done.
-  iIntros (σ1 κ κs n) "[Hσ Hκs] !>". iDestruct (@gen_heap_valid with "Hσ Hl") as %?.
-  iSplit; first by eauto. iNext; iIntros (v2 σ2 efs Hstep); inv_head_step.
-  iModIntro; iSplit=> //. iFrame. by iApply "HΦ".
+  iIntros (Φ) "_ HΦ". iApply (twp_wp_step with "HΦ").
+  iApply twp_alloc; [auto..|]; iIntros (l) "H HΦ". by iApply "HΦ".
 Qed.
+
 Lemma twp_load s E l q v :
   [[{ l ↦{q} v }]] Load (Val $ LitV $ LitLoc l) @ s; E [[{ RET v; l ↦{q} v }]].
 Proof.
@@ -337,17 +320,13 @@ Proof.
   iSplit; first by eauto. iIntros (κ v2 σ2 efs Hstep); inv_head_step.
   iModIntro; iSplit=> //. iSplit; first done. iFrame. by iApply "HΦ".
 Qed.
-
-Lemma wp_store s E l v' v :
-  {{{ ▷ l ↦ v' }}} Store (Val $ LitV (LitLoc l)) (Val v) @ s; E
-  {{{ RET LitV LitUnit; l ↦ v }}}.
+Lemma wp_load s E l q v :
+  {{{ ▷ l ↦{q} v }}} Load (Val $ LitV $ LitLoc l) @ s; E {{{ RET v; l ↦{q} v }}}.
 Proof.
-  iIntros (Φ) ">Hl HΦ". iApply wp_lift_atomic_head_step_no_fork; first done.
-  iIntros (σ1 κ κs n) "[Hσ Hκs] !>". iDestruct (@gen_heap_valid with "Hσ Hl") as %?.
-  iSplit; first by eauto. iNext; iIntros (v2 σ2 efs Hstep); inv_head_step.
-  iMod (@gen_heap_update with "Hσ Hl") as "[$ Hl]".
-  iModIntro. iSplit=>//. iFrame. by iApply "HΦ".
+  iIntros (Φ) ">H HΦ". iApply (twp_wp_step with "HΦ").
+  iApply (twp_load with "H"); [auto..|]; iIntros "H HΦ". by iApply "HΦ".
 Qed.
+
 Lemma twp_store s E l v' v :
   [[{ l ↦ v' }]] Store (Val $ LitV $ LitLoc l) (Val v) @ s; E
   [[{ RET LitV LitUnit; l ↦ v }]].
@@ -358,18 +337,14 @@ Proof.
   iMod (@gen_heap_update with "Hσ Hl") as "[$ Hl]".
   iModIntro. iSplit=>//. iSplit; first done. iFrame. by iApply "HΦ".
 Qed.
-
-Lemma wp_cmpxchg_fail s E l q v' v1 v2 :
-  v' ≠ v1 → vals_compare_safe v' v1 →
-  {{{ ▷ l ↦{q} v' }}} CmpXchg (Val $ LitV $ LitLoc l) (Val v1) (Val v2) @ s; E
-  {{{ RET PairV v' (LitV $ LitBool false); l ↦{q} v' }}}.
+Lemma wp_store s E l v' v :
+  {{{ ▷ l ↦ v' }}} Store (Val $ LitV (LitLoc l)) (Val v) @ s; E
+  {{{ RET LitV LitUnit; l ↦ v }}}.
 Proof.
-  iIntros (?? Φ) ">Hl HΦ". iApply wp_lift_atomic_head_step_no_fork; first done.
-  iIntros (σ1 κ κs n) "[Hσ Hκs] !>". iDestruct (@gen_heap_valid with "Hσ Hl") as %?.
-  iSplit; first by eauto. iNext; iIntros (v2' σ2 efs Hstep); inv_head_step.
-  rewrite bool_decide_false //.
-  iModIntro; iSplit=> //. iFrame. by iApply "HΦ".
+  iIntros (Φ) ">H HΦ". iApply (twp_wp_step with "HΦ").
+  iApply (twp_store with "H"); [auto..|]; iIntros "H HΦ". by iApply "HΦ".
 Qed.
+
 Lemma twp_cmpxchg_fail s E l q v' v1 v2 :
   v' ≠ v1 → vals_compare_safe v' v1 →
   [[{ l ↦{q} v' }]] CmpXchg (Val $ LitV $ LitLoc l) (Val v1) (Val v2) @ s; E
@@ -381,19 +356,15 @@ Proof.
   rewrite bool_decide_false //.
   iModIntro; iSplit=> //. iSplit; first done. iFrame. by iApply "HΦ".
 Qed.
-
-Lemma wp_cmpxchg_suc s E l v1 v2 v' :
-  v' = v1 → vals_compare_safe v' v1 →
-  {{{ ▷ l ↦ v' }}} CmpXchg (Val $ LitV $ LitLoc l) (Val v1) (Val v2) @ s; E
-  {{{ RET PairV v' (LitV $ LitBool true); l ↦ v2 }}}.
+Lemma wp_cmpxchg_fail s E l q v' v1 v2 :
+  v' ≠ v1 → vals_compare_safe v' v1 →
+  {{{ ▷ l ↦{q} v' }}} CmpXchg (Val $ LitV $ LitLoc l) (Val v1) (Val v2) @ s; E
+  {{{ RET PairV v' (LitV $ LitBool false); l ↦{q} v' }}}.
 Proof.
-  iIntros (?? Φ) ">Hl HΦ". iApply wp_lift_atomic_head_step_no_fork; first done.
-  iIntros (σ1 κ κs n) "[Hσ Hκs] !>". iDestruct (@gen_heap_valid with "Hσ Hl") as %?.
-  iSplit; first by eauto. iNext; iIntros (v2' σ2 efs Hstep); inv_head_step.
-  rewrite bool_decide_true //.
-  iMod (@gen_heap_update with "Hσ Hl") as "[$ Hl]".
-  iModIntro. iSplit=>//. iFrame. by iApply "HΦ".
+  iIntros (?? Φ) ">H HΦ". iApply (twp_wp_step with "HΦ").
+  iApply (twp_cmpxchg_fail with "H"); [auto..|]; iIntros "H HΦ". by iApply "HΦ".
 Qed.
+
 Lemma twp_cmpxchg_suc s E l v1 v2 v' :
   v' = v1 → vals_compare_safe v' v1 →
   [[{ l ↦ v' }]] CmpXchg (Val $ LitV $ LitLoc l) (Val v1) (Val v2) @ s; E
@@ -406,17 +377,15 @@ Proof.
   iMod (@gen_heap_update with "Hσ Hl") as "[$ Hl]".
   iModIntro. iSplit=>//. iSplit; first done. iFrame. by iApply "HΦ".
 Qed.
-
-Lemma wp_faa s E l i1 i2 :
-  {{{ ▷ l ↦ LitV (LitInt i1) }}} FAA (Val $ LitV $ LitLoc l) (Val $ LitV $ LitInt i2) @ s; E
-  {{{ RET LitV (LitInt i1); l ↦ LitV (LitInt (i1 + i2)) }}}.
+Lemma wp_cmpxchg_suc s E l v1 v2 v' :
+  v' = v1 → vals_compare_safe v' v1 →
+  {{{ ▷ l ↦ v' }}} CmpXchg (Val $ LitV $ LitLoc l) (Val v1) (Val v2) @ s; E
+  {{{ RET PairV v' (LitV $ LitBool true); l ↦ v2 }}}.
 Proof.
-  iIntros (Φ) ">Hl HΦ". iApply wp_lift_atomic_head_step_no_fork; first done.
-  iIntros (σ1 κ κs n) "[Hσ Hκs] !>". iDestruct (@gen_heap_valid with "Hσ Hl") as %?.
-  iSplit; first by eauto. iNext; iIntros (v2' σ2 efs Hstep); inv_head_step.
-  iMod (@gen_heap_update with "Hσ Hl") as "[$ Hl]".
-  iModIntro. iSplit=>//. iFrame. by iApply "HΦ".
+  iIntros (?? Φ) ">H HΦ". iApply (twp_wp_step with "HΦ").
+  iApply (twp_cmpxchg_suc with "H"); [auto..|]; iIntros "H HΦ". by iApply "HΦ".
 Qed.
+
 Lemma twp_faa s E l i1 i2 :
   [[{ l ↦ LitV (LitInt i1) }]] FAA (Val $ LitV $ LitLoc l) (Val $ LitV $ LitInt i2) @ s; E
   [[{ RET LitV (LitInt i1); l ↦ LitV (LitInt (i1 + i2)) }]].
@@ -427,6 +396,13 @@ Proof.
   iMod (@gen_heap_update with "Hσ Hl") as "[$ Hl]".
   iModIntro. iSplit=>//. iSplit; first done. iFrame. by iApply "HΦ".
 Qed.
+Lemma wp_faa s E l i1 i2 :
+  {{{ ▷ l ↦ LitV (LitInt i1) }}} FAA (Val $ LitV $ LitLoc l) (Val $ LitV $ LitInt i2) @ s; E
+  {{{ RET LitV (LitInt i1); l ↦ LitV (LitInt (i1 + i2)) }}}.
+Proof.
+  iIntros (Φ) ">H HΦ". iApply (twp_wp_step with "HΦ").
+  iApply (twp_faa with "H"); [auto..|]; iIntros "H HΦ". by iApply "HΦ".
+Qed.
 
 Lemma wp_new_proph s E :
   {{{ True }}}
diff --git a/theories/heap_lang/proofmode.v b/theories/heap_lang/proofmode.v
index 2851575e316f241e211dd63dde579bfc62b41d47..98ea0e6758ebc15bcc7577466f5034ad664f91ba 100644
--- a/theories/heap_lang/proofmode.v
+++ b/theories/heap_lang/proofmode.v
@@ -1,8 +1,7 @@
 From iris.proofmode Require Import coq_tactics reduction.
 From iris.proofmode Require Export tactics.
-From iris.program_logic Require Export weakestpre total_weakestpre.
 From iris.program_logic Require Import atomic.
-From iris.heap_lang Require Export tactics lifting array.
+From iris.heap_lang Require Export tactics array.
 From iris.heap_lang Require Import notation.
 Set Default Proof Using "Type".
 Import uPred.
diff --git a/theories/program_logic/hoare.v b/theories/program_logic/hoare.v
index 9776e183ebab776f6c1b80115743636ff774c33b..a8f5850aef2ec8cae40d26bf891a30f03ae1b316 100644
--- a/theories/program_logic/hoare.v
+++ b/theories/program_logic/hoare.v
@@ -120,7 +120,7 @@ Lemma ht_frame_r s E P Φ R e :
 Proof. iIntros "#Hwp !> [HP $]". by iApply "Hwp". Qed.
 
 Lemma ht_frame_step_l s E1 E2 P R1 R2 e Φ :
-  to_val e = None → E2 ⊆ E1 →
+  TCEq (to_val e) None → E2 ⊆ E1 →
   (R1 ={E1,E2}=> ▷ |={E2,E1}=> R2) ∧ {{ P }} e @ s; E2 {{ Φ }}
   ⊢ {{ R1 ∗ P }} e @ s; E1 {{ λ v, R2 ∗ Φ v }}.
 Proof.
@@ -130,7 +130,7 @@ Proof.
 Qed.
 
 Lemma ht_frame_step_r s E1 E2 P R1 R2 e Φ :
-  to_val e = None → E2 ⊆ E1 →
+  TCEq (to_val e) None → E2 ⊆ E1 →
   (R1 ={E1,E2}=> ▷ |={E2,E1}=> R2) ∧ {{ P }} e @ s; E2 {{ Φ }}
   ⊢ {{ P ∗ R1 }} e @ s; E1 {{ λ v, Φ v ∗ R2 }}.
 Proof.
@@ -140,7 +140,7 @@ Proof.
 Qed.
 
 Lemma ht_frame_step_l' s E P R e Φ :
-  to_val e = None →
+  TCEq (to_val e) None →
   {{ P }} e @ s; E {{ Φ }} ⊢ {{ ▷ R ∗ P }} e @ s; E {{ v, R ∗ Φ v }}.
 Proof.
   iIntros (?) "#Hwp !> [HR HP]".
@@ -148,7 +148,7 @@ Proof.
 Qed.
 
 Lemma ht_frame_step_r' s E P Φ R e :
-  to_val e = None →
+  TCEq (to_val e) None →
   {{ P }} e @ s; E {{ Φ }} ⊢ {{ P ∗ ▷ R }} e @ s; E {{ v, Φ v ∗ R }}.
 Proof.
   iIntros (?) "#Hwp !> [HP HR]".
diff --git a/theories/program_logic/total_weakestpre.v b/theories/program_logic/total_weakestpre.v
index 0f25045e7acfdf3f1967e2a167dba87d5640a630..5a34be3273f20775ee548b7950bb908eb42def1f 100644
--- a/theories/program_logic/total_weakestpre.v
+++ b/theories/program_logic/total_weakestpre.v
@@ -253,6 +253,15 @@ Proof. iIntros "[H Hwp]". iApply (twp_wand with "Hwp H"). Qed.
 Lemma twp_wand_r s E e Φ Ψ :
   WP e @ s; E [{ Φ }] ∗ (∀ v, Φ v -∗ Ψ v) -∗ WP e @ s; E [{ Ψ }].
 Proof. iIntros "[Hwp H]". iApply (twp_wand with "Hwp H"). Qed.
+
+Lemma twp_wp_step s E e P Φ :
+  TCEq (to_val e) None →
+  ▷ P -∗
+  WP e @ s; E [{ v, P ={E}=∗ Φ v }] -∗ WP e @ s; E {{ Φ }}.
+Proof.
+  iIntros (?) "HP Hwp".
+  iApply (wp_step_fupd _ _ E _ P with "[HP]"); [auto..|]. by iApply twp_wp.
+Qed.
 End twp.
 
 (** Proofmode class instances *)
diff --git a/theories/program_logic/weakestpre.v b/theories/program_logic/weakestpre.v
index 5bc2534b1352a99fc53423aba051ea24aaef84db..210cac5070e9264beff00fc2fa01089783b9ad05 100644
--- a/theories/program_logic/weakestpre.v
+++ b/theories/program_logic/weakestpre.v
@@ -140,7 +140,7 @@ Proof.
 Qed.
 
 Lemma wp_step_fupd s E1 E2 e P Φ :
-  to_val e = None → E2 ⊆ E1 →
+  TCEq (to_val e) None → E2 ⊆ E1 →
   (|={E1,E2}▷=> P) -∗ WP e @ s; E2 {{ v, P ={E1}=∗ Φ v }} -∗ WP e @ s; E1 {{ Φ }}.
 Proof.
   rewrite !wp_unfold /wp_pre. iIntros (-> ?) "HR H".
@@ -221,24 +221,24 @@ Lemma wp_frame_r s E e Φ R : WP e @ s; E {{ Φ }} ∗ R ⊢ WP e @ s; E {{ v, 
 Proof. iIntros "[H ?]". iApply (wp_strong_mono with "H"); auto with iFrame. Qed.
 
 Lemma wp_frame_step_l s E1 E2 e Φ R :
-  to_val e = None → E2 ⊆ E1 →
+  TCEq (to_val e) None → E2 ⊆ E1 →
   (|={E1,E2}▷=> R) ∗ WP e @ s; E2 {{ Φ }} ⊢ WP e @ s; E1 {{ v, R ∗ Φ v }}.
 Proof.
   iIntros (??) "[Hu Hwp]". iApply (wp_step_fupd with "Hu"); try done.
   iApply (wp_mono with "Hwp"). by iIntros (?) "$$".
 Qed.
 Lemma wp_frame_step_r s E1 E2 e Φ R :
-  to_val e = None → E2 ⊆ E1 →
+  TCEq (to_val e) None → E2 ⊆ E1 →
   WP e @ s; E2 {{ Φ }} ∗ (|={E1,E2}▷=> R) ⊢ WP e @ s; E1 {{ v, Φ v ∗ R }}.
 Proof.
   rewrite [(WP _ @ _; _ {{ _ }} ∗ _)%I]comm; setoid_rewrite (comm _ _ R).
   apply wp_frame_step_l.
 Qed.
 Lemma wp_frame_step_l' s E e Φ R :
-  to_val e = None → ▷ R ∗ WP e @ s; E {{ Φ }} ⊢ WP e @ s; E {{ v, R ∗ Φ v }}.
+  TCEq (to_val e) None → ▷ R ∗ WP e @ s; E {{ Φ }} ⊢ WP e @ s; E {{ v, R ∗ Φ v }}.
 Proof. iIntros (?) "[??]". iApply (wp_frame_step_l s E E); try iFrame; eauto. Qed.
 Lemma wp_frame_step_r' s E e Φ R :
-  to_val e = None → WP e @ s; E {{ Φ }} ∗ ▷ R ⊢ WP e @ s; E {{ v, Φ v ∗ R }}.
+  TCEq (to_val e) None → WP e @ s; E {{ Φ }} ∗ ▷ R ⊢ WP e @ s; E {{ v, Φ v ∗ R }}.
 Proof. iIntros (?) "[??]". iApply (wp_frame_step_r s E E); try iFrame; eauto. Qed.
 
 Lemma wp_wand s E e Φ Ψ :