如何实现Rust中的try_retain()？带错误中断的集合元素过滤

是否存在类似try_retain()的方法？能否对集合元素应用FnMut(&mut V) -> Result<bool>，并根据结果处理：

• Ok(false)——移除元素并继续
• Ok(true)——保留元素并继续
• Err(_)——停止并返回错误

本质上是如何编写try_retain()，等价于以下C++代码：

struct S {
    // ...
};

// 更新`s`（可能还会修改外部状态），返回true表示`s`不再需要
// 出错时抛出异常
bool bar(S& s);

void foo(std::map<int, S>& m) {
    try
    {
        for(auto it = m.begin(), it_end = m.end(); it != it_end; ) {
            if (bar(it->second))
                m.erase(it++);
            else
                ++it;
        }
    }
    catch(...)
    {
        printf("bailed early\n");
        throw;
    }
}

我多次遇到该模式，但不知道如何在Rust中无需两次遍历或额外内存实现它。

Rust标准库目前没有内置try_retain方法，但可以通过扩展trait为集合类型实现该功能，避免额外内存分配和二次遍历。下面以Vec和HashMap为例给出实现：

1. 定义通用扩展Trait #

use std::collections::HashMap;

pub trait TryRetain<V, E> {
    /// 尝试保留集合中满足条件的元素，处理过程中出错则立即返回错误
    /// - `Ok(true)`: 保留元素，继续遍历
    /// - `Ok(false)`: 移除元素，继续遍历
    /// - `Err(e)`: 停止遍历，返回错误
    fn try_retain<F>(&mut self, mut f: F) -> Result<(), E>
    where
        F: FnMut(&mut V) -> Result<bool, E>;
}

2. 为Vec实现TryRetain #

利用Vec的索引操作原地修改，无额外内存开销：

impl<V, E> TryRetain<V, E> for Vec<V> {
    fn try_retain<F>(&mut self, mut f: F) -> Result<(), E>
    where
        F: FnMut(&mut V) -> Result<bool, E>,
    {
        let mut idx = 0;
        while idx < self.len() {
            let should_retain = f(&mut self[idx])?;
            if !should_retain {
                // 移除元素后后续元素前移，idx不递增
                self.remove(idx);
            } else {
                idx += 1;
            }
        }
        Ok(())
    }
}

3. 为HashMap实现TryRetain #

使用raw_entry_mut实现无临时内存的原地修改（Rust 1.61+稳定）：

impl<K, V, E> TryRetain<V, E> for HashMap<K, V>
where
    K: Eq + std::hash::Hash,
{
    fn try_retain<F>(&mut self, mut f: F) -> Result<(), E>
    where
        F: FnMut(&mut V) -> Result<bool, E>,
    {
        let mut iter = self.raw_entry_mut();
        while let Some((_, value)) = iter.next() {
            let should_retain = f(value)?;
            if !should_retain {
                // 删除当前条目，迭代器自动指向下一个元素
                iter.remove_entry();
            }
        }
        Ok(())
    }
}

如果使用低版本Rust（无raw_entry_mut），可以临时存储待删除键，开销极小：

impl<K, V, E> TryRetain<V, E> for HashMap<K, V>
where
    K: Eq + std::hash::Hash + Clone,
{
    fn try_retain<F>(&mut self, mut f: F) -> Result<(), E>
    where
        F: FnMut(&mut V) -> Result<bool, E>,
    {
        let mut to_remove = Vec::new();
        for (key, value) in self.iter_mut() {
            let should_retain = f(value)?;
            if !should_retain {
                to_remove.push(key.clone());
            }
        }
        for key in to_remove {
            self.remove(&key);
        }
        Ok(())
    }
}

4. 使用示例 #

#[derive(Debug)]
struct S {
    value: i32,
}

// 模拟业务逻辑：value>5保留，<=5移除，等于3时返回错误
fn bar(s: &mut S) -> Result<bool, &'static str> {
    if s.value == 3 {
        Err("遇到错误值3")
    } else if s.value > 5 {
        Ok(true)
    } else {
        Ok(false)
    }
}

fn main() -> Result<(), &'static str> {
    let mut vec = vec![
        S { value: 2 },
        S { value: 6 },
        S { value: 3 },
        S { value: 7 },
    ];

    // 遇到错误会立即终止
    match vec.try_retain(bar) {
        Ok(_) => println!("处理完成: {:?}", vec),
        Err(e) => println!("提前终止: {}", e),
    }

    let mut map = HashMap::from([
        (1, S { value: 2 }),
        (2, S { value: 6 }),
        (3, S { value: 8 }),
    ]);
    map.try_retain(bar)?;
    println!("处理后的HashMap: {:?}", map);

    Ok(())
}

说明 #

• Vec的实现完全原地操作，遍历一次完成，无额外内存开销。
• HashMap的raw_entry_mut实现和C++版本逻辑一致，遍历一次完成，无临时内存分配。

内容的提问来源于stack exchange，提问作者C.M.