Rust: 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?