perf(buffer): exclude alignment slack from growth

Signed-off-by: Nicholas Gates <nick@nickgates.com>

Author: gatesn

vortex-buffer/src/buffer.rs

@@ -35,6 +35,8 @@ pub struct Buffer<T> {
     pub(crate) length: usize,
     pub(crate) alignment: Alignment,
     pub(crate) physical_alignment: Alignment,
+    // One physical-alignment block is reserved outside the logical capacity.
+    pub(crate) overallocated: bool,
     pub(crate) backing: Arc<BufferBacking>,
 }
 
@@ -67,6 +69,7 @@ impl<T> Default for Buffer<T> {
             length: 0,
             alignment: Alignment::of::<T>(),
             physical_alignment: Alignment::MAX,
+            overallocated: false,
             backing: EMPTY_BACKING.clone(),
         }
     }
@@ -109,6 +112,7 @@ impl<T> Buffer<T> {
         length: usize,
         alignment: Alignment,
         physical_alignment: Alignment,
+        overallocated: bool,
     ) -> Self {
         // SAFETY: BufferMut keeps offset within allocation, including for empty buffers.
         let ptr = unsafe { allocation.ptr().add(offset).cast() };
@@ -117,6 +121,7 @@ impl<T> Buffer<T> {
             length,
             alignment,
             physical_alignment,
+            overallocated,
             backing: Arc::new(BufferBacking::Owned(allocation)),
         }
     }
@@ -134,6 +139,7 @@ impl<T> Buffer<T> {
             length,
             alignment,
             physical_alignment: alignment,
+            overallocated: false,
             backing: Arc::new(BufferBacking::External { _owner: owner }),
         }
     }
@@ -229,6 +235,7 @@ impl<T> Buffer<T> {
             length: 0,
             alignment,
             physical_alignment: Alignment::MAX,
+            overallocated: false,
             backing: EMPTY_BACKING.clone(),
         }
     }
@@ -289,6 +296,7 @@ impl<T> Buffer<T> {
             length: buffer.length / size_of::<T>(),
             alignment,
             physical_alignment: buffer.physical_alignment,
+            overallocated: buffer.overallocated,
             backing: buffer.backing,
         }
     }
@@ -494,6 +502,7 @@ impl<T> Buffer<T> {
             length: end - begin,
             alignment,
             physical_alignment: self.physical_alignment,
+            overallocated: self.overallocated,
             backing: Arc::clone(&self.backing),
         }
     }
@@ -548,6 +557,7 @@ impl<T> Buffer<T> {
             length: subset.len(),
             alignment,
             physical_alignment: self.physical_alignment,
+            overallocated: self.overallocated,
             backing: Arc::clone(&self.backing),
         }
     }
@@ -568,6 +578,7 @@ impl<T> Buffer<T> {
             length: self.length * size_of::<T>(),
             alignment: self.alignment,
             physical_alignment: self.physical_alignment,
+            overallocated: self.overallocated,
             backing: self.backing,
         }
     }
@@ -579,17 +590,25 @@ impl<T> Buffer<T> {
             length,
             alignment,
             physical_alignment,
+            overallocated,
             backing,
         } = self;
         match Arc::try_unwrap(backing) {
             Ok(BufferBacking::Owned(allocation)) => {
                 let offset = ptr.addr().get() - allocation.ptr().addr().get();
+                let overallocated = overallocated
+                    && offset
+                        == allocation
+                            .ptr()
+                            .as_ptr()
+                            .align_offset(physical_alignment.as_usize());
                 Ok(BufferMut {
                     allocation,
                     offset,
                     length,
                     alignment,
                     physical_alignment,
+                    overallocated,
                     _marker: Default::default(),
                 })
             }
@@ -598,13 +617,15 @@ impl<T> Buffer<T> {
                 length,
                 alignment,
                 physical_alignment,
+                overallocated,
                 backing: Arc::new(backing),
             }),
             Err(backing) => Err(Self {
                 ptr,
                 length,
                 alignment,
                 physical_alignment,
+                overallocated,
                 backing,
             }),
         }
@@ -677,6 +698,7 @@ impl<T> Buffer<T> {
             length: self.length,
             alignment: self.alignment,
             physical_alignment: self.physical_alignment,
+            overallocated: self.overallocated,
             backing: self.backing,
         }
     }
@@ -782,7 +804,14 @@ where
         if std::mem::needs_drop::<T>() {
             Self::from_owner(Wrapper(value), alignment)
         } else {
-            Self::from_allocation(Allocation::from_vec(value), 0, length, alignment, alignment)
+            Self::from_allocation(
+                Allocation::from_vec(value),
+                0,
+                length,
+                alignment,
+                alignment,
+                false,
+            )
         }
     }
 }
@@ -1048,6 +1077,34 @@ mod test {
         assert_eq!(buffer.allocation.alignment(), align_of::<u32>());
     }
 
+    #[test]
+    fn from_u8_vec_preserves_capacity() {
+        let mut vec = Vec::with_capacity(16);
+        vec.extend([1u8, 2, 3]);
+
+        let buffer = Buffer::from(vec);
+        let Ok(buffer) = buffer.try_into_mut() else {
+            panic!("Vec-backed buffer should be uniquely owned")
+        };
+        assert_eq!(buffer.capacity(), 16);
+    }
+
+    #[test]
+    fn sliced_buffer_into_mut_has_safe_capacity() {
+        let mut original = crate::BufferMut::with_capacity(128);
+        original.extend(0u32..100);
+        let original = original.freeze();
+        let sliced = original.slice(64..96);
+        drop(original);
+
+        let Ok(mut sliced) = sliced.try_into_mut() else {
+            panic!("uniquely owned slice should become mutable")
+        };
+        let capacity = sliced.capacity();
+        sliced.push_n(0, capacity - sliced.len());
+        assert_eq!(sliced.len(), capacity);
+    }
+
     #[test]
     fn from_vec_preserves_drop_glue() {
         struct DropValue(Arc<AtomicUsize>);

vortex-buffer/src/buffer_mut.rs

@@ -33,6 +33,8 @@ pub struct BufferMut<T> {
     pub(crate) length: usize,
     pub(crate) alignment: Alignment,
     pub(crate) physical_alignment: Alignment,
+    // One physical-alignment block is reserved outside the logical capacity.
+    pub(crate) overallocated: bool,
     pub(crate) _marker: std::marker::PhantomData<T>,
 }
 
@@ -134,6 +136,7 @@ impl<T> BufferMut<T> {
             length: 0,
             alignment,
             physical_alignment: actual,
+            overallocated: true,
             _marker: Default::default(),
         }
     }
@@ -222,6 +225,7 @@ impl<T> BufferMut<T> {
             length: len,
             alignment,
             physical_alignment: actual_alignment,
+            overallocated: true,
             _marker: Default::default(),
         }
     }
@@ -387,7 +391,15 @@ impl<T> BufferMut<T> {
     /// Returns the capacity of the buffer.
     #[inline]
     pub fn capacity(&self) -> usize {
-        (self.allocation.size() - self.offset) / size_of::<T>()
+        if self.allocation.size() == 0 {
+            return 0;
+        }
+
+        if !self.overallocated {
+            return (self.allocation.size() - self.offset) / size_of::<T>();
+        }
+
+        (self.allocation.size() - self.physical_alignment.as_usize()) / size_of::<T>()
     }
 
     /// Returns a raw pointer to the buffer's data.
@@ -461,7 +473,10 @@ impl<T> BufferMut<T> {
             .checked_mul(size_of::<T>())
             .vortex_expect("buffer capacity overflow");
         let physical_alignment = max(self.alignment, self.physical_alignment);
-        let current_size = self.allocation.size() - self.offset;
+        let current_size = self
+            .capacity()
+            .checked_mul(size_of::<T>())
+            .vortex_expect("buffer capacity overflow");
         let logical_size = required_size
             .max(current_size.saturating_mul(2))
             .max(physical_alignment.as_usize());
@@ -517,6 +532,7 @@ impl<T> BufferMut<T> {
         };
         self.offset = new_offset;
         self.physical_alignment = physical_alignment;
+        self.overallocated = true;
     }
 
     /// Returns the spare capacity of the buffer as a slice of `MaybeUninit<T>`.
@@ -666,6 +682,7 @@ impl<T> BufferMut<T> {
             length: self.length * size_of::<T>(),
             alignment: self.alignment,
             physical_alignment: self.physical_alignment,
+            overallocated: self.overallocated,
             _marker: Default::default(),
         }
     }
@@ -678,6 +695,7 @@ impl<T> BufferMut<T> {
             self.length,
             self.alignment,
             self.physical_alignment,
+            self.overallocated,
         )
     }
 
@@ -742,6 +760,7 @@ impl<T> BufferMut<T> {
             length: self.length,
             alignment: self.alignment,
             physical_alignment: self.physical_alignment,
+            overallocated: self.overallocated,
             _marker: std::marker::PhantomData,
         }
     }
@@ -1093,10 +1112,10 @@ mod test {
 
         buffer.push(0);
         let capacity = buffer.capacity();
-        assert!(capacity >= alignment.as_usize());
+        assert_eq!(capacity, Alignment::DEFAULT_ALIGNMENT.as_usize());
 
         buffer.reserve(capacity);
-        assert!(buffer.capacity() >= capacity * 2);
+        assert_eq!(buffer.capacity(), capacity * 2);
     }
 
     #[test]