From 4001f5933eb0f959232afce75005fc34d94803f1 Mon Sep 17 00:00:00 2001
From: Ralf Jung <post@ralfj.de>
Date: Mon, 19 Aug 2024 08:25:48 +0200
Subject: [PATCH] make the cleanup functions private

also move "is there a borrow tracker" check out of the loop
---
 src/tools/miri/src/provenance_gc.rs | 50 ++++++++++++++---------------
 1 file changed, 25 insertions(+), 25 deletions(-)

diff --git a/src/tools/miri/src/provenance_gc.rs b/src/tools/miri/src/provenance_gc.rs
index 5c6edf4ddced..d4bed69c6707 100644
--- a/src/tools/miri/src/provenance_gc.rs
+++ b/src/tools/miri/src/provenance_gc.rs
@@ -184,6 +184,29 @@ impl LiveAllocs<'_, '_> {
     }
 }
 
+fn remove_unreachable_tags<'tcx>(this: &mut MiriInterpCx<'tcx>, tags: FxHashSet<BorTag>) {
+    // Avoid iterating all allocations if there's no borrow tracker anyway.
+    if this.machine.borrow_tracker.is_some() {
+        this.memory.alloc_map().iter(|it| {
+            for (_id, (_kind, alloc)) in it {
+                alloc.extra.borrow_tracker.as_ref().unwrap().remove_unreachable_tags(&tags);
+            }
+        });
+    }
+}
+
+fn remove_unreachable_allocs<'tcx>(this: &mut MiriInterpCx<'tcx>, allocs: FxHashSet<AllocId>) {
+    let allocs = LiveAllocs { ecx: this, collected: allocs };
+    this.machine.allocation_spans.borrow_mut().retain(|id, _| allocs.is_live(*id));
+    this.machine.symbolic_alignment.borrow_mut().retain(|id, _| allocs.is_live(*id));
+    this.machine.alloc_addresses.borrow_mut().remove_unreachable_allocs(&allocs);
+    if let Some(borrow_tracker) = &this.machine.borrow_tracker {
+        borrow_tracker.borrow_mut().remove_unreachable_allocs(&allocs);
+    }
+    // Clean up core (non-Miri-specific) state.
+    this.remove_unreachable_allocs(&allocs.collected);
+}
+
 impl<'tcx> EvalContextExt<'tcx> for crate::MiriInterpCx<'tcx> {}
 pub trait EvalContextExt<'tcx>: MiriInterpCxExt<'tcx> {
     fn run_provenance_gc(&mut self) {
@@ -202,30 +225,7 @@ pub trait EvalContextExt<'tcx>: MiriInterpCxExt<'tcx> {
         });
 
         // Based on this, clean up the interpreter state.
-        self.remove_unreachable_tags(tags);
-        self.remove_unreachable_allocs(alloc_ids);
-    }
-
-    fn remove_unreachable_tags(&mut self, tags: FxHashSet<BorTag>) {
-        let this = self.eval_context_mut();
-        this.memory.alloc_map().iter(|it| {
-            for (_id, (_kind, alloc)) in it {
-                if let Some(bt) = &alloc.extra.borrow_tracker {
-                    bt.remove_unreachable_tags(&tags);
-                }
-            }
-        });
-    }
-
-    fn remove_unreachable_allocs(&mut self, allocs: FxHashSet<AllocId>) {
-        let this = self.eval_context_mut();
-        let allocs = LiveAllocs { ecx: this, collected: allocs };
-        this.machine.allocation_spans.borrow_mut().retain(|id, _| allocs.is_live(*id));
-        this.machine.symbolic_alignment.borrow_mut().retain(|id, _| allocs.is_live(*id));
-        this.machine.alloc_addresses.borrow_mut().remove_unreachable_allocs(&allocs);
-        if let Some(borrow_tracker) = &this.machine.borrow_tracker {
-            borrow_tracker.borrow_mut().remove_unreachable_allocs(&allocs);
-        }
-        this.remove_unreachable_allocs(&allocs.collected);
+        remove_unreachable_tags(this, tags);
+        remove_unreachable_allocs(this, alloc_ids);
     }
 }