miri: avoid cloning AllocExtra

Author: RalfJung

Diff for: compiler/rustc_const_eval/src/const_eval/machine.rs

@@ -140,7 +140,7 @@ impl<K: Hash + Eq, V> interpret::AllocMap<K, V> for FxIndexMap<K, V> {
 
     #[inline(always)]
     fn filter_map_collect<T>(&self, mut f: impl FnMut(&K, &V) -> Option<T>) -> Vec<T> {
-        self.iter().filter_map(move |(k, v)| f(k, &*v)).collect()
+        self.iter().filter_map(move |(k, v)| f(k, v)).collect()
     }
 
     #[inline(always)]

Diff for: compiler/rustc_const_eval/src/interpret/memory.rs

@@ -993,11 +993,14 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
         bytes
     }
 
-    /// Find leaked allocations. Allocations reachable from `static_roots` or a `Global` allocation
-    /// are not considered leaked, as well as leaks whose kind's `may_leak()` returns true.
-    pub fn find_leaked_allocations(
-        &self,
-        static_roots: &[AllocId],
+    /// Find leaked allocations, remove them from memory and return them. Allocations reachable from
+    /// `static_roots` or a `Global` allocation are not considered leaked, as well as leaks whose
+    /// kind's `may_leak()` returns true.
+    ///
+    /// This is highly destructive, no more execution can happen after this!
+    pub fn take_leaked_allocations(
+        &mut self,
+        static_roots: impl FnOnce(&Self) -> &[AllocId],
     ) -> Vec<(AllocId, MemoryKind<M::MemoryKind>, Allocation<M::Provenance, M::AllocExtra, M::Bytes>)>
     {
         // Collect the set of allocations that are *reachable* from `Global` allocations.
@@ -1008,7 +1011,7 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
                 self.memory.alloc_map.filter_map_collect(move |&id, &(kind, _)| {
                     if Some(kind) == global_kind { Some(id) } else { None }
                 });
-            todo.extend(static_roots);
+            todo.extend(static_roots(self));
             while let Some(id) = todo.pop() {
                 if reachable.insert(id) {
                     // This is a new allocation, add the allocation it points to `todo`.
@@ -1023,13 +1026,15 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
         };
 
         // All allocations that are *not* `reachable` and *not* `may_leak` are considered leaking.
-        self.memory.alloc_map.filter_map_collect(|id, (kind, alloc)| {
-            if kind.may_leak() || reachable.contains(id) {
-                None
-            } else {
-                Some((*id, *kind, alloc.clone()))
-            }
-        })
+        let leaked: Vec<_> = self.memory.alloc_map.filter_map_collect(|&id, &(kind, _)| {
+            if kind.may_leak() || reachable.contains(&id) { None } else { Some(id) }
+        });
+        let mut result = Vec::new();
+        for &id in leaked.iter() {
+            let (kind, alloc) = self.memory.alloc_map.remove(&id).unwrap();
+            result.push((id, kind, alloc));
+        }
+        result
     }
 
     /// Runs the closure in "validation" mode, which means the machine's memory read hooks will be

Diff for: src/tools/miri/src/diagnostics.rs

@@ -473,14 +473,14 @@ pub fn report_leaks<'tcx>(
     leaks: Vec<(AllocId, MemoryKind, Allocation<Provenance, AllocExtra<'tcx>, MiriAllocBytes>)>,
 ) {
     let mut any_pruned = false;
-    for (id, kind, mut alloc) in leaks {
+    for (id, kind, alloc) in leaks {
         let mut title = format!(
             "memory leaked: {id:?} ({}, size: {:?}, align: {:?})",
             kind,
             alloc.size().bytes(),
             alloc.align.bytes()
         );
-        let Some(backtrace) = alloc.extra.backtrace.take() else {
+        let Some(backtrace) = alloc.extra.backtrace else {
             ecx.tcx.dcx().err(title);
             continue;
         };

Diff for: src/tools/miri/src/eval.rs

@@ -476,7 +476,7 @@ pub fn eval_entry<'tcx>(
         }
         // Check for memory leaks.
         info!("Additional static roots: {:?}", ecx.machine.static_roots);
-        let leaks = ecx.find_leaked_allocations(&ecx.machine.static_roots);
+        let leaks = ecx.take_leaked_allocations(|ecx| &ecx.machine.static_roots);
         if !leaks.is_empty() {
             report_leaks(&ecx, leaks);
             tcx.dcx().note("set `MIRIFLAGS=-Zmiri-ignore-leaks` to disable this check");

Diff for: src/tools/miri/src/machine.rs

@@ -321,7 +321,7 @@ impl ProvenanceExtra {
 }
 
 /// Extra per-allocation data
-#[derive(Debug, Clone)]
+#[derive(Debug)]
 pub struct AllocExtra<'tcx> {
     /// Global state of the borrow tracker, if enabled.
     pub borrow_tracker: Option<borrow_tracker::AllocState>,
@@ -338,6 +338,14 @@ pub struct AllocExtra<'tcx> {
     pub backtrace: Option<Vec<FrameInfo<'tcx>>>,
 }
 
+// We need a `Clone` impl because the machine passes `Allocation` through `Cow`...
+// but that should never end up actually cloning our `AllocExtra`.
+impl<'tcx> Clone for AllocExtra<'tcx> {
+    fn clone(&self) -> Self {
+        panic!("our allocations should never be cloned");
+    }
+}
+
 impl VisitProvenance for AllocExtra<'_> {
     fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
         let AllocExtra { borrow_tracker, data_race, weak_memory, backtrace: _ } = self;