diff --git a/compiler/rustc_data_structures/src/functor.rs b/compiler/rustc_data_structures/src/functor.rs index 71ff762c714bf..a3d3f988344c6 100644 --- a/compiler/rustc_data_structures/src/functor.rs +++ b/compiler/rustc_data_structures/src/functor.rs @@ -34,38 +34,43 @@ impl IdFunctor for Vec { type Inner = T; #[inline] - fn try_map_id(mut self, mut f: F) -> Result + fn try_map_id(self, mut f: F) -> Result where F: FnMut(Self::Inner) -> Result, { - // 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 { + vec: Vec>, + hole: Option, + } - // returning will drop self, releasing the allocation - // (len is 0 so elements will not be re-dropped) - return Err(err); + impl Drop for HoleVec { + 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 e17724b72f8b8..181e5180d53d5 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)]