如何根据集合中的消息对 mpsc::Receivers 进行排序？

Question

I have a collection containing std::sync::mpsc::Receiver<Message> and want to sort them into other collections depending on the message in it.

I'm iterating into the receivers, testing if they received a message, if so I test the message nature, and depending on it, I move the receiver to another collection:

for receiver in receivers.iter() {
    if let Ok(msg) = receiver.try_recv() {
        match msg {
            Message::Foo => {
                // TODO remove receiver from receivers
                foos.push(receiver);
            },
            Message::Bar => {
                // TODO remove receiver from receivers
                bars.push(receiver);
            },
        }
    }
}

I do not know how to remove receiver from the receivers collection. Since I'm borrowing it as immutable while iterating over it, I can not modify it without encountering compilation error

I have the feeling that I'm not using the right Rust tools to complete this very basic algorithm.

To illustrate my problem, imagine that I'm sorting the colored balls from collection A to collections B and C, but I can not clone the balls and I can only see the color of a ball when taking it, which destroys the color when I take the next ball.

The solution of labeling the balls from number 0 to number 6, writing down its color aside to its index and then taking the balls to put them in the good collection is inconvenient.

Answer 1

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 .

Answer 2

You can use Vec::drain to move items out of a vector. This moves all of the items, but we can handle the ones to be kept by moving them into a new Vec temporarily.

fn distribute(receivers: &mut Vec<Receiver<Message>>) -> (Vec<Receiver<Message>>, Vec<Receiver<Message>>) {
    let mut foos: Vec<Receiver<Message>> = Vec::new();
    let mut bars: Vec<Receiver<Message>> = Vec::new();
    let mut go_back: Vec<Receiver<Message>> = Vec::new();
    for receiver in receivers.drain(..) {
        if let Ok(msg) = receiver.try_recv() {
            match msg {
                Message::Foo => &mut foos,
                Message::Bar => &mut bars,
                _ => &mut go_back,
            }.push(receiver);
        }
    }
    *receivers = go_back;
    (foos, bars)
}

Of course, this allocates and resizes a new vector every time; a slight improvement would be to swap between two vectors instead of allocating one each time. I thought the not-yet-stable Vec::drain_filter might be helpful but the filter can only return a bool , so it can't communicate where to move the item.

Compilable version in Rust Playground

Answer 3

Ideally you'd combine Vec::retain (originally shown by @AliceRyhl, but later edited to include more complex approaches) with either the crossbeam channels, which are Clone , or with Rc<Receiver> , which is also (trivially) clonable, but incurs a miniscule run-time cost and a heap allocation. If those are unacceptable, I propose two variants:

[✓] Variant 1: Assuming the collection is a vector, iterate using indices. It's not super-elegant, but it's not that bad, and it gets the job done with minimum hassle and with code anyone can understand:

// [✓] this is the accepted variant
let mut pos = 0;
while pos < receivers.len() {
    if let Ok(msg) = receivers[pos].try_recv() {
        match msg {
            Message::Foo => foos.push(receivers.remove(pos)),
            Message::Bar => bars.push(receivers.remove(pos)),
            _ => pos += 1,
        }
    } else {
        pos += 1;
    }
}

Variant 2: Store receivers in Vec<Option<Receiver>> . Then you can easily remove a receiver without mutating the Vec by taking it from the Option and leaving None in its place:

for receiver in receivers.iter_mut() {
    if receiver.is_none() {
        continue;
    }
    if let Ok(msg) = receiver.as_ref().unwrap().try_recv() {
        match msg {
            Message::Foo => foos.push(receiver.take().unwrap()),
            Message::Bar => bars.push(receiver.take().unwrap()),
        }
    }
}

If None s remaining in receivers are an issue, remove them with something like receivers.retain(Option::is_some) .