rust - 我可以使用在函数内部创建的值来扩展迭代器吗？

Question

我有一个 program它对其参数执行 scan 操作:

struct A(usize);
struct B(usize);

fn scan_something<'a>(xs: &'a [A]) -> impl Iterator<Item = B> + 'a {
    let accum = 0;

    xs.iter().scan(accum, |accum, x| {
        *accum += x.0;
        Some(B(*accum))
    })
}

我想用函数内部生成的一些值来扩展迭代器:

fn scan_something<'a>(xs: &'a [A]) -> impl Iterator<Item = B> + 'a {
    let accum = 0;
    let head: A = A(xs.len());

    use std::iter::once;
    once(head).chain(xs.iter()).scan(accum, |accum, x| {
        *accum += x.0;
        Some(B(*accum))
    })
}

这不会编译，因为 once(head) 是 A 的迭代器，而 xs.iter() 是 A 的迭代器&A.

我可以为 A 实现 Clone 并在 xs.iter() 之后放置 .cloned() 到解决这个问题，但我不想克隆整个 xs，因为它可能很长，而且实际程序中的 A 克隆起来并不便宜。

我正在寻找一种方法将 once(head) 转换为 &A 的迭代器，但找不到任何方法。

是否可以让它工作？

Answer 1

Can I extend an iterator with values created inside a function?

是的:

fn example<'a>(input: impl Iterator<Item = i32> + 'a) -> impl Iterator<Item = i32> + 'a {
    Some(42).into_iter().chain(input).chain(Some(99))
}

fn main() {
    for i in example(vec![1, 2, 3].into_iter()) {
        println!("{}", i);
    }
}

I'm looking for a way to turn once(head) into an iterator of &A

引用值:

iter::once(&head)

Is it possible to get [this specific code] to work?

没有。编译器甚至会告诉你:

error[E0515]: cannot return value referencing local variable `head`
  --> src/lib.rs:10:5
   |
10 |       iter::once(&head).chain(xs.iter()).scan(accum, |accum, x| {
   |       ^          ----- `head` is borrowed here
   |  _____|
   | |
11 | |         *accum += x.0;
12 | |         Some(B(*accum))
13 | |     })
   | |______^ returns a value referencing data owned by the current function

另见:

• Is there any way to return a reference to a variable created in a function?

Is it possible to get [something close to this code] to work?

也许吧。由于 scan 从累加器值开始，您可以直接使用它而不是将其粘贴到迭代器上:

fn scan_something<'a>(xs: &'a [A]) -> impl Iterator<Item = B> + 'a {
    xs.iter().scan(xs.len(), |accum, x| {
        *accum += x.0;
        Some(B(*accum))
    })
}

这意味着生成的迭代器少了一项。这是否可以接受取决于您的使用情况。

一个更复杂的解决方案是使用一个枚举来表示借用的值或拥有的值。然后，您可以根据输入 和 本地值创建这些枚举的迭代器。本地值的所有权转移到返回的迭代器:

struct A(usize);
struct B(usize);

use std::iter;

// `A` doesn't implement `Clone`; if it did, use `Cow`
enum OwnedOrBorrowed<'a, T> {
    Owned(T),
    Borrowed(&'a T),
}

impl<'a, T> std::ops::Deref for OwnedOrBorrowed<'a, T> {
    type Target = T;
    fn deref(&self) -> &T {
        match self {
            OwnedOrBorrowed::Owned(t) => t,
            OwnedOrBorrowed::Borrowed(t) => t,
        }
    }
}

fn scan_something<'a>(xs: &'a [A]) -> impl Iterator<Item = B> + 'a {
    let accum = 0;
    let head = OwnedOrBorrowed::Owned(A(xs.len()));

    let borrowed = xs.iter().map(OwnedOrBorrowed::Borrowed);

    iter::once(head).chain(borrowed).scan(accum, |accum, x| {
        *accum += x.0;
        Some(B(*accum))
    })
}

这不是免费的——每次调用 scan 的闭包都会执行条件逻辑来测试该值是被拥有还是被借用。

另见:

• Is it possible to return either a borrowed or owned type in Rust?
• How to abstract over a reference to a value or a value itself?