如何根據集合中的消息對 mpsc::Receivers 進行排序？

Question

我有一個包含std::sync::mpsc::Receiver<Message>集合，並希望根據其中的消息將它們分類到其他集合中。

我正在迭代接收器，測試他們是否收到消息，如果是，我測試消息性質，並根據它，我將接收器移動到另一個集合：

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);
            },
        }
    }
}

我不知道如何從receivers集合中刪除receiver receivers 。 由於我在迭代它時將它借用為不可變的，因此我無法在不遇到編譯錯誤的情況下修改它

我感覺我沒有使用正確的 Rust 工具來完成這個非常基本的算法。

為了說明我的問題，假設我正在將集合 A 中的彩色球分類到集合 B 和 C，但是我無法克隆這些球，而且我在拿球時只能看到球的顏色，這會破壞顏色接下一個球。

將球從 0 號標記到 6 號，將其顏色寫在其索引旁邊，然后將這些球放入好的集合中的解決方案是不方便的。

Answer 1

您可以使用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
    }
});

---

要將丟棄的項目移動到其他向量，您可以編寫自己的retain函數。 下面找到的算法會將丟棄的項目作為迭代器返回。 無法保證它們返回的順序。不幸的是，它不支持將通道移動到多個其他向量。

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 ..)
}

這可以像這樣使用：

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);

---

或者，您可以准確地獲得問題中描述的行為，但這無法使用標准庫中的工具安全地完成。

為了以正確處理恐慌的方式做到這一點，我們可以引入一種保護類型：

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,
}

我們現在可以為這種類型定義兩個助手：

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);
    }
}

我們現在可以定義一個析構函數來正確設置向量的長度，並且在發生恐慌的情況下，刪除我們尚未檢查的項目。

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)
        }
    }
}

最后我們可以使用我們的守衛來實現以下方法。 請注意，析構函數在函數返回之前正確設置了向量的長度。

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);
            }
        }
    }
}

使用這種方法，實現所需的功能應該是相當直接的。 在此處查看完整代碼。

Answer 2

您可以使用Vec::drain將項目移出向量。 這會移動所有項目，但我們可以通過將它們臨時移動到新的Vec來處理要保留的項目。

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)
}

當然，這每次都會分配和調整一個新的向量； 一個小小的改進是在兩個向量之間交換而不是每次分配一個。 我認為尚未穩定的Vec::drain_filter可能會有所幫助，但過濾器只能返回bool ，因此它無法傳達將項目移動到何處。

Rust Playground 中的可編譯版本

Answer 3

理想情況下，您可以將Vec::retain （最初由@AliceRyhl 展示，但后來編輯以包含更復雜的方法）與橫梁通道（ Clone ）或Rc<Receiver> ，后者也是（簡單地）可克隆的，但會產生極小的運行時成本和堆分配。 如果這些是不可接受的，我提出兩種變體：

[✓] 變體 1：假設集合是一個向量，使用索引進行迭代。 它不是超級優雅，但也不是那么糟糕，它以最少的麻煩和任何人都可以理解的代碼完成工作：

// [✓] 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;
    }
}

變體 2：將接收器存儲在Vec<Option<Receiver>> 。 然后你就可以輕松去除接收器沒有突變的Vec通過采取從它Option ，留下None在它的位置：

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()),
        }
    }
}

如果receivers中剩余的None是一個問題，請使用receivers.retain(Option::is_some)類的東西刪除它們。