Skip to main content

vortex_buffer/
buffer.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::any::type_name;
5use std::cmp::Ordering;
6use std::collections::Bound;
7use std::fmt::Debug;
8use std::fmt::Formatter;
9use std::hash::Hash;
10use std::hash::Hasher;
11use std::ops::Deref;
12use std::ops::RangeBounds;
13use std::ptr::NonNull;
14use std::sync::Arc;
15
16use bytes::Buf;
17use bytes::Bytes;
18use vortex_error::VortexExpect;
19use vortex_error::vortex_panic;
20
21use crate::Alignment;
22use crate::Allocation;
23use crate::BufferAllocatorRef;
24use crate::BufferBacking;
25use crate::BufferMut;
26use crate::ByteBuffer;
27use crate::debug::TruncatedDebug;
28use crate::trusted_len::TrustedLen;
29
30/// An immutable buffer of items of `T`.
31///
32/// Zero-sized element types are rejected at compile time when constructing a buffer.
33///
34/// ```compile_fail
35/// use vortex_buffer::Buffer;
36/// let _ = Buffer::<()>::empty();
37/// ```
38///
39/// ```compile_fail
40/// use vortex_buffer::Buffer;
41/// let _ = Buffer::from(vec![(); 3]);
42/// ```
43///
44/// ```compile_fail
45/// use vortex_buffer::{Buffer, ByteBuffer};
46/// let _ = Buffer::<()>::from_byte_buffer(ByteBuffer::empty());
47/// ```
48///
49/// ```compile_fail
50/// use bytes::Bytes;
51/// use vortex_buffer::{Alignment, Buffer};
52/// let _ = Buffer::<()>::from_bytes_aligned(Bytes::new(), Alignment::none());
53/// ```
54#[derive(Clone)]
55pub struct Buffer<T> {
56    /// The first element in this view; may dangle for an empty buffer.
57    pub(crate) ptr: NonNull<T>,
58    /// The number of initialized `T` values visible from `ptr`.
59    pub(crate) length: usize,
60    /// The minimum alignment promised for `ptr` and preserved by aligned slices.
61    pub(crate) alignment: Alignment,
62    /// Shared ownership of the storage containing `ptr`, if any; `Buffer::empty` has no backing.
63    pub(crate) backing: Option<Arc<BufferBacking>>,
64}
65
66// SAFETY: Buffer is an immutable view over backing memory. Its pointer remains valid while the
67// backing is live, and sharing elements follows the same bounds as sharing a slice.
68unsafe impl<T: Send> Send for Buffer<T> {}
69// SAFETY: see the Send implementation above.
70unsafe impl<T: Sync> Sync for Buffer<T> {}
71
72impl<T> Default for Buffer<T> {
73    fn default() -> Self {
74        Self {
75            ptr: empty_ptr(),
76            length: 0,
77            alignment: Alignment::of::<T>(),
78            backing: None,
79        }
80    }
81}
82
83impl<T: PartialEq> PartialEq for Buffer<T> {
84    #[inline]
85    fn eq(&self, other: &Self) -> bool {
86        self.as_slice() == other.as_slice()
87    }
88}
89
90impl<T: Eq> Eq for Buffer<T> {}
91
92impl<T: Ord> Ord for Buffer<T> {
93    #[inline]
94    fn cmp(&self, other: &Self) -> Ordering {
95        self.as_slice().cmp(other.as_slice())
96    }
97}
98
99impl<T: PartialOrd> PartialOrd for Buffer<T> {
100    #[inline]
101    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
102        self.as_slice().partial_cmp(other.as_slice())
103    }
104}
105
106impl<T: Hash> Hash for Buffer<T> {
107    #[inline]
108    fn hash<H: Hasher>(&self, state: &mut H) {
109        self.as_slice().hash(state)
110    }
111}
112
113impl<T> Buffer<T> {
114    pub(crate) fn from_allocation(
115        allocation: Allocation,
116        offset: usize,
117        length: usize,
118        alignment: Alignment,
119    ) -> Self {
120        // SAFETY: BufferMut keeps offset within allocation, including for empty buffers.
121        let ptr = unsafe { allocation.ptr().add(offset).cast() };
122        Self {
123            ptr,
124            length,
125            alignment,
126            backing: Some(Arc::new(BufferBacking::Owned(allocation))),
127        }
128    }
129
130    fn from_bytes(bytes: Bytes, alignment: Alignment) -> Self {
131        let length = bytes.len() / size_of::<T>();
132        if length == 0 {
133            return Self::empty_aligned(alignment);
134        }
135        let ptr =
136            NonNull::new(bytes.as_ptr().cast_mut().cast()).vortex_expect("Bytes pointer is null");
137        Self {
138            ptr,
139            length,
140            alignment,
141            backing: Some(Arc::new(BufferBacking::Bytes(bytes))),
142        }
143    }
144
145    #[cfg(feature = "arrow")]
146    pub(crate) fn from_arrow_owner(
147        arrow: arrow_buffer::Buffer,
148        length: usize,
149        alignment: Alignment,
150    ) -> Self {
151        if length == 0 {
152            return Self::empty_aligned(alignment);
153        }
154        let ptr = NonNull::new(arrow.as_ptr().cast_mut().cast())
155            .vortex_expect("Arrow buffer pointer is null");
156        Self {
157            ptr,
158            length,
159            alignment,
160            backing: Some(Arc::new(BufferBacking::Arrow(arrow))),
161        }
162    }
163
164    /// Returns a new `Buffer<T>` copied from the provided `Vec<T>`, `&[T]`, etc.
165    ///
166    /// Due to our underlying usage of `bytes::Bytes`, we are unable to take zero-copy ownership
167    /// of the provided `Vec<T>` while maintaining the ability to convert it back into a mutable
168    /// buffer. We could fix this by forking `Bytes`, or in many other complex ways, but for now
169    /// callers should prefer to construct `Buffer<T>` from a `BufferMut<T>`.
170    pub fn copy_from(values: impl AsRef<[T]>) -> Self
171    where
172        T: Copy,
173    {
174        BufferMut::copy_from(values).freeze()
175    }
176
177    /// Returns a new `Buffer<T>` copied with the provided allocator.
178    pub fn copy_from_in(values: impl AsRef<[T]>, allocator: BufferAllocatorRef) -> Self
179    where
180        T: Copy,
181    {
182        BufferMut::copy_from_in(values, allocator).freeze()
183    }
184
185    /// Returns a new `Buffer<T>` copied from the provided slice and with the requested alignment.
186    ///
187    /// The allocation is over-aligned to [`Alignment::DEFAULT_ALIGNMENT`] when that is larger than
188    /// `alignment`. Use [`copy_from_preferred_aligned`] to control the over-alignment.
189    ///
190    /// [`copy_from_preferred_aligned`]: Self::copy_from_preferred_aligned
191    pub fn copy_from_aligned(values: impl AsRef<[T]>, alignment: Alignment) -> Self
192    where
193        T: Copy,
194    {
195        Self::copy_from_preferred_aligned(values, alignment, Some(Alignment::DEFAULT_ALIGNMENT))
196    }
197
198    /// Returns a new `Buffer<T>` copied from the provided slice and with the requested alignment.
199    ///
200    /// The buffer reports `alignment`, but the underlying allocation is over-aligned to the larger
201    /// of `alignment` and `preferred_alignment`.
202    pub fn copy_from_preferred_aligned(
203        values: impl AsRef<[T]>,
204        alignment: Alignment,
205        preferred_alignment: Option<Alignment>,
206    ) -> Self
207    where
208        T: Copy,
209    {
210        BufferMut::copy_from_preferred_aligned(values, alignment, preferred_alignment).freeze()
211    }
212
213    /// Create a new zeroed `Buffer` with the given value.
214    pub fn zeroed(len: usize) -> Self {
215        Self::zeroed_aligned(len, Alignment::of::<T>())
216    }
217
218    /// Create a new zeroed `Buffer` with the provided allocator.
219    pub fn zeroed_in(len: usize, allocator: BufferAllocatorRef) -> Self {
220        BufferMut::zeroed_in(len, allocator).freeze()
221    }
222
223    /// Create a new zeroed `Buffer` with the requested alignment.
224    ///
225    /// The allocation is over-aligned to [`Alignment::DEFAULT_ALIGNMENT`] when that is larger than
226    /// `alignment`. Use [`zeroed_preferred_aligned`] to control the over-alignment.
227    ///
228    /// [`zeroed_preferred_aligned`]: Self::zeroed_preferred_aligned
229    pub fn zeroed_aligned(len: usize, alignment: Alignment) -> Self {
230        Self::zeroed_preferred_aligned(len, alignment, Some(Alignment::DEFAULT_ALIGNMENT))
231    }
232
233    /// Create a new zeroed `Buffer` with the requested alignment.
234    ///
235    /// The buffer reports `alignment`, but the underlying allocation is over-aligned to the larger
236    /// of `alignment` and `preferred_alignment`.
237    pub fn zeroed_preferred_aligned(
238        len: usize,
239        alignment: Alignment,
240        preferred_alignment: Option<Alignment>,
241    ) -> Self {
242        BufferMut::zeroed_preferred_aligned(len, alignment, preferred_alignment).freeze()
243    }
244
245    /// Create a new empty `ByteBuffer` with the provided alignment.
246    pub fn empty() -> Self {
247        Self::empty_aligned(Alignment::of::<T>())
248    }
249
250    /// Create a new empty `ByteBuffer` with the provided alignment.
251    ///
252    /// This does not allocate. Empty buffers use an aligned dangling pointer.
253    pub fn empty_aligned(alignment: Alignment) -> Self {
254        const { assert!(size_of::<T>() != 0, "ZSTs are not supported") };
255        if !alignment.is_aligned_to(Alignment::of::<T>()) {
256            misaligned_scalar_type(alignment, Alignment::of::<T>());
257        }
258        Self {
259            ptr: empty_ptr(),
260            length: 0,
261            alignment,
262            backing: None,
263        }
264    }
265
266    /// Create a new full `ByteBuffer` with the given value.
267    pub fn full(item: T, len: usize) -> Self
268    where
269        T: Copy,
270    {
271        BufferMut::full(item, len).freeze()
272    }
273
274    /// Create a full `Buffer` with the given value and allocator.
275    pub fn full_in(item: T, len: usize, allocator: BufferAllocatorRef) -> Self
276    where
277        T: Copy,
278    {
279        BufferMut::full_in(item, len, allocator).freeze()
280    }
281
282    /// Create a `Buffer<T>` zero-copy from a `ByteBuffer`.
283    ///
284    /// ## Panics
285    ///
286    /// Panics if the buffer is not aligned to the size of `T`, or the length is not a multiple of
287    /// the size of `T`.
288    pub fn from_byte_buffer(buffer: ByteBuffer) -> Self {
289        // TODO(ngates): should this preserve the current alignment of the buffer?
290        Self::from_byte_buffer_aligned(buffer, Alignment::of::<T>())
291    }
292
293    /// Create a `Buffer<T>` zero-copy from a `ByteBuffer`.
294    ///
295    /// ## Panics
296    ///
297    /// Panics if the buffer is not aligned to the given alignment, if the length is not a multiple
298    /// of the size of `T`, or if the given alignment is not aligned to that of `T`.
299    pub fn from_byte_buffer_aligned(buffer: ByteBuffer, alignment: Alignment) -> Self {
300        const { assert!(size_of::<T>() != 0, "ZSTs are not supported") };
301        if !alignment.is_aligned_to(Alignment::of::<T>()) {
302            incompatible_scalar_alignment(alignment, Alignment::of::<T>());
303        }
304        if !alignment.is_ptr_aligned(buffer.as_ptr()) {
305            buffer_not_aligned(alignment);
306        }
307        if !buffer.len().is_multiple_of(size_of::<T>()) {
308            buffer_len_not_multiple(buffer.len(), size_of::<T>());
309        }
310        Self {
311            ptr: buffer.ptr.cast(),
312            length: buffer.length / size_of::<T>(),
313            alignment,
314            backing: buffer.backing,
315        }
316    }
317
318    /// Create a `Buffer<T>` zero-copy from a `Bytes`.
319    ///
320    /// ## Panics
321    ///
322    /// Panics if the buffer is not aligned to the size of `T`, or the length is not a multiple of
323    /// the size of `T`.
324    pub fn from_bytes_aligned(bytes: Bytes, alignment: Alignment) -> Self {
325        const { assert!(size_of::<T>() != 0, "ZSTs are not supported") };
326        if !alignment.is_aligned_to(Alignment::of::<T>()) {
327            incompatible_scalar_alignment(alignment, Alignment::of::<T>());
328        }
329        if !alignment.is_ptr_aligned(bytes.as_ptr()) {
330            bytes_not_aligned(alignment);
331        }
332        if !bytes.len().is_multiple_of(size_of::<T>()) {
333            bytes_len_not_multiple(bytes.len(), size_of::<T>());
334        }
335        Self::from_bytes(bytes, alignment)
336    }
337
338    /// Create a buffer with values from the TrustedLen iterator.
339    /// Should be preferred over `from_iter` when the iterator is known to be `TrustedLen`.
340    pub fn from_trusted_len_iter<I: TrustedLen<Item = T>>(iter: I) -> Self {
341        BufferMut::from_trusted_len_iter(iter).freeze()
342    }
343
344    /// Map each element of the buffer with a closure.
345    pub fn map_each_in_place<R, F>(self, mut f: F) -> BufferMut<R>
346    where
347        T: Copy,
348        F: FnMut(T) -> R,
349    {
350        match self.try_into_mut() {
351            Ok(mut_buf) => mut_buf.map_each_in_place(f),
352            Err(buf) => {
353                let len = buf.len();
354                let allocator = buf.allocator().clone();
355                let mut out_buf = BufferMut::with_capacity_in(len, allocator);
356                out_buf
357                    .spare_capacity_mut()
358                    .iter_mut()
359                    .zip(buf)
360                    .for_each(|(out, in_)| {
361                        out.write(f(in_));
362                    });
363                // Safety: just assigned to each value
364                unsafe { out_buf.set_len(len) }
365                out_buf
366            }
367        }
368    }
369
370    /// Clear the buffer, preserving existing capacity.
371    pub fn clear(&mut self) {
372        self.length = 0;
373    }
374
375    /// Returns the length of the buffer in elements of type T.
376    #[allow(clippy::inline_always)]
377    #[inline(always)]
378    pub fn len(&self) -> usize {
379        self.length
380    }
381
382    /// Returns whether the buffer is empty.
383    #[allow(clippy::inline_always)]
384    #[inline(always)]
385    pub fn is_empty(&self) -> bool {
386        self.length == 0
387    }
388
389    /// Returns the alignment of the buffer.
390    #[allow(clippy::inline_always)]
391    #[inline(always)]
392    pub fn alignment(&self) -> Alignment {
393        self.alignment
394    }
395
396    pub(crate) fn allocator(&self) -> &BufferAllocatorRef {
397        match self.backing.as_deref() {
398            Some(backing) => backing.allocator(),
399            None => BufferAllocatorRef::static_ref(),
400        }
401    }
402
403    /// Returns a raw pointer to the buffer's data.
404    #[allow(clippy::inline_always)]
405    #[inline(always)]
406    pub fn as_ptr(&self) -> *const T {
407        self.ptr.as_ptr()
408    }
409
410    /// Returns a slice over the buffer of elements of type T.
411    #[allow(clippy::inline_always)]
412    #[inline(always)]
413    pub fn as_slice(&self) -> &[T] {
414        // SAFETY: ptr points into the live backing and construction checks its alignment.
415        unsafe { std::slice::from_raw_parts(self.ptr.as_ptr(), self.length) }
416    }
417
418    /// Return a view over the buffer as an opaque byte slice.
419    #[allow(clippy::inline_always)]
420    #[inline(always)]
421    pub fn as_bytes(&self) -> &[u8] {
422        // SAFETY: the element range is initialized and remains live through backing.
423        unsafe {
424            std::slice::from_raw_parts(self.ptr.as_ptr().cast(), size_of_val(self.as_slice()))
425        }
426    }
427
428    /// Returns an iterator over the buffer of elements of type T.
429    pub fn iter(&self) -> Iter<'_, T> {
430        Iter {
431            inner: self.as_slice().iter(),
432        }
433    }
434
435    /// Returns a slice of self for the provided range.
436    ///
437    /// # Panics
438    ///
439    /// Requires that `begin <= end` and `end <= self.len()`.
440    /// Also requires that both `begin` and `end` are aligned to the buffer's required alignment.
441    #[allow(clippy::inline_always)]
442    #[inline(always)]
443    pub fn slice(&self, range: impl RangeBounds<usize>) -> Self {
444        self.slice_with_alignment(range, self.alignment)
445    }
446
447    /// Returns a slice of self for the provided range, with no guarantees about the resulting
448    /// alignment.
449    ///
450    /// # Panics
451    ///
452    /// Requires that `begin <= end` and `end <= self.len()`.
453    #[allow(clippy::inline_always)]
454    #[inline(always)]
455    pub fn slice_unaligned(&self, range: impl RangeBounds<usize>) -> Self {
456        self.slice_with_alignment(range, Alignment::of::<u8>())
457    }
458
459    /// Returns a slice of self for the provided range, ensuring the resulting slice has the
460    /// given alignment.
461    ///
462    /// # Panics
463    ///
464    /// Requires that `begin <= end` and `end <= self.len()`.
465    /// Also requires that both `begin` and `end` are aligned to the given alignment.
466    pub fn slice_with_alignment(
467        &self,
468        range: impl RangeBounds<usize>,
469        alignment: Alignment,
470    ) -> Self {
471        let len = self.len();
472        let begin = match range.start_bound() {
473            Bound::Included(&n) => n,
474            Bound::Excluded(&n) => n.checked_add(1).vortex_expect("out of range"),
475            Bound::Unbounded => 0,
476        };
477        let end = match range.end_bound() {
478            Bound::Included(&n) => n.checked_add(1).vortex_expect("out of range"),
479            Bound::Excluded(&n) => n,
480            Bound::Unbounded => len,
481        };
482
483        if begin > end {
484            slice_range_inverted(begin, end);
485        }
486        if end > len {
487            slice_range_out_of_bounds(end, len);
488        }
489
490        if end == begin {
491            // We prefer to return a new empty buffer instead of sharing this one and creating a
492            // strong reference just to hold an empty slice.
493            return Self::empty_aligned(alignment);
494        }
495
496        let begin_byte = begin * size_of::<T>();
497        if !alignment.is_offset_aligned(begin_byte) {
498            slice_start_misaligned(alignment, begin_byte);
499        }
500        if !alignment.is_aligned_to(Alignment::of::<T>()) {
501            vortex_panic!("Slice alignment must at least align to type T")
502        }
503
504        Self {
505            // SAFETY: begin is in bounds and the alignment check applies to the new pointer.
506            ptr: unsafe { self.ptr.add(begin) },
507            length: end - begin,
508            alignment,
509            backing: self.backing.clone(),
510        }
511    }
512
513    /// Returns a slice of self that is equivalent to the given subset.
514    ///
515    /// When processing the buffer you will often end up with `&[T]` that is a subset
516    /// of the underlying buffer. This function turns the slice into a slice of the buffer
517    /// it has been taken from.
518    ///
519    /// # Panics:
520    /// Requires that the given sub slice is in fact contained within the Bytes buffer; otherwise this function will panic.
521    #[allow(clippy::inline_always)]
522    #[inline(always)]
523    pub fn slice_ref(&self, subset: &[T]) -> Self {
524        self.slice_ref_with_alignment(subset, Alignment::of::<T>())
525    }
526
527    /// Returns a slice of self that is equivalent to the given subset.
528    ///
529    /// When processing the buffer you will often end up with `&[T]` that is a subset
530    /// of the underlying buffer. This function turns the slice into a slice of the buffer
531    /// it has been taken from.
532    ///
533    /// # Panics:
534    /// Requires that the given sub slice is in fact contained within the Bytes buffer; otherwise this function will panic.
535    /// Also requires that the given alignment aligns to the type of slice and is smaller or equal to the buffers alignment
536    pub fn slice_ref_with_alignment(&self, subset: &[T], alignment: Alignment) -> Self {
537        if !alignment.is_aligned_to(Alignment::of::<T>()) {
538            vortex_panic!("slice_ref alignment must at least align to type T")
539        }
540
541        if !self.alignment.is_aligned_to(alignment) {
542            vortex_panic!("slice_ref subset alignment must at least align to the buffer alignment")
543        }
544
545        if !alignment.is_ptr_aligned(subset.as_ptr()) {
546            vortex_panic!("slice_ref subset must be aligned to {:?}", alignment);
547        }
548
549        let start = self.as_ptr().addr();
550        let end = start + size_of_val(self.as_slice());
551        let subset_start = subset.as_ptr().addr();
552        let subset_end = subset_start
553            .checked_add(size_of_val(subset))
554            .vortex_expect("slice_ref address overflow");
555        if subset_start < start || subset_end > end {
556            vortex_panic!("slice_ref subset must be contained in the buffer");
557        }
558
559        Self {
560            ptr: NonNull::new(subset.as_ptr().cast_mut()).vortex_expect("slice pointer is null"),
561            length: subset.len(),
562            alignment,
563            backing: self.backing.clone(),
564        }
565    }
566
567    /// Returns the underlying bytes without copying.
568    pub fn into_bytes(self) -> Bytes {
569        if let Some(backing) = self.backing.as_ref()
570            && let BufferBacking::Bytes(bytes) = backing.as_ref()
571        {
572            let offset = self.ptr.cast::<u8>().addr().get() - bytes.as_ptr().addr();
573            let length = self.length * size_of::<T>();
574            if offset == 0 && length == bytes.len() && Arc::strong_count(backing) == 1 {
575                return match self.backing {
576                    Some(backing) => match Arc::try_unwrap(backing) {
577                        Ok(BufferBacking::Bytes(bytes)) => bytes,
578                        _ => unreachable!(),
579                    },
580                    None => unreachable!(),
581                };
582            }
583            return bytes.slice(offset..offset + length);
584        }
585        match self.backing {
586            Some(backing) => Bytes::from_owner(BufferBytesOwner {
587                ptr: self.ptr.cast(),
588                length: self.length * size_of::<T>(),
589                backing,
590            }),
591            None => Bytes::new(),
592        }
593    }
594
595    /// Return the ByteBuffer for this `Buffer<T>`.
596    pub fn into_byte_buffer(self) -> ByteBuffer {
597        ByteBuffer {
598            ptr: self.ptr.cast(),
599            length: self.length * size_of::<T>(),
600            alignment: self.alignment,
601            backing: self.backing,
602        }
603    }
604
605    /// Try to convert self into `BufferMut<T>` if there is only a single strong reference.
606    pub fn try_into_mut(self) -> Result<BufferMut<T>, Self> {
607        let Self {
608            ptr,
609            length,
610            alignment,
611            backing,
612        } = self;
613        let Some(backing) = backing else {
614            return Ok(BufferMut::empty_aligned(alignment));
615        };
616        if !matches!(backing.as_ref(), BufferBacking::Owned(_)) {
617            return Err(Self {
618                ptr,
619                length,
620                alignment,
621                backing: Some(backing),
622            });
623        }
624        match Arc::try_unwrap(backing) {
625            Ok(BufferBacking::Owned(allocation)) => {
626                let offset = ptr.addr().get() - allocation.ptr().addr().get();
627                let capacity = if allocation.size() == 0 {
628                    0
629                } else {
630                    (allocation.size() - offset) / size_of::<T>()
631                };
632                Ok(BufferMut {
633                    allocation,
634                    ptr,
635                    length,
636                    capacity,
637                    alignment,
638                    _marker: Default::default(),
639                })
640            }
641            Ok(_) => unreachable!(),
642            Err(backing) => Err(Self {
643                ptr,
644                length,
645                alignment,
646                backing: Some(backing),
647            }),
648        }
649    }
650
651    /// Convert self into `BufferMut<T>`, cloning the data if there are multiple strong references.
652    pub fn into_mut(self) -> BufferMut<T>
653    where
654        T: Copy,
655    {
656        self.try_into_mut().unwrap_or_else(|buffer| {
657            let allocator = buffer.allocator().clone();
658            BufferMut::<T>::copy_from_aligned_in(&buffer, buffer.alignment, allocator)
659        })
660    }
661
662    /// Returns whether a `Buffer<T>` is aligned to the given alignment.
663    pub fn is_aligned(&self, alignment: Alignment) -> bool {
664        alignment.is_ptr_aligned(self.as_ptr())
665    }
666
667    /// Return a `Buffer<T>` with the given alignment. Where possible, this will be zero-copy.
668    pub fn aligned(mut self, alignment: Alignment) -> Self
669    where
670        T: Copy,
671    {
672        if alignment.is_ptr_aligned(self.as_ptr()) {
673            self.alignment = alignment;
674            self
675        } else {
676            #[cfg(feature = "warn-copy")]
677            {
678                let bt = std::backtrace::Backtrace::capture();
679                tracing::warn!(
680                    "Buffer is not aligned to requested alignment {alignment}, copying: {bt}"
681                )
682            }
683            let allocator = self.allocator().clone();
684            BufferMut::copy_from_aligned_in(self, alignment, allocator).freeze()
685        }
686    }
687
688    /// Return a `Buffer<T>` with the given alignment. Panics if the buffer is not aligned.
689    pub fn ensure_aligned(mut self, alignment: Alignment) -> Self {
690        if alignment.is_ptr_aligned(self.as_ptr()) {
691            self.alignment = alignment;
692            self
693        } else {
694            vortex_panic!("Buffer is not aligned to requested alignment {}", alignment)
695        }
696    }
697}
698
699impl<T> Buffer<T> {
700    /// Transmute a `Buffer<T>` into a `Buffer<U>`.
701    ///
702    /// # Safety
703    ///
704    /// The caller must ensure that all possible bit representations of type `T` are valid when
705    /// interpreted as type `U`.
706    /// See [`std::mem::transmute`] for more details.
707    ///
708    /// # Panics
709    ///
710    /// Panics if the type `U` does not have the same size and alignment as `T`.
711    pub unsafe fn transmute<U>(self) -> Buffer<U> {
712        assert_eq!(size_of::<T>(), size_of::<U>(), "Buffer type size mismatch");
713        assert_eq!(
714            align_of::<T>(),
715            align_of::<U>(),
716            "Buffer type alignment mismatch"
717        );
718
719        Buffer {
720            ptr: self.ptr.cast(),
721            length: self.length,
722            alignment: self.alignment,
723            backing: self.backing,
724        }
725    }
726}
727
728/// An iterator over Buffer elements.
729///
730/// This is an analog to the `std::slice::Iter` type.
731pub struct Iter<'a, T> {
732    inner: std::slice::Iter<'a, T>,
733}
734
735impl<'a, T> Iterator for Iter<'a, T> {
736    type Item = &'a T;
737
738    #[inline]
739    fn next(&mut self) -> Option<Self::Item> {
740        self.inner.next()
741    }
742
743    #[inline]
744    fn size_hint(&self) -> (usize, Option<usize>) {
745        self.inner.size_hint()
746    }
747
748    #[inline]
749    fn count(self) -> usize {
750        self.inner.count()
751    }
752
753    #[inline]
754    fn last(self) -> Option<Self::Item> {
755        self.inner.last()
756    }
757
758    #[inline]
759    fn nth(&mut self, n: usize) -> Option<Self::Item> {
760        self.inner.nth(n)
761    }
762}
763
764impl<T> ExactSizeIterator for Iter<'_, T> {
765    #[inline]
766    fn len(&self) -> usize {
767        self.inner.len()
768    }
769}
770
771impl<T: Debug> Debug for Buffer<T> {
772    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
773        f.debug_struct(&format!("Buffer<{}>", type_name::<T>()))
774            .field("length", &self.length)
775            .field("alignment", &self.alignment)
776            .field("as_slice", &TruncatedDebug(self.as_slice()))
777            .finish()
778    }
779}
780
781impl<T> Deref for Buffer<T> {
782    type Target = [T];
783
784    #[inline]
785    fn deref(&self) -> &Self::Target {
786        self.as_slice()
787    }
788}
789
790impl<T> AsRef<[T]> for Buffer<T> {
791    #[inline]
792    fn as_ref(&self) -> &[T] {
793        self.as_slice()
794    }
795}
796
797impl<T> FromIterator<T> for Buffer<T> {
798    #[inline]
799    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
800        BufferMut::from_iter(iter).freeze()
801    }
802}
803
804impl<T> From<Vec<T>> for Buffer<T>
805where
806    T: Send + Sync + 'static,
807{
808    fn from(value: Vec<T>) -> Self {
809        const { assert!(size_of::<T>() != 0, "ZSTs are not supported") };
810        let length = value.len();
811        let alignment = Alignment::of::<T>();
812        if std::mem::needs_drop::<T>() {
813            // Keep the typed owner so its elements are dropped.
814            Self {
815                ptr: NonNull::new(value.as_ptr().cast_mut())
816                    .vortex_expect("a Vec always has a non-null pointer"),
817                length,
818                alignment,
819                backing: Some(Arc::new(BufferBacking::External {
820                    _owner: Box::new(value),
821                })),
822            }
823        } else {
824            Self::from_allocation(Allocation::from_vec(value), 0, length, alignment)
825        }
826    }
827}
828
829impl From<Bytes> for ByteBuffer {
830    fn from(bytes: Bytes) -> Self {
831        Self::from_bytes(bytes, Alignment::of::<u8>())
832    }
833}
834
835impl Buf for ByteBuffer {
836    #[inline]
837    fn remaining(&self) -> usize {
838        self.len()
839    }
840
841    #[inline]
842    fn chunk(&self) -> &[u8] {
843        self.as_slice()
844    }
845
846    #[inline]
847    fn advance(&mut self, cnt: usize) {
848        if !self.alignment.is_offset_aligned(cnt) {
849            vortex_panic!(
850                "Cannot advance buffer by {} items, resulting alignment is not {}",
851                cnt,
852                self.alignment
853            );
854        }
855        assert!(cnt <= self.length, "cannot advance past the buffer length");
856        // SAFETY: cnt is within the initialized byte range.
857        self.ptr = unsafe { self.ptr.add(cnt) };
858        self.length -= cnt;
859    }
860}
861
862struct BufferBytesOwner {
863    ptr: NonNull<u8>,
864    length: usize,
865    backing: Arc<BufferBacking>,
866}
867
868// SAFETY: the owner exposes immutable initialized bytes and keeps their backing live.
869unsafe impl Send for BufferBytesOwner {}
870unsafe impl Sync for BufferBytesOwner {}
871
872impl AsRef<[u8]> for BufferBytesOwner {
873    fn as_ref(&self) -> &[u8] {
874        let _ = &self.backing;
875        // SAFETY: ptr and length came from a live Buffer.
876        unsafe { std::slice::from_raw_parts(self.ptr.as_ptr(), self.length) }
877    }
878}
879
880fn empty_ptr<T>() -> NonNull<T> {
881    let addr = 1usize << (usize::BITS - 1);
882    NonNull::new(std::ptr::without_provenance_mut(addr)).vortex_expect("empty pointer is non-null")
883}
884
885/// Owned iterator over a [`Buffer`].
886pub struct BufferIterator<T: Copy> {
887    // Keep the buffer alive for the duration of the iteration.
888    _buffer: Buffer<T>,
889    ptr: *const T,
890    end: *const T,
891}
892
893// SAFETY: `BufferIterator` is a `Buffer<T>` plus two cursors into it, so it can safely be
894// `Send`/`Sync` exactly when `Buffer<T>` is. Same bounds as `std::vec::IntoIter`.
895unsafe impl<T: Copy + Send> Send for BufferIterator<T> {}
896unsafe impl<T: Copy + Sync> Sync for BufferIterator<T> {}
897
898impl<T: Copy> Iterator for BufferIterator<T> {
899    type Item = T;
900
901    #[inline]
902    fn next(&mut self) -> Option<Self::Item> {
903        if self.ptr == self.end {
904            None
905        } else {
906            // SAFETY: ptr is within the buffer and has not reached end.
907            let value = unsafe { self.ptr.read() };
908            self.ptr = unsafe { self.ptr.add(1) };
909            Some(value)
910        }
911    }
912
913    #[inline]
914    fn size_hint(&self) -> (usize, Option<usize>) {
915        let remaining = unsafe { self.end.offset_from(self.ptr) } as usize;
916        (remaining, Some(remaining))
917    }
918}
919
920impl<T: Copy> ExactSizeIterator for BufferIterator<T> {}
921
922impl<T: Copy> IntoIterator for Buffer<T> {
923    type Item = T;
924    type IntoIter = BufferIterator<T>;
925
926    #[inline]
927    fn into_iter(self) -> Self::IntoIter {
928        let ptr = self.as_slice().as_ptr();
929        let end = unsafe { ptr.add(self.len()) };
930        BufferIterator {
931            _buffer: self,
932            ptr,
933            end,
934        }
935    }
936}
937
938impl<T> From<BufferMut<T>> for Buffer<T> {
939    #[inline]
940    fn from(value: BufferMut<T>) -> Self {
941        value.freeze()
942    }
943}
944
945// If panic helpers are inlined, this makes callers of BufferMut's functions
946// not inline-eligible. This is more ugly than direct panics but helps with
947// performance
948
949#[cold]
950#[inline(never)]
951fn misaligned_scalar_type(alignment: Alignment, scalar_align: Alignment) -> ! {
952    vortex_panic!("Alignment {alignment} must align to the scalar type's alignment {scalar_align}")
953}
954
955#[cold]
956#[inline(never)]
957fn incompatible_scalar_alignment(alignment: Alignment, scalar_align: Alignment) -> ! {
958    vortex_panic!(
959        "Alignment {alignment} must be compatible with the scalar type's alignment {scalar_align}"
960    )
961}
962
963#[cold]
964#[inline(never)]
965fn buffer_not_aligned(alignment: Alignment) -> ! {
966    vortex_panic!("Buffer must align to the requested alignment {alignment}")
967}
968
969#[cold]
970#[inline(never)]
971fn bytes_not_aligned(alignment: Alignment) -> ! {
972    vortex_panic!("Bytes alignment must align to the requested alignment {alignment}")
973}
974
975#[cold]
976#[inline(never)]
977fn buffer_len_not_multiple(len: usize, scalar_size: usize) -> ! {
978    vortex_panic!("Buffer length {len} must be a multiple of the scalar type's size {scalar_size}")
979}
980
981#[cold]
982#[inline(never)]
983fn bytes_len_not_multiple(len: usize, scalar_size: usize) -> ! {
984    vortex_panic!("Bytes length {len} must be a multiple of the scalar type's size {scalar_size}")
985}
986
987#[cold]
988#[inline(never)]
989fn slice_range_inverted(begin: usize, end: usize) -> ! {
990    vortex_panic!("range start must not be greater than end: {begin:?} <= {end:?}")
991}
992
993#[cold]
994#[inline(never)]
995fn slice_range_out_of_bounds(end: usize, len: usize) -> ! {
996    vortex_panic!("range end out of bounds: {end:?} > {len:?}")
997}
998
999#[cold]
1000#[inline(never)]
1001fn slice_start_misaligned(alignment: Alignment, begin_byte: usize) -> ! {
1002    vortex_panic!("range start must be aligned to {alignment:?}, byte {begin_byte}")
1003}
1004
1005#[cfg(test)]
1006mod test {
1007    use std::mem::align_of;
1008    use std::sync::Arc;
1009    use std::sync::atomic::AtomicUsize;
1010    use std::sync::atomic::Ordering;
1011
1012    use bytes::Buf;
1013    use bytes::Bytes;
1014
1015    use crate::Alignment;
1016    use crate::Buffer;
1017    use crate::BufferBacking;
1018    use crate::ByteBuffer;
1019    use crate::buffer;
1020
1021    #[test]
1022    fn align() {
1023        let buf = buffer![0u8, 1, 2];
1024        let aligned = buf.aligned(Alignment::new(32));
1025        assert_eq!(aligned.alignment(), Alignment::new(32));
1026        assert_eq!(aligned.as_slice(), &[0, 1, 2]);
1027    }
1028
1029    #[test]
1030    fn buffer_iterator_send_sync() {
1031        fn assert_send_sync<T: Send + Sync>(_: &T) {}
1032
1033        let mut iter = buffer![0i32, 1, 2, 3].into_iter();
1034        assert_send_sync(&iter);
1035        iter.next();
1036        let remaining: Vec<i32> = std::thread::spawn(move || iter.collect()).join().unwrap();
1037        assert_eq!(remaining, vec![1, 2, 3]);
1038    }
1039
1040    #[test]
1041    fn slice() {
1042        let buf = buffer![0, 1, 2, 3, 4];
1043        assert_eq!(buf.slice(1..3).as_slice(), &[1, 2]);
1044        assert_eq!(buf.slice(1..=3).as_slice(), &[1, 2, 3]);
1045    }
1046
1047    #[test]
1048    fn slice_unaligned() {
1049        let buf = buffer![0i32, 1, 2, 3, 4].into_byte_buffer();
1050        // With a regular slice, this would panic. See [`slice_bad_alignment`].
1051        let sliced = buf.slice_unaligned(1..2);
1052        // Verify the slice has the expected length (1 byte from index 1 to 2).
1053        assert_eq!(sliced.len(), 1);
1054        // The original buffer has i32 values [0, 1, 2, 3, 4].
1055        // In little-endian bytes, 0i32 = [0, 0, 0, 0], so byte at index 1 is 0.
1056        assert_eq!(sliced.as_slice(), &[0]);
1057    }
1058
1059    #[test]
1060    #[should_panic]
1061    fn slice_bad_alignment() {
1062        let buf = buffer![0i32, 1, 2, 3, 4].into_byte_buffer();
1063        // We should only be able to slice this buffer on 4-byte (i32) boundaries.
1064        buf.slice(1..2);
1065    }
1066
1067    #[test]
1068    fn bytes_buf() {
1069        let mut buf = ByteBuffer::copy_from("helloworld".as_bytes());
1070        assert_eq!(buf.remaining(), 10);
1071        assert_eq!(buf.chunk(), b"helloworld");
1072
1073        buf.advance(5);
1074        assert_eq!(buf.remaining(), 5);
1075        assert_eq!(buf.as_slice(), b"world");
1076        assert_eq!(buf.chunk(), b"world");
1077    }
1078
1079    #[test]
1080    fn buffer_zeroed() {
1081        const LEN: usize = 17;
1082
1083        let buf = Buffer::<u32>::zeroed(LEN);
1084
1085        assert!(buf.is_aligned(Alignment::of::<u32>()));
1086        assert_eq!(buf.as_slice(), &[0; LEN]);
1087    }
1088
1089    #[test]
1090    fn buffer_zeroed_aligned() {
1091        const LEN: usize = 17;
1092        let alignment = Alignment::new(64);
1093
1094        let buf = Buffer::<u32>::zeroed_aligned(LEN, alignment);
1095
1096        assert!(buf.is_aligned(alignment));
1097        assert_eq!(buf.as_slice(), &[0; LEN]);
1098    }
1099
1100    #[test]
1101    fn copy_from_over_aligns_to_default() {
1102        let values = [1u32, 2, 3];
1103        let buf = Buffer::<u32>::copy_from(values);
1104
1105        // The buffer reports the scalar type's alignment, ...
1106        assert_eq!(buf.alignment(), Alignment::of::<u32>());
1107        // ... but the underlying allocation is over-aligned to DEFAULT_ALIGNMENT.
1108        assert!(buf.is_aligned(Alignment::DEFAULT_ALIGNMENT));
1109        assert_eq!(buf.as_slice(), &values);
1110    }
1111
1112    #[test]
1113    fn zeroed_over_aligns_to_default() {
1114        const LEN: usize = 17;
1115
1116        let buf = Buffer::<u32>::zeroed(LEN);
1117
1118        assert_eq!(buf.alignment(), Alignment::of::<u32>());
1119        assert!(buf.is_aligned(Alignment::DEFAULT_ALIGNMENT));
1120        assert_eq!(buf.as_slice(), &[0; LEN]);
1121    }
1122
1123    #[test]
1124    fn from_vec() {
1125        let vec = vec![1, 2, 3, 4, 5];
1126        let buff = Buffer::from(vec.clone());
1127        assert!(buff.is_aligned(Alignment::of::<i32>()));
1128        assert_eq!(vec, buff.as_ref());
1129    }
1130
1131    #[test]
1132    fn from_vec_adopts_allocation() {
1133        let mut vec = Vec::with_capacity(16);
1134        vec.extend([1u32, 2, 3, 4, 5]);
1135        let ptr = vec.as_ptr();
1136        let capacity = vec.capacity();
1137
1138        let buffer = Buffer::from(vec);
1139        assert_eq!(buffer.as_ptr(), ptr);
1140
1141        let Ok(mut buffer) = buffer.try_into_mut() else {
1142            panic!("Vec-backed buffer should be uniquely owned")
1143        };
1144        assert_eq!(buffer.capacity(), capacity);
1145        assert_eq!(buffer.allocation.alignment(), align_of::<u32>());
1146
1147        buffer.extend(6..=32);
1148        assert_eq!(buffer.as_slice(), (1..=32).collect::<Vec<_>>());
1149        assert_eq!(buffer.allocation.alignment(), align_of::<u32>());
1150    }
1151
1152    #[test]
1153    fn byte_owner_preserves_slice_and_lifetime() {
1154        struct Owner {
1155            values: Vec<u8>,
1156            drops: Arc<AtomicUsize>,
1157        }
1158
1159        impl AsRef<[u8]> for Owner {
1160            fn as_ref(&self) -> &[u8] {
1161                &self.values[1..4]
1162            }
1163        }
1164
1165        impl Drop for Owner {
1166            fn drop(&mut self) {
1167                self.drops.fetch_add(1, Ordering::Relaxed);
1168            }
1169        }
1170
1171        let drops = Arc::new(AtomicUsize::new(0));
1172        let owner = Owner {
1173            values: vec![0, 1, 2, 3, 4],
1174            drops: Arc::clone(&drops),
1175        };
1176        let ptr = owner.as_ref().as_ptr();
1177        let buffer = ByteBuffer::from(Bytes::from_owner(owner));
1178        assert_eq!(buffer.as_ptr(), ptr);
1179        assert_eq!(buffer.as_slice(), [1, 2, 3]);
1180        let view = buffer.slice(1..);
1181        drop(buffer);
1182        assert_eq!(drops.load(Ordering::Relaxed), 0);
1183        assert_eq!(view.as_slice(), [2, 3]);
1184        drop(view);
1185        assert_eq!(drops.load(Ordering::Relaxed), 1);
1186    }
1187
1188    #[test]
1189    fn bytes_round_trip_reuses_owner() {
1190        let bytes = Bytes::from_static(&[1, 2, 3, 4]);
1191        let ptr = bytes.as_ptr();
1192
1193        let buffer = ByteBuffer::from(bytes);
1194        assert!(matches!(
1195            buffer.backing.as_deref(),
1196            Some(BufferBacking::Bytes(_))
1197        ));
1198        let bytes = buffer.into_bytes();
1199
1200        assert_eq!(bytes.as_ptr(), ptr);
1201        assert_eq!(bytes.as_ref(), &[1, 2, 3, 4]);
1202    }
1203
1204    #[test]
1205    fn external_try_into_mut_preserves_backing() {
1206        let buffer = ByteBuffer::from(Bytes::from_static(&[1, 2, 3, 4]));
1207        let Some(original_backing) = buffer.backing.as_ref() else {
1208            panic!("external buffer has no backing")
1209        };
1210        let backing = Arc::as_ptr(original_backing);
1211
1212        let Err(buffer) = buffer.try_into_mut() else {
1213            panic!("external buffer became mutable")
1214        };
1215
1216        let Some(new_backing) = buffer.backing.as_ref() else {
1217            panic!("external buffer has no backing")
1218        };
1219        assert_eq!(Arc::as_ptr(new_backing), backing);
1220    }
1221
1222    #[test]
1223    fn from_u8_vec_preserves_capacity() {
1224        let mut vec = Vec::with_capacity(16);
1225        vec.extend([1u8, 2, 3]);
1226
1227        let buffer = Buffer::from(vec);
1228        let Ok(buffer) = buffer.try_into_mut() else {
1229            panic!("Vec-backed buffer should be uniquely owned")
1230        };
1231        assert_eq!(buffer.capacity(), 16);
1232    }
1233
1234    #[test]
1235    fn sliced_buffer_into_mut_has_safe_capacity() {
1236        let mut original = crate::BufferMut::with_capacity(128);
1237        original.extend(0u32..100);
1238        let original = original.freeze();
1239        let sliced = original.slice(64..96);
1240        drop(original);
1241
1242        let Ok(mut sliced) = sliced.try_into_mut() else {
1243            panic!("uniquely owned slice should become mutable")
1244        };
1245        let ptr = sliced.as_ptr();
1246        let capacity = sliced.capacity();
1247        sliced.push_n(0, capacity - sliced.len());
1248        assert_eq!(sliced.len(), capacity);
1249        assert_eq!(sliced.as_ptr(), ptr);
1250        sliced.push(42);
1251        assert_eq!(&sliced[..32], (64u32..96).collect::<Vec<_>>());
1252        assert_eq!(&sliced[32..capacity], vec![0; capacity - 32]);
1253        assert_eq!(sliced[capacity], 42);
1254    }
1255
1256    #[test]
1257    fn from_vec_preserves_drop_glue() {
1258        struct DropValue(Arc<AtomicUsize>);
1259
1260        impl Drop for DropValue {
1261            fn drop(&mut self) {
1262                self.0.fetch_add(1, Ordering::Relaxed);
1263            }
1264        }
1265
1266        let drops = Arc::new(AtomicUsize::new(0));
1267        let values = (0..3)
1268            .map(|_| DropValue(Arc::clone(&drops)))
1269            .collect::<Vec<_>>();
1270        let buffer = Buffer::from(values);
1271
1272        assert_eq!(drops.load(Ordering::Relaxed), 0);
1273        drop(buffer);
1274        assert_eq!(drops.load(Ordering::Relaxed), 3);
1275    }
1276
1277    #[test]
1278    fn empty_aligned_max_alignment() {
1279        // Empty buffers are backed by a static and must satisfy any valid alignment.
1280        let buf = Buffer::<u8>::empty_aligned(Alignment::MAX);
1281        assert!(buf.is_empty());
1282        assert!(buf.is_aligned(Alignment::MAX));
1283    }
1284
1285    #[test]
1286    fn empty_has_no_backing() {
1287        assert!(Buffer::<u8>::empty().backing.is_none());
1288    }
1289
1290    #[test]
1291    fn empty_slice_preserves_alignment() {
1292        let buf = Buffer::<u64>::zeroed_aligned(8, Alignment::new(64));
1293        let sliced = buf.slice(0..0);
1294        assert!(sliced.is_empty());
1295        assert_eq!(sliced.alignment(), Alignment::new(64));
1296        assert!(sliced.is_aligned(Alignment::new(64)));
1297    }
1298
1299    #[test]
1300    fn empty_into_mut_preserves_alignment() {
1301        let buf = Buffer::<u8>::empty_aligned(Alignment::new(64));
1302        let buf_mut = buf.into_mut();
1303        assert_eq!(buf_mut.alignment(), Alignment::new(64));
1304        assert!(buf_mut.is_empty());
1305    }
1306
1307    #[test]
1308    fn test_slice_unaligned_end_pos() {
1309        let data = vec![0u8; 2];
1310        // Overalign the u8 vector.
1311        let aligned_buffer = Buffer::copy_from_aligned(&data, Alignment::new(8));
1312        // Previously, `Buffer::slice` incorrectly asserted that the end position
1313        // must be aligned. That assertion has been removed such that the end
1314        // position can be arbitrary and only the beginning of the slice needs
1315        // to be aligned.
1316        aligned_buffer.slice(0..1);
1317    }
1318
1319    #[test]
1320    fn test_empty_equality() {
1321        let a = Buffer::<u16>::empty();
1322        let b = Buffer::<u16>::empty();
1323
1324        assert_eq!(a, b);
1325    }
1326}