Wrapping iterator to always return at least one element

[jgarvin]

Say I have an iterator b that yields T. I want to wrap it in another iterator a that yields Option<T> with the following behavior:

• If b is empty, then a yields None exactly once
• If b has elements, then a yields Some(x) for all x in b

Does this already exist somewhere?

This is handy when when iterating a tree from some node back up to the root -- starting at some node and going up the tree back to the root, I need to do work for every (parent, child) pair, but the root doesn't have a parent. So if I just do for parent in node.parents() I accidentally skip doing any work for the root. Functions in the loop body that take a node and its parent allow passing None for the parent to make this seamless.

[jgarvin]

I wrote my own implementation, here in case anybody needs it, but still curious if I'm reinventing the wheel. Also wouldn't mind review feedback

enum AtLeastOnceState {
    Start,
    Yielding,
    Done,
}

struct AtLeastOnce<T: Iterator> {
    underlying: T,
    state: AtLeastOnceState,
}

impl<T: Iterator> AtLeastOnce<T> {
    fn new(underlying: T) -> Self {
        Self {
            underlying,
            state: AtLeastOnceState::Start,
        }
    }
}

impl<T: Iterator> Iterator for AtLeastOnce<T> {
    type Item = Option<T::Item>;

    fn next(&mut self) -> Option<Self::Item> {
        match self.state {
            AtLeastOnceState::Start => {
                let first = self.underlying.next();
                match &first {
                    None => {
                        self.state = AtLeastOnceState::Done;
                        Some(None)
                    }
                    Some(x) => {
                        self.state = AtLeastOnceState::Yielding;
                        Some(first)
                    }
                }
            }
            AtLeastOnceState::Yielding => {
                let elem = self.underlying.next();
                match &elem {
                    None => {
                        self.state = AtLeastOnceState::Done;
                        None
                    }
                    Some(x) => Some(elem),
                }
            }
            AtLeastOnceState::Done => None,
        }
    }
}

[2e71828]

I din’t think there’s anything builtin, but you can combine some of the standard adapters to make it happen:

fn at_least_once<I:Iterator>(iter:I)->impl Iterator<Item=Option<I::Item>> {
    iter.map(Some)
        .chain(std::iter::once(None))
        .enumerate()
        .filter_map(|x| match x {
            (0,None) => Some(None),
            (_,None) => None,
            (_,Some(x)) => Some(Some(x)),
        })
}

fn main() {
    dbg!(at_least_once(0..4).collect::<Vec<_>>());
    dbg!(at_least_once(0..0).collect::<Vec<_>>());
}

Rust Playground

[jgarvin]

Ah that's way more concise, I hadn't thought about using enumerate to enable special first case behavior!

[2e71828]

Your version can be cleaned up a lot by matching on tuples instead of writing nested matches:

(untested)

fn next(&mut self)->… {
    use AtLeastOnceState::*;
    let next = self.underlying.next();
    let (state, ret) = match (self.state, next) {
        ( Done,     _       ) => ( Done,     None          ),
        ( Start,    None    ) => ( Done,     Some(None)    ),
        ( Yielding, None    ) => ( Done,     None          ),
        ( _,        Some(x) ) => ( Yielding, Some(Some(x)) ),
    }
    self.state = state;
    ret
}

[afetisov]

You can also get the desired behaviour with from_fn:

fn transform<I: Iterator>(mut it: I) -> impl Iterator<Item=Option<I::Item>> {
    let mut polled = false;
    std::iter::from_fn(move || {
        if !polled {
            polled = true;
            Some(it.next())
        } else {
            it.next().map(Some)
        }
    })
}

[jgarvin]

@afetisov wasn't familiar with from_fn, nice