Make IdFunctor::try_map_id panic-safe

2 files changed, 30 insertions, 24 deletions

diff --git a/compiler/rustc_data_structures/src/functor.rs b/compiler/rustc_data_structures/src/functor.rs
index 71ff762c714..a3d3f988344 100644
--- a/compiler/rustc_data_structures/src/functor.rs
+++ b/compiler/rustc_data_structures/src/functor.rs
@@ -34,38 +34,43 @@ impl<T> IdFunctor for Vec<T> {
     type Inner = T;
 
     #[inline]
-    fn try_map_id<F, E>(mut self, mut f: F) -> Result<Self, E>
+    fn try_map_id<F, E>(self, mut f: F) -> Result<Self, E>
     where
         F: FnMut(Self::Inner) -> Result<Self::Inner, E>,
     {
-        // FIXME: We don't really care about panics here and leak
-        // far more than we should, but that should be fine for now.
-        let len = self.len();
-        unsafe {
-            self.set_len(0);
-            let start = self.as_mut_ptr();
-            for i in 0..len {
-                let p = start.add(i);
-                match f(p.read()) {
-                    Ok(val) => p.write(val),
-                    Err(err) => {
-                        // drop all other elements in self
-                        // (current element was "moved" into the call to f)
-                        for j in (0..i).chain(i + 1..len) {
-                            start.add(j).drop_in_place();
-                        }
+        struct HoleVec<T> {
+            vec: Vec<mem::ManuallyDrop<T>>,
+            hole: Option<usize>,
+        }
 
-                        // returning will drop self, releasing the allocation
-                        // (len is 0 so elements will not be re-dropped)
-                        return Err(err);
+        impl<T> Drop for HoleVec<T> {
+            fn drop(&mut self) {
+                unsafe {
+                    for (index, slot) in self.vec.iter_mut().enumerate() {
+                        if self.hole != Some(index) {
+                            mem::ManuallyDrop::drop(slot);
+                        }
                     }
                 }
             }
-            // Even if we encountered an error, set the len back
-            // so we don't leak memory.
-            self.set_len(len);
         }
-        Ok(self)
+
+        unsafe {
+            let (ptr, length, capacity) = self.into_raw_parts();
+            let vec = Vec::from_raw_parts(ptr.cast(), length, capacity);
+            let mut hole_vec = HoleVec { vec, hole: None };
+
+            for (index, slot) in hole_vec.vec.iter_mut().enumerate() {
+                hole_vec.hole = Some(index);
+                let original = mem::ManuallyDrop::take(slot);
+                let mapped = f(original)?;
+                *slot = mem::ManuallyDrop::new(mapped);
+                hole_vec.hole = None;
+            }
+
+            mem::forget(hole_vec);
+            Ok(Vec::from_raw_parts(ptr, length, capacity))
+        }
     }
 }
 
diff --git a/compiler/rustc_data_structures/src/lib.rs b/compiler/rustc_data_structures/src/lib.rs
index e17724b72f8..181e5180d53 100644
--- a/compiler/rustc_data_structures/src/lib.rs
+++ b/compiler/rustc_data_structures/src/lib.rs
@@ -23,6 +23,7 @@
 #![feature(once_cell)]
 #![feature(test)]
 #![feature(thread_id_value)]
+#![feature(vec_into_raw_parts)]
 #![allow(rustc::default_hash_types)]
 #![deny(unaligned_references)]