diff --git a/theories/heap_lang/adequacy.v b/theories/heap_lang/adequacy.v
index 486186e68bfc01a0208af921a08388102499604e..2ab23f0ecdef6b16299034a12782022694f0f670 100644
--- a/theories/heap_lang/adequacy.v
+++ b/theories/heap_lang/adequacy.v
@@ -4,17 +4,17 @@ From iris.heap_lang Require Import proofmode notation proph_map.
 From iris.proofmode Require Import tactics.
 Set Default Proof Using "Type".
 
-Class heap_prophPreG Σ := HeapProphPreG {
-  heap_proph_preG_iris :> invPreG Σ;
-  heap_proph_preG_heap :> gen_heapPreG loc val Σ;
-  heap_proph_preG_proph :> proph_mapPreG proph val Σ
+Class heapPreG Σ := HeapPreG {
+  heap_preG_iris :> invPreG Σ;
+  heap_preG_heap :> gen_heapPreG loc val Σ;
+  heap_preG_proph :> proph_mapPreG proph val Σ
 }.
 
 Definition heapΣ : gFunctors := #[invΣ; gen_heapΣ loc val; proph_mapΣ proph val].
-Instance subG_heapPreG {Σ} : subG heapΣ Σ → heap_prophPreG Σ.
+Instance subG_heapPreG {Σ} : subG heapΣ Σ → heapPreG Σ.
 Proof. solve_inG. Qed.
 
-Definition heap_adequacy Σ `{heap_prophPreG Σ} s e σ φ :
+Definition heap_adequacy Σ `{heapPreG Σ} s e σ φ :
   (∀ `{heapG Σ}, WP e @ s; ⊤ {{ v, ⌜φ v⌝ }}%I) →
   adequate s e σ (λ v _, φ v).
 Proof.
diff --git a/theories/heap_lang/total_adequacy.v b/theories/heap_lang/total_adequacy.v
index 7cb9e7932a783d75a3eee270efce3a174f327e24..9dcf3b23772bfeb8fa87e1a0c5a976308855c9bf 100644
--- a/theories/heap_lang/total_adequacy.v
+++ b/theories/heap_lang/total_adequacy.v
@@ -4,15 +4,14 @@ From iris.heap_lang Require Import proofmode notation.
 From iris.proofmode Require Import tactics.
 Set Default Proof Using "Type".
 
-(* TODO (MR) re-introduce lemma after we prove twp_total *)
-(*
 Definition heap_total Σ `{heapPreG Σ} s e σ φ :
   (∀ `{heapG Σ}, WP e @ s; ⊤ [{ v, ⌜φ v⌝ }]%I) →
   sn erased_step ([e], σ).
 Proof.
-  intros Hwp; eapply (twp_total _ _) with (κs := []); iIntros (?) "".
-  iMod (gen_heap_init σ) as (?) "Hh".
-  iModIntro. iExists (fun σ _ => gen_heap_ctx σ); iFrame.
-  iApply (Hwp (HeapG _ _ _)).
+  intros Hwp; eapply (twp_total _ _); iIntros (?) "".
+  iMod (gen_heap_init σ.1) as (?) "Hh".
+  iMod (proph_map_init [] σ.2) as (?) "Hp".
+  iModIntro.
+  iExists (fun σ κs => (gen_heap_ctx σ.1 ∗ proph_map_ctx κs σ.2)%I). iFrame.
+  iApply (Hwp (HeapG _ _ _ _)).
 Qed.
-*)