Collecting into a Result<Vec<_>> doesn't reserve the capacity in advance

[scurest]

Couldn't find an issue for this and don't know if it counts but filing anyway.

---

If you have

fn foo(s: &[i8]) -> Vec<u8> {
    s.iter()
    .map(|&x| x as u8)
    .collect()
}

the SpecExtend machinery ensures that the vector has s.len() space reserved in advance. However if you change it to return a result

fn foo(s: &[i8]) -> Result<Vec<u8>> {
    s.iter()
    .map(|&x| if x < 0 { Err(...) } else { Ok(x as u8) } )
    .collect()
}

then (based on examining the LLVM IR and heaptracker's "Temporary" measurements) that optimization has quietly been lost.

This is technically correct in the sense that the first element yielded could be an Err of course (the size hint for the Adapter in Result's FromIterator impl has a lower bound of 0). But this pessimizes the good path to take more memory and be slower in favor of possibly saving memory on the bad one, which seems backwards.

Is there a specialization that could be added to fix this?

[scottmcm]

I wonder if there's a more general tweak. It looks like collect only looks at the low end of the hint (https://doc.rust-lang.org/src/alloc/vec.rs.html#1804), so it feels like there ought to be a way to get some use out of the high end as well.

Strawman proposal:

• Do the same as today if high is None
  • Mostly to avoid trying to do anything with the default (0, None) "hint"
• Otherwise reserve a rough approximation of √(low × high)
  • So (0, 1GiB) pre-reserves 32KiB, for example
  • Very cheap with leading_zeros and shifting
  • Hopefully never so big as to be a problem (and could try_reserve to never be, I suppose)
• At the end, shrink_to_fit if len() <= cap()/2
  • So if there really is only a handful of elements, it doesn't waste too much space

Probably needs the benchmark set to decide whether it's constructive, though...

[JulianKnodt]

I don't know if this issue is dead or not, but I was wondering if this could be solved by creating a newtype that optimistically allocates capacity for the vec, such as this?

[dtolnay]

This is not closed by #92138. That PR involves a totally different impl.

[dylni]

Thanks @dtolnay!

[TennyZhuang]

Paste my potential solution from
https://internals.rust-lang.org/t/an-optimistic-version-of-collect-result-t-e/17584/4, do we need an RFC for that?

I’d prefer to implement two versions, collect_optimistic and collect_pessimistic, which will use size_hint.0 and size_hint.1, and leave collect as an opaque implementation (even leave for PGO?). When in the critical path, users can choose the version manually.

The collect can use a specialized implementation for Result<I, E> that forward to optimistic version, and keep the original behavior for other types.

We can also introduce a new unsafe trait TrustedUpperBound for optimization.

[scottmcm]

My instinct is that the code doing the collect isn't the one that knows whether to call optimistic or pessimistic. It feels to me like filter_optimistic or filter_pessimistic is the place that's more likely to know whether it'd be better to use a small or large reservation in a later collect. And that way it can be passed to other things too.

[stepancheg]

Can someone advice how it can be implemented in Rust?

There are two possibilities:

Easy: add expected_size_lower_bound to Iterator trait

trait Iterator {
  // unsable, doc-hidden etc
  fn expected_size_lower_bound(&self) -> usize {
    self.size_hint().0
  }
}

And then patch GenericShunt Iterator implementation to override estimated_capacity_lower_bound:

rust/library/core/src/iter/adapters/mod.rs

Line 148 in c528357

pub(crate) struct GenericShunt<'a, I, R> {

And finally, patch Vec::from_iter to use expected_size_lower_bound instead of size_hint to reserve capacity.

Hard: use some specialization

Introduce a trait like:

trait ExpectedSize: Iterator {
  fn expected_size_lower_bound(&self) -> usize;
}

And then implement it for GenericShunt.

My specialization-fu is not good enough to implement a function like this:

trait ExpectedSize {
  fn expected_size_lower_bound(&self) -> usize;
}

impl<I: Iterator> ExpectedSize for I {
  fn expected_size_lower_bound(&self) -> usize {
    self.size_hint().0
  }
}

// And here I get compilation errors.
impl<I, R> ExpectedSize for GenericShunt<I, R> { ... }

Full example on playground, where I get either "conflicting implementations of trait" or "cannot specialize on trait" error depending on whether I add I: Iterator constraint on implementation.