diff options
Diffstat (limited to 'compiler/rustc_data_structures/src/flat_map_in_place.rs')
-rw-r--r-- | compiler/rustc_data_structures/src/flat_map_in_place.rs | 72 |
1 files changed, 72 insertions, 0 deletions
diff --git a/compiler/rustc_data_structures/src/flat_map_in_place.rs b/compiler/rustc_data_structures/src/flat_map_in_place.rs new file mode 100644 index 000000000..f58844f28 --- /dev/null +++ b/compiler/rustc_data_structures/src/flat_map_in_place.rs @@ -0,0 +1,72 @@ +use smallvec::{Array, SmallVec}; +use std::ptr; +use thin_vec::ThinVec; + +pub trait FlatMapInPlace<T>: Sized { + fn flat_map_in_place<F, I>(&mut self, f: F) + where + F: FnMut(T) -> I, + I: IntoIterator<Item = T>; +} + +// The implementation of this method is syntactically identical for all the +// different vector types. +macro_rules! flat_map_in_place { + () => { + fn flat_map_in_place<F, I>(&mut self, mut f: F) + where + F: FnMut(T) -> I, + I: IntoIterator<Item = T>, + { + 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 { + // 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 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); + 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); + + read_i += 1; + write_i += 1; + } + } + } + + // write_i tracks the number of actually written new items. + self.set_len(write_i); + } + } + }; +} + +impl<T> FlatMapInPlace<T> for Vec<T> { + flat_map_in_place!(); +} + +impl<T, A: Array<Item = T>> FlatMapInPlace<T> for SmallVec<A> { + flat_map_in_place!(); +} + +impl<T> FlatMapInPlace<T> for ThinVec<T> { + flat_map_in_place!(); +} |