Is there something like `Iterator::fold` which returns all computed elements?

map computes a sequence y_i = f(x_i) and return that sequence as a new iterator. fold computes a sequence y_1 = f(y_0, x_1), ..., y_n = f(y_{n-1}, x_n) and returns just the last element y_n. Is there an Iterator function similar to fold which returns the whole sequence?

Here is an attempt:

fn fold_list<T, A, F>(items: &[T], a: &A, f: F) -> Vec<A>
where
    F: Fn(&A, &T) -> A
{
    let mut result = Vec::new();
    
    let mut it = items.iter();

    if let Some(first) = it.next() {
        result.push(f(&a, &first));
    }

    for (i, x) in it.enumerate() {
        result.push(f(&result[i], &x));
    }

    result
}

fn main() {
    println!("{:?}", fold_list(&[1, 2, 3, 4], &1, |a, &x| a * x));
    println!("{:?}", fold_list(&[1], &2, |a, &x| a * x));
    println!("{:?}", fold_list(&[], &0, |a, &x: &i32| a * x));
}

are you looking for Iterator::scan()?

5 Likes

Here’d be a possible iterator version of your code

fn fold_list<T, A>(
    items: impl IntoIterator<Item = T>,
    a: &A,
    mut f: impl FnMut(&A, T) -> A,
) -> impl Iterator<Item = A> {

    let mut items = items.into_iter();
    let mut previous = items.next().map(|t| f(a, t));
    std::iter::from_fn(move || {
        previous.take().map(|p| {
            previous = items.next().map(|t| f(&p, t));
            p
        })
    })
}

use itertools::Itertools; // <- used for .format(…sep…) method to print the iterator like a list

fn main() {
    println!("[{:?}]", fold_list(&[1, 2, 3, 4], &1, |a, &x| a * x).format(", "));
    println!("[{:?}]", fold_list(&[1], &2, |a, &x| a * x).format(", "));
    println!("[{:?}]", fold_list(&[], &0, |a, &x: &i32| a * x).format(", "));
}
3 Likes

Actually, with std::iter::successors, it’s even nicer

fn fold_list<T, A>(
    items: impl IntoIterator<Item = T>,
    a: &A,
    mut f: impl FnMut(&A, T) -> A,
) -> impl Iterator<Item = A> {

    let mut items = items.into_iter();
    let first = items.next().map(|t| f(a, t));
    std::iter::successors(first, move |p| items.next().map(|t| f(p, t)))
}

Edit: Less DRY version:

fn fold_list<T, A>(
    items: impl IntoIterator<Item = T>,
    a: &A,
    mut f: impl FnMut(&A, T) -> A,
) -> impl Iterator<Item = A> {

    let mut items = items.into_iter();
    let mut next = move |a: &A| items.next().map(|t| f(a, t));
    std::iter::successors(next(a), next)
}
3 Likes

Yes exactly that! I swear I went through the Iterator documentation at least twice before posting :laughing:

1 Like

Wow, I just read the docs for scan, but I would never have guessed from the first read that that is what it does. Maybe adding @danvil's explanation of the adaptor to the doc would be helpful.

1 Like

This topic was automatically closed 90 days after the last reply. We invite you to open a new topic if you have further questions or comments.