From 76f763c2e9673afb399876c0dee532ce12a5c9ad Mon Sep 17 00:00:00 2001
From: Scott McMurray <scottmcm@users.noreply.github.com>
Date: Sun, 2 Mar 2025 18:46:13 -0800
Subject: [PATCH] Attempt to use the high part of the `size_hint` in `collect`

---
 library/alloc/src/lib.rs                      |  1 +
 .../alloc/src/vec/spec_from_iter_nested.rs    | 46 +++++++++++++------
 .../vec-collect-reserve-upper-bound.rs        | 16 +++++++
 3 files changed, 49 insertions(+), 14 deletions(-)
 create mode 100644 tests/codegen/vec-collect-reserve-upper-bound.rs

diff --git a/library/alloc/src/lib.rs b/library/alloc/src/lib.rs
index 2e9dd98571537..32a5188ac3c21 100644
--- a/library/alloc/src/lib.rs
+++ b/library/alloc/src/lib.rs
@@ -173,6 +173,7 @@
 #![feature(hashmap_internals)]
 #![feature(intrinsics)]
 #![feature(lang_items)]
+#![feature(let_chains)]
 #![feature(min_specialization)]
 #![feature(multiple_supertrait_upcastable)]
 #![feature(negative_impls)]
diff --git a/library/alloc/src/vec/spec_from_iter_nested.rs b/library/alloc/src/vec/spec_from_iter_nested.rs
index 22eed238798cf..d404fd9da42d2 100644
--- a/library/alloc/src/vec/spec_from_iter_nested.rs
+++ b/library/alloc/src/vec/spec_from_iter_nested.rs
@@ -22,21 +22,39 @@ where
         // empty, but the loop in extend_desugared() is not going to see the
         // vector being full in the few subsequent loop iterations.
         // So we get better branch prediction.
-        let mut vector = match iterator.next() {
-            None => return Vec::new(),
-            Some(element) => {
-                let (lower, _) = iterator.size_hint();
-                let initial_capacity =
-                    cmp::max(RawVec::<T>::MIN_NON_ZERO_CAP, lower.saturating_add(1));
-                let mut vector = Vec::with_capacity(initial_capacity);
-                unsafe {
-                    // SAFETY: We requested capacity at least 1
-                    ptr::write(vector.as_mut_ptr(), element);
-                    vector.set_len(1);
-                }
-                vector
-            }
+        let (low, high) = iterator.size_hint();
+        let Some(first) = iterator.next() else {
+            return Vec::new();
         };
+        // `push`'s growth strategy is (currently) to double the capacity if
+        // there's no space available, so it can have up to 50% "wasted" space.
+        // Thus if the upper-bound on the size_hint also wouldn't waste more
+        // than that, just allocate it from the start. (After all, it's silly
+        // to allocate 254 for a hint of `(254, Some(255)`.)
+        let initial_capacity = {
+            // This is written like this to not overflow on any well-behaved iterator,
+            // even things like `repeat_n(val, isize::MAX as usize + 10)`
+            // where `low * 2` would need checking.
+            // A bad (but safe) iterator might have `low > high`, but if so it'll
+            // produce a huge `extra` that'll probably fail the following check.
+            let hint = if let Some(high) = high
+                && let extra = high - low
+                && extra < low
+            {
+                high
+            } else {
+                low
+            };
+            cmp::max(RawVec::<T>::MIN_NON_ZERO_CAP, hint)
+        };
+        let mut vector = Vec::with_capacity(initial_capacity);
+        // SAFETY: We requested capacity at least MIN_NON_ZERO_CAP, which
+        // is never zero, so there's space for at least one element.
+        unsafe {
+            ptr::write(vector.as_mut_ptr(), first);
+            vector.set_len(1);
+        }
+
         // must delegate to spec_extend() since extend() itself delegates
         // to spec_from for empty Vecs
         <Vec<T> as SpecExtend<T, I>>::spec_extend(&mut vector, iterator);
diff --git a/tests/codegen/vec-collect-reserve-upper-bound.rs b/tests/codegen/vec-collect-reserve-upper-bound.rs
new file mode 100644
index 0000000000000..778704fdfb98b
--- /dev/null
+++ b/tests/codegen/vec-collect-reserve-upper-bound.rs
@@ -0,0 +1,16 @@
+//@ compile-flags: -Copt-level=3
+#![crate_type = "lib"]
+
+#[no_mangle]
+pub fn should_use_low(a: [i32; 10], b: [i32; 100], p: fn(i32) -> bool) -> Vec<i32> {
+    // CHECK-LABEL: define void @should_use_low
+    // CHECK: call{{.+}}dereferenceable_or_null(40){{.+}}@__rust_alloc(
+    a.iter().copied().chain(b.iter().copied().filter(|x| p(*x))).collect()
+}
+
+#[no_mangle]
+pub fn should_use_high(a: [i32; 100], b: [i32; 10], p: fn(i32) -> bool) -> Vec<i32> {
+    // CHECK-LABEL: define void @should_use_high
+    // CHECK: call{{.+}}dereferenceable_or_null(440){{.+}}@__rust_alloc(
+    a.iter().copied().chain(b.iter().copied().filter(|x| p(*x))).collect()
+}