How can a access the first and last item of an iterator after iterating on all of them?

I have a function that does more or less:

fn process(data: &[i32]) {
    let vec: Vec<_> = data.map(|x| f(x)).collect();
    if data.first() == data.last() {
        foo(vec)
    } else {
        bar(vec)
    }
}

I would like to replace the type of data by an Iterator<Item=i32> to make it more flexible. However, How can I rewrite the if data.first() == data.last()? Since I iterate on the whole collection, I will already access both, but I don't know how I can save them during the map().collect() to be able to compare them later.

I think something like this should work:

let mut iter = data.map(|x| f(x));
let first = iter.next();
let last = iter.last();
if first == last {
    ...
}

Though, if you only need first & last, it might be better to first extract them, and then map an f.

EDIT: Ah, if I actually read the question, I see why yo do this, you need first & last before mapping :sweat_smile:

I'd probably just do something side-effecting then:

let mut res = Vec::new();
let first = data.next();
res.extend(first.map(f));
let mut last = None;
res.extend(data.inspect(|it| last = Some(it)).map(f));
if first == last {
    
}
1 Like

For @matklad solution to be error free you need to call fuse on the iterator first.

This is essentially the same I was trying out few minutes ago. Many alternatives, can see someone posting with something cleaner.

fn process(mut data: impl Iterator<Item=i32>) {
    let first = data.next();
    let mut last = first.clone();
    
    let vec = match first {
        Some(fi) => std::iter::once(fi).chain(data).inspect(|&i| last = Some(i)).map(|x| f(x)).collect(),
        None => vec![],
    };
    
    if first == last {
        foo(vec)
    } else {
        bar(vec)
    }
}

Edit below (small addition), written after steffahn's posting.
Alternative to using match is;

    let vec = first
        .into_iter()
        .chain(data)
        .inspect(|&i| last = Some(i))
        .map(|x| f(x))
        .collect();

Does not need a fuse since next not called on data if first is None.

2 Likes

For still including the possibility of using &[i32], you need to be a bit more general than Iterator. For example this works:

use std::borrow::Borrow;

fn f(x: i32) -> i32 {
    x
}

fn foo(x: Vec<i32>) {}
fn bar(x: Vec<i32>) {}

fn process(data: impl IntoIterator<Item = impl Borrow<i32>>) {
    let mut data = data.into_iter().peekable();
    // first: Option<i32>
    let first = data.peek().map(|i| *i.borrow());
    let mut last = None;
    let vec = data
        .map(|i| {
            let i = *i.borrow();
            last = Some(i);
            f(i)
        })
        .collect();
    if first == last {
        foo(vec)
    } else {
        bar(vec)
    }
}

// still works on Iterator<Item = i32>
fn process_iterator(data: impl Iterator<Item = i32>) {
    process(data)
}

// still works on &[i32]
fn process_slice(data: &[i32]) {
    process(data)
}

Also I would probably use peekable() which should behave very similar to the solution of @jonh while being a bit easier to use IMO.

1 Like

Thanks all, it was really informative (and sorry, I read your answers on my phone and forgot to answer when I came back to my PC). I think I will keep the &[i32] because the readability suffers too much, but I will keep your answer in mind if I ever face this issue again.

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.