如何在Rust动态迭代器适配器链中统计各阶段元素数量

问题背景 #

我有一条编译时数量未知的迭代器适配器链，基于初始迭代器构建。简化示例为按可整除性过滤数字范围并统计最终剩余元素数量：

let n: usize = 10;
let mut iter_chain: Box<dyn std::iter::Iterator<Item = usize>> = Box::new(1..=n);

for i in 2..n {
    iter_chain = Box::new(iter_chain.filter(move |j| j % i != 0));
}

println!("{}", iter_chain.count());

现在需要统计该过滤链每个阶段的元素数量，原本计划在每个filter之间插入inspect适配器递增对应计数器，但每个inspect需要借用外部可变计数器，导致存储计数器的Vector出现多次可变借用冲突：

let n: usize = 10;
let mut iter_chain: Box<dyn std::iter::Iterator<Item = usize>> = Box::new(1..=n);

let mut filter_overview = vec![0; n];
for i in 2..n {
    // 编译错误：cannot borrow `filter_overview` as mutable more than once at a time
    iter_chain = Box::new(iter_chain.inspect(|_| filter_overview[i] += 1)); 
    iter_chain = Box::new(iter_chain.filter(move |j| j % i != 0));
}

println!("{filter_overview:?} {}", iter_chain.count());

示例中filter_overview的预期值为[10, 5, 3, 3, 2, 2, 1, 1]，需要解决借用冲突或找到替代实现思路。

解决方案 #

方法1：利用RefCell+Arc实现共享可变访问 #

通过Arc共享所有权，结合RefCell的运行时借用检查，绕开编译期的可变借用限制。因为迭代链是顺序执行的，inspect的闭包不会同时访问计数器，运行时不会出现借用冲突：

use std::cell::RefCell;
use std::sync::Arc;

fn main() {
    let n: usize = 10;
    let mut iter_chain: Box<dyn std::iter::Iterator<Item = usize>> = Box::new(1..=n);

    let filter_overview = Arc::new(RefCell::new(vec![0; n]));
    // 初始化初始阶段元素数量
    filter_overview.borrow_mut()[0] = 10;

    for i in 2..n {
        let idx = i - 1;
        let overview_clone = Arc::clone(&filter_overview);
        
        iter_chain = Box::new(iter_chain.inspect(move |_| {
            overview_clone.borrow_mut()[idx] += 1;
        }));
        iter_chain = Box::new(iter_chain.filter(move |j| j % i != 0));
    }

    let final_count = iter_chain.count();
    println!("{:?} {}", filter_overview.borrow(), final_count);
}

方法2：原子类型实现无锁计数 #

使用AtomicUsize原子类型，配合Arc共享，通过原子操作实现线程安全的计数（即使单线程场景也能解决借用问题）：

use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;

fn main() {
    let n: usize = 10;
    let mut iter_chain: Box<dyn std::iter::Iterator<Item = usize>> = Box::new(1..=n);

    // 用Arc包裹原子计数器，方便闭包捕获
    let filter_overview: Vec<Arc<AtomicUsize>> = (0..n)
        .map(|_| Arc::new(AtomicUsize::new(0)))
        .collect();
    filter_overview[0].store(10, Ordering::Relaxed);

    for i in 2..n {
        let idx = i - 1;
        let counter = Arc::clone(&filter_overview[idx]);
        
        iter_chain = Box::new(iter_chain.inspect(move |_| {
            counter.fetch_add(1, Ordering::Relaxed);
        }));
        iter_chain = Box::new(iter_chain.filter(move |j| j % i != 0));
    }

    let final_count = iter_chain.count();
    // 转换为普通usize数组输出
    let overview: Vec<usize> = filter_overview
        .iter()
        .map(|counter| counter.load(Ordering::Relaxed))
        .collect();
    println!("{overview:?} {final_count}");
}

方法3：直接统计各阶段元素数（最简洁） #

放弃插入inspect的思路，每次过滤后直接克隆迭代器并计数，完全避免借用问题，代码更简洁：

fn main() {
    let n: usize = 10;
    let mut iter = 1..=n;
    let mut filter_overview = vec![iter.clone().count()];

    for i in 2..n {
        iter = iter.filter(move |j| j % i != 0);
        filter_overview.push(iter.clone().count());
    }

    let final_count = iter.count();
    println!("{filter_overview:?} {final_count}");
}

该方法的输出与预期完全一致，且无需处理复杂的借用逻辑，适合大多数场景。

内容的提问来源于stack exchange，提问作者lpewewq