diff --git a/compiler/rustc_data_structures/src/flat_map_in_place.rs b/compiler/rustc_data_structures/src/flat_map_in_place.rs index e66b00b755760..6d718059f9c81 100644 --- a/compiler/rustc_data_structures/src/flat_map_in_place.rs +++ b/compiler/rustc_data_structures/src/flat_map_in_place.rs @@ -1,4 +1,4 @@ -use std::ptr; +use std::{mem, ptr}; use smallvec::{Array, SmallVec}; use thin_vec::ThinVec; @@ -13,39 +13,44 @@ pub trait FlatMapInPlace: Sized { // The implementation of this method is syntactically identical for all the // different vector types. macro_rules! flat_map_in_place { - () => { + ($vec:ident $( where T: $bound:path)?) => { fn flat_map_in_place(&mut self, mut f: F) where F: FnMut(T) -> I, I: IntoIterator, { + struct LeakGuard<'a, T $(: $bound)?>(&'a mut $vec); + + impl<'a, T $(: $bound)?> Drop for LeakGuard<'a, T> { + fn drop(&mut self) { + unsafe { + self.0.set_len(0); // make sure we just leak elements in case of panic + } + } + } + + let this = LeakGuard(self); + let mut read_i = 0; let mut write_i = 0; unsafe { - let mut old_len = self.len(); - self.set_len(0); // make sure we just leak elements in case of panic - - while read_i < old_len { + while read_i < this.0.len() { // move the read_i'th item out of the vector and map it // to an iterator - let e = ptr::read(self.as_ptr().add(read_i)); + let e = ptr::read(this.0.as_ptr().add(read_i)); let iter = f(e).into_iter(); read_i += 1; for e in iter { if write_i < read_i { - ptr::write(self.as_mut_ptr().add(write_i), e); + ptr::write(this.0.as_mut_ptr().add(write_i), e); write_i += 1; } else { // If this is reached we ran out of space // in the middle of the vector. // However, the vector is in a valid state here, // so we just do a somewhat inefficient insert. - self.set_len(old_len); - self.insert(write_i, e); - - old_len = self.len(); - self.set_len(0); + this.0.insert(write_i, e); read_i += 1; write_i += 1; @@ -54,20 +59,23 @@ macro_rules! flat_map_in_place { } // write_i tracks the number of actually written new items. - self.set_len(write_i); + this.0.set_len(write_i); + + // The ThinVec is in a sane state again. Prevent the LeakGuard from leaking the data. + mem::forget(this); } } }; } impl FlatMapInPlace for Vec { - flat_map_in_place!(); + flat_map_in_place!(Vec); } impl> FlatMapInPlace for SmallVec { - flat_map_in_place!(); + flat_map_in_place!(SmallVec where T: Array); } impl FlatMapInPlace for ThinVec { - flat_map_in_place!(); + flat_map_in_place!(ThinVec); }