diff --git a/theories/heap_lang/array.v b/theories/heap_lang/array.v
index 43c64c41f916da7da52dca6e623a4895cbb4d4ca..e3769b7017c9ed245801eb3b9fdb4e57f52c1bc0 100644
--- a/theories/heap_lang/array.v
+++ b/theories/heap_lang/array.v
@@ -1,7 +1,7 @@
 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.program_logic Require Export total_weakestpre weakestpre.
 From iris.heap_lang Require Export lifting.
 From iris.heap_lang Require Import tactics notation.
 Set Default Proof Using "Type".
@@ -92,17 +92,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 +103,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 +123,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.
+
+(** Access to 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 +145,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 +169,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 +203,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 +225,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 +251,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 +270,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 +288,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 +305,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..4660d56046257f538f37df0dac256c7b5c8e8905 100644
--- a/theories/heap_lang/lifting.v
+++ b/theories/heap_lang/lifting.v
@@ -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 }}}