Optimize buffer ops

Signed-off-by: Adam Gutglick <adam@spiraldb.com>

Author: AdamGS

vortex-buffer/benches/vortex_buffer.rs

@@ -140,6 +140,42 @@ fn push_n_vortex_buffer<T: PrimInt>(bencher: Bencher, length: usize) {
         });
 }
 
+const SLICE_WINDOW: usize = 64;
+
+#[divan::bench(args = &[65_536usize])]
+fn slice_tight_loop_vortex(bencher: Bencher, len: usize) {
+    let buf = Buffer::<i32>::from_iter(
+        (0..i32::try_from(len).vortex_expect("len fits into i32")).map(|i| i % i32::MAX),
+    );
+    bencher.bench(|| {
+        let mut offset = 0;
+        while offset + SLICE_WINDOW <= len {
+            divan::black_box(buf.slice(offset..offset + SLICE_WINDOW));
+            offset += SLICE_WINDOW;
+        }
+    });
+}
+
+#[divan::bench(args = &[65_536usize])]
+fn slice_tight_loop_arrow(bencher: Bencher, len: usize) {
+    let buf = ScalarBuffer::<i32>::from_iter(
+        (0..i32::try_from(len).vortex_expect("len fits into i32")).map(|i| i % i32::MAX),
+    );
+    bencher.bench(|| {
+        let mut offset = 0;
+        while offset + SLICE_WINDOW <= len {
+            divan::black_box(buf.slice(offset, SLICE_WINDOW));
+            offset += SLICE_WINDOW;
+        }
+    });
+}
+
+#[divan::bench]
+fn slice_empty_vortex(bencher: Bencher) {
+    let buf = Buffer::<i32>::from_iter((0..1024).map(|i| i % i32::MAX));
+    bencher.bench(|| divan::black_box(buf.slice(8..8)));
+}
+
 #[divan::bench(args = INPUT_SIZE)]
 fn map_new_output(bencher: Bencher, n: i32) {
     bencher

vortex-buffer/src/alignment.rs

@@ -54,6 +54,9 @@ impl Alignment {
         Self::new(align_of::<T>())
     }
 
+    /// The largest valid alignment: the greatest power of 2 that fits into a `u16`.
+    pub const MAX: Alignment = Alignment::new(1 << 15);
+
     /// Check if `self` alignment is a "larger" than `other` alignment.
     ///
     /// ## Example
@@ -67,10 +70,32 @@ impl Alignment {
     /// assert!(!b.is_aligned_to(a));
     /// ```
     #[inline]
-    pub fn is_aligned_to(&self, other: Alignment) -> bool {
-        // Since we know alignments are powers of 2, we can compare them by checking if the number
-        // of trailing zeros in the binary representation of the alignment is greater or equal.
-        self.0.trailing_zeros() >= other.0.trailing_zeros()
+    pub const fn is_aligned_to(&self, other: Alignment) -> bool {
+        // Since both alignments are powers of 2, divisibility is equivalent to ordering.
+        self.0 >= other.0
+    }
+
+    /// Check if the given byte offset (or length) is a multiple of this alignment.
+    ///
+    /// ## Example
+    ///
+    /// ```
+    /// use vortex_buffer::Alignment;
+    ///
+    /// let a = Alignment::new(4);
+    /// assert!(a.is_offset_aligned(8));
+    /// assert!(!a.is_offset_aligned(2));
+    /// ```
+    #[inline]
+    pub const fn is_offset_aligned(&self, offset: usize) -> bool {
+        // Alignment is always a power of 2, so a mask test is equivalent to `offset % self == 0`.
+        offset & (self.0 - 1) == 0
+    }
+
+    /// Check if the given pointer is aligned to this alignment.
+    #[inline]
+    pub fn is_ptr_aligned<T>(&self, ptr: *const T) -> bool {
+        self.is_offset_aligned(ptr.addr())
     }
 
     /// Returns the log2 of the alignment.

vortex-buffer/src/buffer.rs

@@ -32,10 +32,17 @@ pub struct Buffer<T> {
     pub(crate) _marker: PhantomData<T>,
 }
 
+/// Zero-length backing memory for empty buffers, aligned to [`Alignment::MAX`] so it satisfies
+/// any valid alignment without allocating.
+#[repr(align(32768))]
+struct EmptyBacking([u8; 0]);
+
+static EMPTY_BACKING: EmptyBacking = EmptyBacking([]);
+
 impl<T> Default for Buffer<T> {
     fn default() -> Self {
         Self {
-            bytes: Default::default(),
+            bytes: Bytes::from_static(&EMPTY_BACKING.0),
             length: 0,
             alignment: Alignment::of::<T>(),
             _marker: PhantomData,
@@ -101,12 +108,27 @@ impl<T> Buffer<T> {
 
     /// Create a new empty `ByteBuffer` with the provided alignment.
     pub fn empty() -> Self {
-        BufferMut::empty().freeze()
+        Self::empty_aligned(Alignment::of::<T>())
     }
 
     /// Create a new empty `ByteBuffer` with the provided alignment.
+    ///
+    /// This does not allocate: empty buffers are backed by a shared static allocation that is
+    /// aligned to [`Alignment::MAX`].
     pub fn empty_aligned(alignment: Alignment) -> Self {
-        BufferMut::empty_aligned(alignment).freeze()
+        if !alignment.is_aligned_to(Alignment::of::<T>()) {
+            vortex_panic!(
+                "Alignment {} must align to the scalar type's alignment {}",
+                alignment,
+                Alignment::of::<T>(),
+            );
+        }
+        Self {
+            bytes: Bytes::from_static(&EMPTY_BACKING.0),
+            length: 0,
+            alignment,
+            _marker: PhantomData,
+        }
     }
 
     /// Create a new full `ByteBuffer` with the given value.
@@ -152,7 +174,7 @@ impl<T> Buffer<T> {
                 Alignment::of::<T>(),
             );
         }
-        if bytes.as_ptr().align_offset(*alignment) != 0 {
+        if !alignment.is_ptr_aligned(bytes.as_ptr()) {
             vortex_panic!(
                 "Bytes alignment must align to the requested alignment {}",
                 alignment,
@@ -320,7 +342,7 @@ impl<T> Buffer<T> {
         let begin_byte = begin * size_of::<T>();
         let end_byte = end * size_of::<T>();
 
-        if !begin_byte.is_multiple_of(*alignment) {
+        if !alignment.is_offset_aligned(begin_byte) {
             vortex_panic!(
                 "range start must be aligned to {alignment:?}, byte {}",
                 begin_byte
@@ -369,7 +391,7 @@ impl<T> Buffer<T> {
             vortex_panic!("slice_ref subset alignment must at least align to the buffer alignment")
         }
 
-        if subset.as_ptr().align_offset(*alignment) != 0 {
+        if !alignment.is_ptr_aligned(subset.as_ptr()) {
             vortex_panic!("slice_ref subset must be aligned to {:?}", alignment);
         }
 
@@ -435,17 +457,17 @@ impl<T> Buffer<T> {
     /// Convert self into `BufferMut<T>`, cloning the data if there are multiple strong references.
     pub fn into_mut(self) -> BufferMut<T> {
         self.try_into_mut()
-            .unwrap_or_else(|buffer| BufferMut::<T>::copy_from(&buffer))
+            .unwrap_or_else(|buffer| BufferMut::<T>::copy_from_aligned(&buffer, buffer.alignment))
     }
 
     /// Returns whether a `Buffer<T>` is aligned to the given alignment.
     pub fn is_aligned(&self, alignment: Alignment) -> bool {
-        self.bytes.as_ptr().align_offset(*alignment) == 0
+        alignment.is_ptr_aligned(self.bytes.as_ptr())
     }
 
     /// Return a `Buffer<T>` with the given alignment. Where possible, this will be zero-copy.
     pub fn aligned(mut self, alignment: Alignment) -> Self {
-        if self.as_ptr().align_offset(*alignment) == 0 {
+        if alignment.is_ptr_aligned(self.as_ptr()) {
             self.alignment = alignment;
             self
         } else {
@@ -462,7 +484,7 @@ impl<T> Buffer<T> {
 
     /// Return a `Buffer<T>` with the given alignment. Panics if the buffer is not aligned.
     pub fn ensure_aligned(mut self, alignment: Alignment) -> Self {
-        if self.as_ptr().align_offset(*alignment) == 0 {
+        if alignment.is_ptr_aligned(self.as_ptr()) {
             self.alignment = alignment;
             self
         } else {
@@ -634,7 +656,7 @@ impl Buf for ByteBuffer {
 
     #[inline]
     fn advance(&mut self, cnt: usize) {
-        if !cnt.is_multiple_of(*self.alignment) {
+        if !self.alignment.is_offset_aligned(cnt) {
             vortex_panic!(
                 "Cannot advance buffer by {} items, resulting alignment is not {}",
                 cnt,
@@ -786,6 +808,31 @@ mod test {
         assert_eq!(vec, buff.as_ref());
     }
 
+    #[test]
+    fn empty_aligned_max_alignment() {
+        // Empty buffers are backed by a static and must satisfy any valid alignment.
+        let buf = Buffer::<u8>::empty_aligned(Alignment::MAX);
+        assert!(buf.is_empty());
+        assert!(buf.is_aligned(Alignment::MAX));
+    }
+
+    #[test]
+    fn empty_slice_preserves_alignment() {
+        let buf = Buffer::<u64>::zeroed_aligned(8, Alignment::new(64));
+        let sliced = buf.slice(0..0);
+        assert!(sliced.is_empty());
+        assert_eq!(sliced.alignment(), Alignment::new(64));
+        assert!(sliced.is_aligned(Alignment::new(64)));
+    }
+
+    #[test]
+    fn empty_into_mut_preserves_alignment() {
+        let buf = Buffer::<u8>::empty_aligned(Alignment::new(64));
+        let buf_mut = buf.into_mut();
+        assert_eq!(buf_mut.alignment(), Alignment::new(64));
+        assert!(buf_mut.is_empty());
+    }
+
     #[test]
     fn test_slice_unaligned_end_pos() {
         let data = vec![0u8; 2];

vortex-buffer/src/buffer_mut.rs

@@ -358,7 +358,7 @@ impl<T> BufferMut<T> {
         }
 
         let bytes_at = at * size_of::<T>();
-        if !bytes_at.is_multiple_of(*self.alignment) {
+        if !self.alignment.is_offset_aligned(bytes_at) {
             vortex_panic!(
                 "Cannot split buffer at {}, resulting alignment is not {}",
                 at,
@@ -742,7 +742,7 @@ impl Buf for ByteBufferMut {
     }
 
     fn advance(&mut self, cnt: usize) {
-        if !cnt.is_multiple_of(*self.alignment) {
+        if !self.alignment.is_offset_aligned(cnt) {
             vortex_panic!(
                 "Cannot advance buffer by {} items, resulting alignment is not {}",
                 cnt,
@@ -765,7 +765,7 @@ unsafe impl BufMut for ByteBufferMut {
 
     #[inline]
     unsafe fn advance_mut(&mut self, cnt: usize) {
-        if !cnt.is_multiple_of(*self.alignment) {
+        if !self.alignment.is_offset_aligned(cnt) {
             vortex_panic!(
                 "Cannot advance buffer by {} items, resulting alignment is not {}",
                 cnt,