连接两个迭代器,同时懒洋洋地构建第二个迭代器

时间:2018-03-23 18:04:06

标签: rust

我想要一个像Iterator::chain()这样的方法,它只在需要时计算参数迭代器。在以下代码中,永远不应调用expensive_function

use std::{thread, time};

fn expensive_function() -> Vec<u64> {
    thread::sleep(time::Duration::from_secs(5));
    vec![4, 5, 6]
}

pub fn main() {
    let nums = [1, 2, 3];
    for &i in nums.iter().chain(expensive_function().iter()) {
        if i > 2 {
            break;
        } else {
            println!("{}", i);
        }
    }
}

2 个答案:

答案 0 :(得分:5)

一种可能的方法:将昂贵的计算委托给迭代器适配器。

let nums = [1, 2, 3];
for i in nums.iter()
    .cloned()
    .chain([()].into_iter().flat_map(|_| expensive_function()))
{
    if i > 2 {
        break;
    } else {
        println!("{}", i);
    }
}

Playground

传递的迭代器是将虚拟单元值()平面映射到值列表的结果,这是惰性的。由于迭代器需要拥有该计算的相应结果,因此我选择从数组中复制数字。

答案 1 :(得分:4)

您可以创建自己的自定义迭代器适配器,仅在原始迭代器耗尽时才计算闭包。

trait IteratorExt: Iterator {
    fn chain_with<F, I>(self, f: F) -> ChainWith<Self, F, I::IntoIter>
    where
        Self: Sized,
        F: FnOnce() -> I,
        I: IntoIterator<Item = Self::Item>,
    {
        ChainWith {
            base: self,
            factory: Some(f),
            iterator: None,
        }
    }
}

impl<I: Iterator> IteratorExt for I {}

struct ChainWith<B, F, I> {
    base: B,
    factory: Option<F>,
    iterator: Option<I>,
}

impl<B, F, I> Iterator for ChainWith<B, F, I::IntoIter>
where
    B: Iterator,
    F: FnOnce() -> I,
    I: IntoIterator<Item = B::Item>,
{
    type Item = I::Item;
    fn next(&mut self) -> Option<Self::Item> {
        if let Some(b) = self.base.next() {
            return Some(b);
        }

        // Exhausted the first, generate the second

        if let Some(f) = self.factory.take() {
            self.iterator = Some(f().into_iter());
        }

        self.iterator
            .as_mut()
            .expect("There must be an iterator")
            .next()
    }
}
use std::{thread, time};

fn expensive_function() -> Vec<u64> {
    panic!("You lose, good day sir");
    thread::sleep(time::Duration::from_secs(5));
    vec![4, 5, 6]
}

pub fn main() {
    let nums = [1, 2, 3];
    for i in nums.iter().cloned().chain_with(|| expensive_function()) {
        if i > 2 {
            break;
        } else {
            println!("{}", i);
        }
    }
}