From f7c1b9748e90ab51c1d1126135dc9a75b5f5aa1c Mon Sep 17 00:00:00 2001
From: Ibraheem Ahmed <ibraheem@ibraheem.ca>
Date: Tue, 12 Aug 2025 21:41:19 -0400
Subject: [PATCH] optimize `Vec::push`

---
 src/buckets.rs | 49 +++++++++++++++++++++++++++++--------------------
 src/vec/raw.rs | 17 ++++++++++++-----
 2 files changed, 41 insertions(+), 25 deletions(-)

diff --git a/src/buckets.rs b/src/buckets.rs
index 8e4f37f..9dab3da 100644
--- a/src/buckets.rs
+++ b/src/buckets.rs
@@ -314,35 +314,46 @@ impl<T, const BUCKETS: usize> Buckets<T, BUCKETS> {
 
         // If we are close to the end of this bucket, we eagerly allocate the next one to reduce
         // later contention.
-        if location.entry == location.bucket_len.get() / 8 * 7 {
-            if let Some(new_index) = index.after_bucket().advance() {
-                allocate_race(self.bucket(new_index), new_index.len());
-            }
-        }
-
-        if let Some(item) = self.get(index) {
-            return item;
+        if location.entry == (location.bucket_len.get() - (location.bucket_len.get() >> 3)) {
+            self.alloc_bucket_after(index);
         }
 
         let bucket = self.bucket(location.bucket);
 
-        // Panics: Let `i` be `index.get()`. To avoid panics, the condition is:
-        //
-        //   location.bucket_len * size_of::<T>() < isize::MAX + 1
-        // ⇔ 2 ^ floor(log2(i + SKIPPED_ENTRIES + 1)) < (isize::MAX + 1) / size_of::<T>()
-        // ⇔ i + SKIPPED_ENTRIES + 1 < ((isize::MAX + 1) / size_of::<T>()).next_power_of_two()
-        // ⇔ i < ((isize::MAX + 1) / size_of::<T>()).next_power_of_two() - SKIPPED_ENTRIES - 1
-        //
-        // Since `SKIPPED_ENTRIES` is an implementation detail, the caller can't enforce this.
-        // But the formula may be useful anyway.
-        let ptr = allocate_race_and_get(bucket, location.bucket_len);
+        // Acquire is necessary because we access the bucket afterward.
+        let mut ptr = bucket.load(atomic::Ordering::Acquire) as *const T;
+        if ptr.is_null() {
+            // Panics: Let `i` be `index.get()`. To avoid panics, the condition is:
+            //
+            //   location.bucket_len * size_of::<T>() < isize::MAX + 1
+            // ⇔ 2 ^ floor(log2(i + SKIPPED_ENTRIES + 1)) < (isize::MAX + 1) / size_of::<T>()
+            // ⇔ i + SKIPPED_ENTRIES + 1 < ((isize::MAX + 1) / size_of::<T>()).next_power_of_two()
+            // ⇔ i < ((isize::MAX + 1) / size_of::<T>()).next_power_of_two() - SKIPPED_ENTRIES - 1
+            //
+            // Since `SKIPPED_ENTRIES` is an implementation detail, the caller can't enforce this.
+            // But the formula may be useful anyway.
+            ptr = allocate_race_and_get(bucket, location.bucket_len);
+        }
 
         // Safety:
+        // - The pointer is non-null.
         // - By our invariants, the index is in bounds.
         // - We loaded the bucket pointer with `Acquire`, allowing us to access the allocation.
         unsafe { &*ptr.add(location.entry) }
     }
 
+    /// Eagerly allocate the bucket after the bucket containing the provided index.
+    #[cold]
+    #[inline(never)]
+    fn alloc_bucket_after(&self, index: Index<BUCKETS>)
+    where
+        T: MaybeZeroable,
+    {
+        if let Some(new_index) = index.after_bucket().advance() {
+            allocate_race(self.bucket(new_index), new_index.len());
+        }
+    }
+
     /// Retrieve a unique reference to the value at the specified index, or allocate the bucket if
     /// it hasn't been allocated yet.
     ///
@@ -522,8 +533,6 @@ fn allocate_race_and_get<T: MaybeZeroable>(bucket: &AtomicPtr<T>, len: NonZeroUs
 /// # Panics
 ///
 /// `len * size_of::<T>()` must not overflow an `isize`.
-#[cold]
-#[inline(never)]
 fn allocate_race<T: MaybeZeroable>(bucket: &AtomicPtr<T>, len: NonZeroUsize) {
     // Panics: Ensured by caller.
     let ptr = Box::into_raw(allocate_slice::<T>(len));
diff --git a/src/vec/raw.rs b/src/vec/raw.rs
index 78bb867..9efe897 100644
--- a/src/vec/raw.rs
+++ b/src/vec/raw.rs
@@ -35,6 +35,7 @@ impl<T> Vec<T> {
         let Some(zero) = <buckets::Index<BUCKETS>>::new(0) else {
             unreachable!();
         };
+
         Vec {
             inflight: AtomicUsize::new(zero.into_raw().get()),
             buckets: Buckets::new(),
@@ -145,16 +146,22 @@ impl<T> Vec<T> {
 
         // Safety: We cannot overflow.
         let Some(index) = (unsafe { buckets::Index::from_raw_checked_above(index) }) else {
-            // We could alternatively abort here, as `Arc` does. But we decrement and panic instead
-            // to keep in line with `Vec`'s behavior. Assuming that `isize::MAX` concurrent threads
-            // don't call this method, it is still impossible for it to overflow.
-            self.inflight.fetch_sub(1, Ordering::Relaxed);
-            panic!("capacity overflow");
+            self.next_index_overflow();
         };
 
         index
     }
 
+    #[cold]
+    #[inline(never)]
+    fn next_index_overflow(&self) -> ! {
+        // We could alternatively abort here, as `Arc` does. But we decrement and panic instead
+        // to keep in line with `Vec`'s behavior. Assuming that `isize::MAX` concurrent threads
+        // don't call this method, it is still impossible for it to overflow.
+        self.inflight.fetch_sub(1, Ordering::Relaxed);
+        panic!("capacity overflow");
+    }
+
     /// Appends the element returned from the closure to the back of the vector
     /// at the index represented by the `usize` passed to closure.
     ///