use std::ptr;
use smallvec::{Array, SmallVec};
pub trait MoveMap<T>: Sized {
fn move_map<F>(self, mut f: F) -> Self where F: FnMut(T) -> T {
self.move_flat_map(|e| Some(f(e)))
}
fn move_flat_map<F, I>(self, f: F) -> Self
where F: FnMut(T) -> I,
I: IntoIterator<Item=T>;
}
impl<T> MoveMap<T> for Vec<T> {
fn move_flat_map<F, I>(mut self, mut f: F) -> Self
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);
while read_i < old_len {
let e = ptr::read(self.get_unchecked(read_i));
let iter = f(e).into_iter();
read_i += 1;
for e in iter {
if write_i < read_i {
ptr::write(self.get_unchecked_mut(write_i), e);
write_i += 1;
} else {
self.set_len(old_len);
self.insert(write_i, e);
old_len = self.len();
self.set_len(0);
read_i += 1;
write_i += 1;
}
}
}
self.set_len(write_i);
}
self
}
}
impl<T> MoveMap<T> for ::ptr::P<[T]> {
fn move_flat_map<F, I>(self, f: F) -> Self
where F: FnMut(T) -> I,
I: IntoIterator<Item=T>
{
::ptr::P::from_vec(self.into_vec().move_flat_map(f))
}
}
impl<T, A: Array<Item = T>> MoveMap<T> for SmallVec<A> {
fn move_flat_map<F, I>(mut self, mut f: F) -> Self
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);
while read_i < old_len {
let e = ptr::read(self.get_unchecked(read_i));
let iter = f(e).into_iter();
read_i += 1;
for e in iter {
if write_i < read_i {
ptr::write(self.get_unchecked_mut(write_i), e);
write_i += 1;
} else {
self.set_len(old_len);
self.insert(write_i, e);
old_len = self.len();
self.set_len(0);
read_i += 1;
write_i += 1;
}
}
}
self.set_len(write_i);
}
self
}
}