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

Question

I have a program which performs a scan operation on its argument:

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

I want to extend the iterator with a few values generated inside the function:

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

This doesn't compile because once(head) is an iterator of A, while xs.iter() is an iterator of &A.

I can implement Clone for A and put.cloned() after xs.iter() to fix this, but I don't want to clone the entire xs because it can be very long and A in the actual program is not cheap to clone.

I'm looking for a way to turn once(head) into an iterator of &A, but cannot find any method.

Is it possible to get it to work?

I think the best you can do is https://play.rust-lang.org/?version=stable&mode=debug&edition=2018&gist=5ecde66452279b0f186383d8ec61ac33, doesn't match the expected output because I miss calculated when I comment AR's answer — Stargateur, Mar 30 '19 at 07:30

Shepmaster · Accepted Answer · 2019-03-30T15:23:44.407

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

Yes:

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

Take a reference to the value:

iter::once(&head)

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

No. The compiler will even tell you that:

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

See also:

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?

Maybe. Since scan starts with the accumulator value, you can just use that instead of sticking it onto the iterator:

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

This would mean that the resulting iterator has one less item. If that's acceptable depends on your usage.

A more complicated solution is to have an enum that represents either a borrowed value or an owned value. You can then create an iterator of these enums from the input and the local value. The ownership of the local value is transferred to the returned iterator:

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

This isn't free — every invocation of scan's closure will perform conditional logic to test if the value is owned or borrowed.

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

1 Answers1