1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 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!();
}