You can use Vec::retain.
receivers.retain(|receiver| {
if let Ok(msg) = receiver.try_recv() {
match msg {
Message::Foo => {
foos.push(receiver);
},
Message::Bar => {
bars.push(receiver);
},
}
// By returning false, the receiver is removed.
false
} else {
// By returning true, the receiver is kept.
true
}
});
To have the discarded items moved to some other vector, you can write your own retain function. The algorithm found below will return the discarded items as an iterator. There is no guarantee of the ordering they are returned in. Unfortunately it does not support moving the channels to more than one other vector.
pub fn retain2<T, F>(vec: &mut Vec<T>, mut f: F) -> impl Iterator<Item = T> + '_
where
F: FnMut(&mut T) -> bool,
{
let len = vec.len();
let mut del = 0;
{
let v = &mut **vec;
for i in 0..len {
if !f(&mut v[i]) {
del += 1;
} else if del > 0 {
v.swap(i - del, i);
}
}
}
vec.drain(len - del ..)
}
This can be used like this:
let iter = retain2(&mut receivers, |receiver| {
if let Ok(msg) = receiver.try_recv() {
match msg {
Message::Foo => {
// By returning false, the receiver is sent to the iterator.
false
},
}
} else {
// By returning true, the receiver is kept in `receivers`.
true
}
});
// Append all of the removed iterators to this other vector.
// They are not guaranteed to be returned in the original order.
foos.extend(iter);
Alternatively, you can obtain the behavior described in the question exactly, but this cannot be done safely using the tools in the standard library.
To do this in a way that properly handles panics, we can introduce a guard type:
struct RetainGuard<'a, T> {
vec: &'a mut Vec<T>,
// How many items have been taken?
taken: usize,
// How many items have been inserted again?
// This must never be larger than `taken`.
inserted: usize,
// The original length of the vector.
vec_len: usize,
}
We can now define two helpers for this type:
impl<'a, T> RetainGuard<'a, T> {
fn take(&mut self) -> Option<T> {
let i = self.taken;
if i == self.vec_len {
return None;
}
self.taken += 1;
unsafe {
Some(std::ptr::read(&mut self.vec[i]))
}
}
// Safety:
// Must not insert more items than has already been taken.
unsafe fn insert(&mut self, val: T) {
std::ptr::write(&mut self.vec[self.inserted], val);
self.inserted += 1;
debug_assert!(self.inserted <= self.taken);
}
}
And we can now define a destructor to properly set the length of the vector, and in the case of panics, drop the items we have yet to inspect.
impl<'a, T> Drop for RetainGuard<'a, T> {
fn drop(&mut self) {
unsafe {
self.vec.set_len(self.inserted);
// This is an empty slice unless the closure panics.
std::ptr::drop_in_place(self.not_taken_slice());
}
}
}
impl<'a, T> RetainGuard<'a, T> {
fn not_taken_slice(&mut self) -> &mut [T] {
unsafe {
let ptr = self.vec.as_mut_ptr();
let start = ptr.add(self.taken);
let len = self.vec_len - self.taken;
std::slice::from_raw_parts_mut(start, len)
}
}
}
Finally we can use our guard to implement the following method. Note that the destructor properly sets the length of the vector before the function returns.
pub fn retain3<T, F>(vec: &mut Vec<T>, mut f: F)
where
F: FnMut(T) -> Option<T>,
{
let mut guard = RetainGuard::new(vec);
while let Some(item) = guard.take() {
if let Some(item) = f(item) {
unsafe {
// SAFETY: We have called take more times than insert.
guard.insert(item);
}
}
}
}
Using this method, it should be rather straight-forward to implement the desired functionality. See the full code here.