#![cfg_attr(not(feature = "std"), no_std)]
#[cfg(not(feature = "std"))]
extern crate alloc;
#[cfg(not(feature = "std"))]
use alloc::collections::vec_deque::VecDeque;
#[cfg(not(feature = "std"))]
use alloc::vec::Vec;
#[cfg(not(feature = "std"))]
use core::ptr;
#[cfg(feature = "std")]
use std::collections::VecDeque;
#[cfg(feature = "std")]
use std::ptr;
pub trait RetainMut<T> {
fn retain_mut<F>(&mut self, f: F)
where
F: FnMut(&mut T) -> bool;
}
impl<T> RetainMut<T> for Vec<T> {
fn retain_mut<F>(&mut self, mut f: F)
where
F: FnMut(&mut T) -> bool,
{
let original_len = self.len();
unsafe { self.set_len(0) };
struct BackshiftOnDrop<'a, T> {
v: &'a mut Vec<T>,
processed_len: usize,
deleted_cnt: usize,
original_len: usize,
}
impl<T> Drop for BackshiftOnDrop<'_, T> {
fn drop(&mut self) {
if self.deleted_cnt > 0 {
unsafe {
ptr::copy(
self.v.as_ptr().add(self.processed_len),
self.v
.as_mut_ptr()
.add(self.processed_len - self.deleted_cnt),
self.original_len - self.processed_len,
);
}
}
unsafe {
self.v.set_len(self.original_len - self.deleted_cnt);
}
}
}
let mut g = BackshiftOnDrop {
v: self,
processed_len: 0,
deleted_cnt: 0,
original_len,
};
while g.processed_len < original_len {
let cur = unsafe { &mut *g.v.as_mut_ptr().add(g.processed_len) };
if !f(cur) {
g.processed_len += 1;
g.deleted_cnt += 1;
unsafe { ptr::drop_in_place(cur) };
continue;
}
if g.deleted_cnt > 0 {
unsafe {
let hole_slot = g.v.as_mut_ptr().add(g.processed_len - g.deleted_cnt);
ptr::copy_nonoverlapping(cur, hole_slot, 1);
}
}
g.processed_len += 1;
}
drop(g);
}
}
impl<T> RetainMut<T> for VecDeque<T> {
fn retain_mut<F>(&mut self, mut f: F)
where
F: FnMut(&mut T) -> bool,
{
let len = self.len();
let mut del = 0;
for i in 0..len {
if !f(&mut self[i]) {
del += 1;
} else if del > 0 {
self.swap(i - del, i);
}
}
if del > 0 {
self.truncate(len - del);
}
}
}