fix(buffer): reject ZSTs and address allocator review comments

vortex-buffer/src/allocation.rs

@@ -4,6 +4,7 @@
 //! Allocator-backed storage for Vortex buffers.
 
 use std::alloc::Layout;
+use std::any::Any;
 use std::fmt;
 use std::fmt::Debug;
 use std::mem::ManuallyDrop;
@@ -27,8 +28,15 @@ pub trait BufferAllocator: Allocator + Debug + Send + Sync + 'static {}
 impl<A> BufferAllocator for A where A: Allocator + Debug + Send + Sync + 'static {}
 
 /// A shared reference to a buffer allocator.
+///
+/// The static allocator does not need shared ownership, so it is stored without an [`Arc`]. This
+/// makes cloning the common static allocator a simple value copy.
 #[derive(Clone)]
-pub struct BufferAllocatorRef(Option<Arc<dyn BufferAllocator>>);
+pub struct BufferAllocatorRef(
+    // `None` selects the static allocator without allocating or updating an Arc reference count.
+    // `Some` keeps a custom allocator alive for as long as its buffers need it.
+    Option<Arc<dyn BufferAllocator>>,
+);
 
 impl BufferAllocatorRef {
     /// Wrap an allocator in a shared reference.
@@ -333,32 +341,13 @@ impl Drop for Allocation {
     }
 }
 
-pub(crate) trait BufferOwner: Send + Sync + 'static {
-    fn as_ptr(&self) -> *const u8;
-
-    fn len(&self) -> usize;
-}
-
-impl<T> BufferOwner for T
-where
-    T: AsRef<[u8]> + Send + Sync + 'static,
-{
-    fn as_ptr(&self) -> *const u8 {
-        self.as_ref().as_ptr()
-    }
-
-    fn len(&self) -> usize {
-        self.as_ref().len()
-    }
-}
-
 pub(crate) enum BufferBacking {
     Owned(Allocation),
     Bytes(bytes::Bytes),
     #[cfg(feature = "arrow")]
     Arrow(arrow_buffer::Buffer),
     External {
-        _owner: Box<dyn BufferOwner>,
+        _owner: Box<dyn Any + Send + Sync>,
     },
 }
 
@@ -386,9 +375,13 @@ mod tests {
     use allocator_api2::alloc::AllocError;
     use allocator_api2::alloc::Allocator;
     use allocator_api2::alloc::Global;
+    use rstest::rstest;
+    use vortex_error::VortexResult;
+    use vortex_error::vortex_err;
 
     use crate::Alignment;
     use crate::BufferAllocatorRef;
+    use crate::BufferMut;
 
     #[derive(Clone, Debug, Default)]
     struct TrackingAllocator {
@@ -462,15 +455,18 @@ mod tests {
         assert_eq!(state.deallocations.load(Ordering::Relaxed), 1);
     }
 
-    #[test]
-    fn buffer_growth_uses_allocator_grow() {
+    #[rstest]
+    fn buffer_growth_uses_allocator_grow(#[values(4, 64, 4096)] alignment: usize) {
         let allocator = TrackingAllocator::default();
         let state = Arc::clone(&allocator.state);
-        let mut buffer = BufferAllocatorRef::new(allocator).with_capacity::<u32>(1);
+        let alignment = Alignment::new(alignment);
+        let mut buffer =
+            BufferAllocatorRef::new(allocator).with_capacity_aligned::<u32>(1, alignment);
         let initial_capacity = buffer.capacity();
         buffer.extend(std::iter::repeat_n(7, initial_capacity));
 
         buffer.push(u32::MAX);
+        assert!(alignment.is_ptr_aligned(buffer.as_ptr()));
 
         assert_eq!(&buffer[..initial_capacity], vec![7; initial_capacity]);
         assert_eq!(buffer[initial_capacity], u32::MAX);
@@ -498,4 +494,44 @@ mod tests {
         assert_eq!(state.allocations.load(Ordering::Relaxed), 1);
         assert_eq!(state.grows.load(Ordering::Relaxed), 0);
     }
+
+    #[test]
+    fn empty_buffers_preserve_allocator_without_allocating() -> VortexResult<()> {
+        let allocator = TrackingAllocator::default();
+        let state = Arc::clone(&allocator.state);
+        let allocator = BufferAllocatorRef::new(allocator);
+        let buffer = BufferMut::<u32>::zeroed_in(0, allocator.clone());
+        let buffer = buffer.freeze();
+        let copy = buffer.clone().into_mut();
+        assert!(copy.allocator().ptr_eq(&allocator));
+        let mut buffer = buffer
+            .try_into_mut()
+            .map_err(|_| vortex_err!("unique buffer"))?;
+        buffer.reserve(0);
+        assert!(buffer.is_empty());
+        assert!(buffer.allocator().ptr_eq(&allocator));
+        drop((copy, buffer));
+        assert_eq!(state.allocations.load(Ordering::Relaxed), 0);
+        assert_eq!(state.grows.load(Ordering::Relaxed), 0);
+        assert_eq!(state.deallocations.load(Ordering::Relaxed), 0);
+        Ok(())
+    }
+
+    #[test]
+    fn shared_into_mut_preserves_allocator() {
+        let allocator = TrackingAllocator::default();
+        let state = Arc::clone(&allocator.state);
+        let allocator = BufferAllocatorRef::new(allocator);
+        let original = allocator.copy_from([1u32, 2, 3]).freeze();
+        let mut copy = original.clone().into_mut();
+        assert!(copy.allocator().ptr_eq(&allocator));
+        copy[0] = 42;
+        assert_eq!(original.as_slice(), [1, 2, 3]);
+        assert_eq!(copy.as_slice(), [42, 2, 3]);
+        assert_eq!(state.allocations.load(Ordering::Relaxed), 2);
+        drop(copy);
+        assert_eq!(state.deallocations.load(Ordering::Relaxed), 1);
+        drop(original);
+        assert_eq!(state.deallocations.load(Ordering::Relaxed), 2);
+    }
 }

vortex-buffer/src/buffer.rs

@@ -28,11 +28,38 @@ use crate::debug::TruncatedDebug;
 use crate::trusted_len::TrustedLen;
 
 /// An immutable buffer of items of `T`.
+///
+/// Zero-sized element types are rejected at compile time when constructing a buffer.
+///
+/// ```compile_fail
+/// use vortex_buffer::Buffer;
+/// let _ = Buffer::<()>::empty();
+/// ```
+///
+/// ```compile_fail
+/// use vortex_buffer::Buffer;
+/// let _ = Buffer::from(vec![(); 3]);
+/// ```
+///
+/// ```compile_fail
+/// use vortex_buffer::{Buffer, ByteBuffer};
+/// let _ = Buffer::<()>::from_byte_buffer(ByteBuffer::empty());
+/// ```
+///
+/// ```compile_fail
+/// use bytes::Bytes;
+/// use vortex_buffer::{Alignment, Buffer};
+/// let _ = Buffer::<()>::from_bytes_aligned(Bytes::new(), Alignment::none());
+/// ```
 #[derive(Clone)]
 pub struct Buffer<T> {
+    /// The first element in this view; may dangle for an empty buffer.
     pub(crate) ptr: NonNull<T>,
+    /// The number of initialized `T` values visible from `ptr`.
     pub(crate) length: usize,
+    /// The minimum alignment promised for `ptr` and preserved by aligned slices.
     pub(crate) alignment: Alignment,
+    /// Shared ownership of the storage containing `ptr`, if any; `Buffer::empty` has no backing.
     pub(crate) backing: Option<Arc<BufferBacking>>,
 }
 
@@ -100,22 +127,6 @@ impl<T> Buffer<T> {
         }
     }
 
-    fn from_owner(owner: impl crate::BufferOwner, alignment: Alignment) -> Self {
-        let owner: Box<dyn crate::BufferOwner> = Box::new(owner);
-        let length = owner.len() / size_of::<T>();
-        let ptr = if length == 0 {
-            empty_ptr()
-        } else {
-            NonNull::new(owner.as_ptr().cast_mut().cast()).vortex_expect("owner pointer is null")
-        };
-        Self {
-            ptr,
-            length,
-            alignment,
-            backing: Some(Arc::new(BufferBacking::External { _owner: owner })),
-        }
-    }
-
     fn from_bytes(bytes: Bytes, alignment: Alignment) -> Self {
         let length = bytes.len() / size_of::<T>();
         if length == 0 {
@@ -228,6 +239,7 @@ impl<T> Buffer<T> {
     ///
     /// This does not allocate. Empty buffers use an aligned dangling pointer.
     pub fn empty_aligned(alignment: Alignment) -> Self {
+        const { assert!(size_of::<T>() != 0, "ZSTs are not supported") };
         if !alignment.is_aligned_to(Alignment::of::<T>()) {
             vortex_panic!(
                 "Alignment {} must align to the scalar type's alignment {}",
@@ -277,6 +289,7 @@ impl<T> Buffer<T> {
     /// Panics if the buffer is not aligned to the given alignment, if the length is not a multiple
     /// of the size of `T`, or if the given alignment is not aligned to that of `T`.
     pub fn from_byte_buffer_aligned(buffer: ByteBuffer, alignment: Alignment) -> Self {
+        const { assert!(size_of::<T>() != 0, "ZSTs are not supported") };
         if !alignment.is_aligned_to(Alignment::of::<T>()) {
             vortex_panic!(
                 "Alignment {} must be compatible with the scalar type's alignment {}",
@@ -309,6 +322,7 @@ impl<T> Buffer<T> {
     /// Panics if the buffer is not aligned to the size of `T`, or the length is not a multiple of
     /// the size of `T`.
     pub fn from_bytes_aligned(bytes: Bytes, alignment: Alignment) -> Self {
+        const { assert!(size_of::<T>() != 0, "ZSTs are not supported") };
         if !alignment.is_aligned_to(Alignment::of::<T>()) {
             vortex_panic!(
                 "Alignment {} must be compatible with the scalar type's alignment {}",
@@ -802,29 +816,25 @@ impl<T> FromIterator<T> for Buffer<T> {
     }
 }
 
-// Helper struct that preserves drop glue for non-native Vec elements.
-#[repr(transparent)]
-struct Wrapper<T>(Vec<T>);
-
-impl<T: Send + Sync + 'static> crate::BufferOwner for Wrapper<T> {
-    fn as_ptr(&self) -> *const u8 {
-        self.0.as_ptr().cast()
-    }
-
-    fn len(&self) -> usize {
-        self.0.len() * size_of::<T>()
-    }
-}
-
 impl<T> From<Vec<T>> for Buffer<T>
 where
     T: Send + Sync + 'static,
 {
     fn from(value: Vec<T>) -> Self {
+        const { assert!(size_of::<T>() != 0, "ZSTs are not supported") };
         let length = value.len();
         let alignment = Alignment::of::<T>();
         if std::mem::needs_drop::<T>() {
-            Self::from_owner(Wrapper(value), alignment)
+            // Keep the typed owner so its elements are dropped.
+            Self {
+                ptr: NonNull::new(value.as_ptr().cast_mut())
+                    .vortex_expect("a Vec always has a non-null pointer"),
+                length,
+                alignment,
+                backing: Some(Arc::new(BufferBacking::External {
+                    _owner: Box::new(value),
+                })),
+            }
         } else {
             Self::from_allocation(Allocation::from_vec(value), 0, length, alignment)
         }
@@ -1094,6 +1104,42 @@ mod test {
         assert_eq!(buffer.allocation.alignment(), align_of::<u32>());
     }
 
+    #[test]
+    fn byte_owner_preserves_slice_and_lifetime() {
+        struct Owner {
+            values: Vec<u8>,
+            drops: Arc<AtomicUsize>,
+        }
+
+        impl AsRef<[u8]> for Owner {
+            fn as_ref(&self) -> &[u8] {
+                &self.values[1..4]
+            }
+        }
+
+        impl Drop for Owner {
+            fn drop(&mut self) {
+                self.drops.fetch_add(1, Ordering::Relaxed);
+            }
+        }
+
+        let drops = Arc::new(AtomicUsize::new(0));
+        let owner = Owner {
+            values: vec![0, 1, 2, 3, 4],
+            drops: Arc::clone(&drops),
+        };
+        let ptr = owner.as_ref().as_ptr();
+        let buffer = ByteBuffer::from(Bytes::from_owner(owner));
+        assert_eq!(buffer.as_ptr(), ptr);
+        assert_eq!(buffer.as_slice(), [1, 2, 3]);
+        let view = buffer.slice(1..);
+        drop(buffer);
+        assert_eq!(drops.load(Ordering::Relaxed), 0);
+        assert_eq!(view.as_slice(), [2, 3]);
+        drop(view);
+        assert_eq!(drops.load(Ordering::Relaxed), 1);
+    }
+
     #[test]
     fn bytes_round_trip_reuses_owner() {
         let bytes = Bytes::from_static(&[1, 2, 3, 4]);
@@ -1151,9 +1197,15 @@ mod test {
         let Ok(mut sliced) = sliced.try_into_mut() else {
             panic!("uniquely owned slice should become mutable")
         };
+        let ptr = sliced.as_ptr();
         let capacity = sliced.capacity();
         sliced.push_n(0, capacity - sliced.len());
         assert_eq!(sliced.len(), capacity);
+        assert_eq!(sliced.as_ptr(), ptr);
+        sliced.push(42);
+        assert_eq!(&sliced[..32], (64u32..96).collect::<Vec<_>>());
+        assert_eq!(&sliced[32..capacity], vec![0; capacity - 32]);
+        assert_eq!(sliced[capacity], 42);
     }
 
     #[test]

vortex-buffer/src/buffer_mut.rs

@@ -23,12 +23,30 @@ use crate::debug::TruncatedDebug;
 use crate::trusted_len::TrustedLen;
 
 /// A mutable buffer that maintains a runtime-defined alignment through resizing operations.
+///
+/// Zero-sized element types are rejected at compile time when constructing a buffer.
+///
+/// ```compile_fail
+/// use vortex_buffer::BufferMut;
+/// let _ = BufferMut::<()>::empty();
+/// ```
+///
+/// ```compile_fail
+/// use vortex_buffer::BufferMut;
+/// let _ = BufferMut::<()>::zeroed(3);
+/// ```
 pub struct BufferMut<T> {
+    /// The owned allocation, including any bytes before `ptr` used for alignment.
     pub(crate) allocation: Allocation,
+    /// The first element, aligned to `alignment`; it may dangle for an empty buffer.
     pub(crate) ptr: std::ptr::NonNull<T>,
+    /// The number of initialized `T` values starting at `ptr`.
     pub(crate) length: usize,
+    /// The number of `T` values that fit from `ptr`.
     pub(crate) capacity: usize,
+    /// The minimum alignment maintained for `ptr` across reallocations.
     pub(crate) alignment: Alignment,
+    /// Marks the buffer as logically owning values of `T` despite storing an erased allocation.
     pub(crate) _marker: std::marker::PhantomData<T>,
 }
 
@@ -100,6 +118,7 @@ impl<T> BufferMut<T> {
         preferred_alignment: Option<Alignment>,
         allocator: BufferAllocatorRef,
     ) -> Self {
+        const { assert!(size_of::<T>() != 0, "ZSTs are not supported") };
         let actual = max(
             alignment,
             preferred_alignment.unwrap_or(Alignment::of::<u8>()),
@@ -131,11 +150,7 @@ impl<T> BufferMut<T> {
         let offset = allocation.ptr().as_ptr().align_offset(actual.as_usize());
         // SAFETY: the allocation includes enough padding to reach this aligned pointer.
         let ptr = unsafe { allocation.ptr().add(offset).cast() };
-        let capacity = if size_of::<T>() == 0 {
-            capacity
-        } else {
-            (allocation.size() - offset) / size_of::<T>()
-        };
+        let capacity = (allocation.size() - offset) / size_of::<T>();
         Self {
             allocation,
             ptr,
@@ -204,6 +219,7 @@ impl<T> BufferMut<T> {
         preferred_alignment: Option<Alignment>,
         allocator: BufferAllocatorRef,
     ) -> Self {
+        const { assert!(size_of::<T>() != 0, "ZSTs are not supported") };
         let preferred_alignment = preferred_alignment.unwrap_or(Alignment::of::<u8>());
         let actual_alignment = max(preferred_alignment, alignment);
         let size = len
@@ -226,11 +242,7 @@ impl<T> BufferMut<T> {
             .align_offset(actual_alignment.as_usize());
         // SAFETY: the allocation includes enough padding to reach this aligned pointer.
         let ptr = unsafe { allocation.ptr().add(offset).cast() };
-        let capacity = if size_of::<T>() == 0 {
-            len
-        } else {
-            (allocation.size() - offset) / size_of::<T>()
-        };
+        let capacity = (allocation.size() - offset) / size_of::<T>();
         Self {
             allocation,
             ptr,