`Option::map_iter()` to create an iterator from `Option`

[ChayimFriedman2]

Proposal

Problem statement

It's common to have an Option containing an iterator, or something that can be transformed into an iterator, and wanting to get an iterator out of it. This can be done with option.into_iter().flat_map(|it| it.iter()). The verbosity is okay (although not great, and this isn't easily discoverable), but this is not as efficient as it can be, because FlatMap/Flatten need to handle iteration for generic iterators, which can be a lot slower. Internal iteration is often fine (the FlatMap can be optimized) but external iteration uses way more instructions than needed. As an example, see this godbolt link, which contains way more instructions than needed.

Code for the godbolt link:

unsafe extern "C" {
    safe fn black_box(v: &i32);
}

#[unsafe(no_mangle)]
pub fn foo(v: Option<&[i32]>) {
    for i in v.into_iter().flatten() {
        black_box(i)
    }
}

In addition, FlatMap/Flatten can't provide an accurate size_hint()/TrustedLen, which can make collecting even slower.

Motivating examples or use cases

The rust-analyzer codebase contain dozens (more than 20) instances of the abovementioned pattern.

Solution sketch

Include a method in Option, that takes a callback that returns an iterator and returns an iterator for that behaves like FlatMap but is specialized for Option.

I chose to call this method map_iter(), but I'm not really pleased with this name, I just don't have a better one. So I'll be happy to hear other opinions.

Alternatives

At least the size_hint() part is fixable I believe with specialization. The next() part is harder, it may not be possible to specialize it to the most efficient codegen.

There's also the argument that adding a method purely for the sake of optimization isn't a good idea. While I do think a special method is also more easily discoverable and nicer to type, the main motivation is indeed perf. But I do think if we can't get this level of perf from alternatives it's worth adding the method, even for perf sake only.

What happens now?

This issue contains an API change proposal (or ACP) and is part of the libs-api team feature lifecycle. Once this issue is filed, the libs-api team will review open proposals as capability becomes available. Current response times do not have a clear estimate, but may be up to several months.

Possible responses

The libs team may respond in various different ways. First, the team will consider the problem (this doesn't require any concrete solution or alternatives to have been proposed):

• We think this problem seems worth solving, and the standard library might be the right place to solve it.
• We think that this probably doesn't belong in the standard library.

Second, if there's a concrete solution:

• We think this specific solution looks roughly right, approved, you or someone else should implement this. (Further review will still happen on the subsequent implementation PR.)
• We're not sure this is the right solution, and the alternatives or other materials don't give us enough information to be sure about that. Here are some questions we have that aren't answered, or rough ideas about alternatives we'd want to see discussed.

[the8472]

Include a method in Option, that takes a callback that returns an iterator and returns an iterator for that behaves like FlatMap but is specialized for Option.

I think anything we could do for that case we could also do via specialization by moving the option into the frontiter part of Flatten at construction time instead of the iter. This way we never need to look at the other fields.

Though your solution sketch is lacking. Can you provide an api signature? And maybe show on godbolt how it's better, at least for next?

[cuviper]

Ooh, this is something I've thought about before, and I even went so far as to prototype it, including a few uses in the compiler. It turns out I still have that branch laying around, so I've pushed it here:

rust-lang/rust@master...cuviper:rust:option_flatten_iter

It's been a long while, so I'm sure that has bitrotted, but the API I added was:

impl<T: IntoIterator> Option<T> {
    /// Converts from `Option<T>` to a flattened iterator, where either
    /// `Some(x)` iterates from `x.into_iter()`, or else `None` is empty.
    ///
    /// This is conceptually similar to `option.into_iter().flatten()`, but is
    /// optimized for the fact that `option.into_iter()` only produces a single
    /// item itself. The iterator type is smaller, and it also reports a more
    /// accurate size than `Flatten`, even implementing `ExactSizeIterator` if
    /// the underlying iterator supports it.
    #[inline]
    #[unstable(feature = "option_flatten_iter", issue = "none")]
    pub fn flatten_iter(self) -> Flatten<T::IntoIter> {
        Flatten { iter: self.map(T::into_iter) }
    }
}

On the other hand, rust-lang/rust#121204 added OneShot specialization for inner iterators that only produce one item. It might be interesting to extend that to deal with the outer iterator too, but I think we would need a lattice rule for specialization, so we could cover inner OneShot, outer OneShot, or both.

I think anything we could do for that case we could also do via specialization by moving the option into the frontiter part of Flatten at construction time instead of the iter. This way we never need to look at the other fields.

That seems doable, although we should be careful about laziness when it comes to the inner IntoIterator conversion, because it might be observable to user code.

[the8472]

That seems doable, although we should be careful about laziness when it comes to the inner IntoIterator conversion, because it might be observable to user code.

We could limit it to TrustedLen or some other specialization traits because we know those should be profitable and don't have into_iter effects.

It might be interesting to extend that to deal with the outer iterator too, but I think we would need a lattice rule for specialization, so we could cover inner OneShot, outer OneShot, or both.

We can often work around the lack of lattice specialization with some tricks, it should be feasible here.

[the8472]

even implementing ExactSizeIterator if the underlying iterator supports it.

That we can do even without specialization. We already do it for the Take<RepeatWith<F>> combination.

[cuviper]

even implementing ExactSizeIterator if the underlying iterator supports it.

That we can do even without specialization. We already do it for the Take<RepeatWith<F>> combination.

Ok, yes, I guess we can do that broadly for Flatten<option::IntoIter<I>> where I: ExactSizedIterator. (also option::{Iter,IterMut}, and maybe all OneShot iterators, although we probably don't want that detail exposed.)

For something that still needs IntoIterator, like the Option<&[i32]> above iterating Flatten<option::IntoIter<&[i32]>>, we would need a specific impl for every not-yet-Iterator type we want to support.

Whereas my flatten_iter makes it trivial even for user types by front-loading the conversion.

[ChayimFriedman2]

@the8472 You can see a simple implementation along with a godbolt link in https://godbolt.org/z/MxhPWoxrE.

[the8472]

For something that still needs IntoIterator, like the Option<&[i32]> above iterating Flatten<option::IntoIter<&[i32]>>, we would need a specific impl for every not-yet-Iterator type we want to support.

Shouldn't this work for both iterators and intoiterators, since there's a blanket impl?

impl<I> ExactSizeIterator for Flatten<option::IntoIter<I>> where
    I: IntoIterator,
    <I as IntoIterator>::IntoIter: ExactSizeIterator,
    Self: Iterator
    {}

[cuviper]

That will compile, but it's wrong, because its size_hint() will return (0, None) if it hasn't advanced yet. The default ExactSizeIterator::len will then fail its assertion too.

Actually, even doing that for I: ExactSizeIterator with manual len would still have a bad size_hint, unless we specialize.

[the8472]

Right, if we want the sizehint for all we need to execute the into_iter for all.