1use 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::marker::PhantomData;
12use std::ops::Deref;
13use std::ops::RangeBounds;
14
15use bytes::Buf;
16use bytes::Bytes;
17use vortex_error::VortexExpect;
18use vortex_error::vortex_panic;
19
20use crate::Alignment;
21use crate::BufferMut;
22use crate::ByteBuffer;
23use crate::debug::TruncatedDebug;
24use crate::trusted_len::TrustedLen;
25
26#[derive(Clone)]
28pub struct Buffer<T> {
29 pub(crate) bytes: Bytes,
30 pub(crate) length: usize,
31 pub(crate) alignment: Alignment,
32 pub(crate) _marker: PhantomData<T>,
33}
34
35const EMPTY_BACKING: &[u8] = {
39 let addr = 1usize << (usize::BITS - 1);
40 assert!(Alignment::MAX.is_offset_aligned(addr));
41 unsafe { std::slice::from_raw_parts(std::ptr::without_provenance(addr), 0) }
43};
44
45impl<T> Default for Buffer<T> {
46 fn default() -> Self {
47 Self {
48 bytes: Bytes::from_static(EMPTY_BACKING),
49 length: 0,
50 alignment: Alignment::of::<T>(),
51 _marker: PhantomData,
52 }
53 }
54}
55
56impl<T> PartialEq for Buffer<T> {
57 #[inline]
58 fn eq(&self, other: &Self) -> bool {
59 self.bytes == other.bytes
60 }
61}
62
63impl<T> Eq for Buffer<T> {}
64
65impl<T> Ord for Buffer<T> {
66 #[inline]
67 fn cmp(&self, other: &Self) -> Ordering {
68 self.bytes.cmp(&other.bytes)
69 }
70}
71
72impl<T> PartialOrd for Buffer<T> {
73 #[inline]
74 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
75 Some(self.cmp(other))
76 }
77}
78
79impl<T> Hash for Buffer<T> {
80 #[inline]
81 fn hash<H: Hasher>(&self, state: &mut H) {
82 self.bytes.as_ref().hash(state)
83 }
84}
85
86impl<T> Buffer<T> {
87 pub fn copy_from(values: impl AsRef<[T]>) -> Self {
94 BufferMut::copy_from(values).freeze()
95 }
96
97 pub fn copy_from_aligned(values: impl AsRef<[T]>, alignment: Alignment) -> Self {
104 Self::copy_from_preferred_aligned(values, alignment, Some(Alignment::DEFAULT_ALIGNMENT))
105 }
106
107 pub fn copy_from_preferred_aligned(
112 values: impl AsRef<[T]>,
113 alignment: Alignment,
114 preferred_alignment: Option<Alignment>,
115 ) -> Self {
116 BufferMut::copy_from_preferred_aligned(values, alignment, preferred_alignment).freeze()
117 }
118
119 pub fn zeroed(len: usize) -> Self {
121 Self::zeroed_aligned(len, Alignment::of::<T>())
122 }
123
124 pub fn zeroed_aligned(len: usize, alignment: Alignment) -> Self {
131 Self::zeroed_preferred_aligned(len, alignment, Some(Alignment::DEFAULT_ALIGNMENT))
132 }
133
134 pub fn zeroed_preferred_aligned(
139 len: usize,
140 alignment: Alignment,
141 preferred_alignment: Option<Alignment>,
142 ) -> Self {
143 BufferMut::zeroed_preferred_aligned(len, alignment, preferred_alignment).freeze()
144 }
145
146 pub fn empty() -> Self {
148 Self::empty_aligned(Alignment::of::<T>())
149 }
150
151 pub fn empty_aligned(alignment: Alignment) -> Self {
156 if !alignment.is_aligned_to(Alignment::of::<T>()) {
157 vortex_panic!(
158 "Alignment {} must align to the scalar type's alignment {}",
159 alignment,
160 Alignment::of::<T>(),
161 );
162 }
163 Self {
164 bytes: Bytes::from_static(EMPTY_BACKING),
165 length: 0,
166 alignment,
167 _marker: PhantomData,
168 }
169 }
170
171 pub fn full(item: T, len: usize) -> Self
173 where
174 T: Copy,
175 {
176 BufferMut::full(item, len).freeze()
177 }
178
179 pub fn from_byte_buffer(buffer: ByteBuffer) -> Self {
186 Self::from_byte_buffer_aligned(buffer, Alignment::of::<T>())
188 }
189
190 pub fn from_byte_buffer_aligned(buffer: ByteBuffer, alignment: Alignment) -> Self {
197 Self::from_bytes_aligned(buffer.into_inner(), alignment)
198 }
199
200 pub fn from_bytes_aligned(bytes: Bytes, alignment: Alignment) -> Self {
207 if !alignment.is_aligned_to(Alignment::of::<T>()) {
208 vortex_panic!(
209 "Alignment {} must be compatible with the scalar type's alignment {}",
210 alignment,
211 Alignment::of::<T>(),
212 );
213 }
214 if !alignment.is_ptr_aligned(bytes.as_ptr()) {
215 vortex_panic!(
216 "Bytes alignment must align to the requested alignment {}",
217 alignment,
218 );
219 }
220 if !bytes.len().is_multiple_of(size_of::<T>()) {
221 vortex_panic!(
222 "Bytes length {} must be a multiple of the scalar type's size {}",
223 bytes.len(),
224 size_of::<T>()
225 );
226 }
227 let length = bytes.len() / size_of::<T>();
228 Self {
229 bytes,
230 length,
231 alignment,
232 _marker: Default::default(),
233 }
234 }
235
236 pub fn from_trusted_len_iter<I: TrustedLen<Item = T>>(iter: I) -> Self {
239 let (_, upper_bound) = iter.size_hint();
240 let mut buffer = BufferMut::with_capacity(
241 upper_bound.vortex_expect("TrustedLen iterator has no upper bound"),
242 );
243 buffer.extend_trusted(iter);
244 buffer.freeze()
245 }
246
247 pub fn map_each_in_place<R, F>(self, mut f: F) -> BufferMut<R>
249 where
250 T: Copy,
251 F: FnMut(T) -> R,
252 {
253 match self.try_into_mut() {
254 Ok(mut_buf) => mut_buf.map_each_in_place(f),
255 Err(buf) => {
256 let len = buf.len();
257 let mut out_buf = BufferMut::with_capacity(len);
258 out_buf
259 .spare_capacity_mut()
260 .iter_mut()
261 .zip(buf)
262 .for_each(|(out, in_)| {
263 out.write(f(in_));
264 });
265 unsafe { out_buf.set_len(len) }
267 out_buf
268 }
269 }
270 }
271
272 pub fn clear(&mut self) {
274 self.bytes.clear();
275 self.length = 0;
276 }
277
278 #[inline(always)]
280 pub fn len(&self) -> usize {
281 self.length
282 }
283
284 #[inline(always)]
286 pub fn is_empty(&self) -> bool {
287 self.length == 0
288 }
289
290 #[inline(always)]
292 pub fn alignment(&self) -> Alignment {
293 self.alignment
294 }
295
296 #[inline(always)]
298 pub fn as_slice(&self) -> &[T] {
299 unsafe { std::slice::from_raw_parts(self.bytes.as_ptr().cast(), self.length) }
301 }
302
303 #[inline(always)]
305 pub fn as_bytes(&self) -> &[u8] {
306 self.bytes.as_ref()
307 }
308
309 pub fn iter(&self) -> Iter<'_, T> {
311 Iter {
312 inner: self.as_slice().iter(),
313 }
314 }
315
316 #[inline(always)]
323 pub fn slice(&self, range: impl RangeBounds<usize>) -> Self {
324 self.slice_with_alignment(range, self.alignment)
325 }
326
327 #[inline(always)]
334 pub fn slice_unaligned(&self, range: impl RangeBounds<usize>) -> Self {
335 self.slice_with_alignment(range, Alignment::of::<u8>())
336 }
337
338 pub fn slice_with_alignment(
346 &self,
347 range: impl RangeBounds<usize>,
348 alignment: Alignment,
349 ) -> Self {
350 let len = self.len();
351 let begin = match range.start_bound() {
352 Bound::Included(&n) => n,
353 Bound::Excluded(&n) => n.checked_add(1).vortex_expect("out of range"),
354 Bound::Unbounded => 0,
355 };
356 let end = match range.end_bound() {
357 Bound::Included(&n) => n.checked_add(1).vortex_expect("out of range"),
358 Bound::Excluded(&n) => n,
359 Bound::Unbounded => len,
360 };
361
362 if begin > end {
363 vortex_panic!(
364 "range start must not be greater than end: {:?} <= {:?}",
365 begin,
366 end
367 );
368 }
369 if end > len {
370 vortex_panic!("range end out of bounds: {:?} > {:?}", end, len);
371 }
372
373 if end == begin {
374 return Self::empty_aligned(alignment);
377 }
378
379 let begin_byte = begin * size_of::<T>();
380 let end_byte = end * size_of::<T>();
381
382 if !alignment.is_offset_aligned(begin_byte) {
383 vortex_panic!(
384 "range start must be aligned to {alignment:?}, byte {}",
385 begin_byte
386 );
387 }
388 if !alignment.is_aligned_to(Alignment::of::<T>()) {
389 vortex_panic!("Slice alignment must at least align to type T")
390 }
391
392 Self {
393 bytes: self.bytes.slice(begin_byte..end_byte),
394 length: end - begin,
395 alignment,
396 _marker: Default::default(),
397 }
398 }
399
400 #[inline(always)]
409 pub fn slice_ref(&self, subset: &[T]) -> Self {
410 self.slice_ref_with_alignment(subset, Alignment::of::<T>())
411 }
412
413 pub fn slice_ref_with_alignment(&self, subset: &[T], alignment: Alignment) -> Self {
423 if !alignment.is_aligned_to(Alignment::of::<T>()) {
424 vortex_panic!("slice_ref alignment must at least align to type T")
425 }
426
427 if !self.alignment.is_aligned_to(alignment) {
428 vortex_panic!("slice_ref subset alignment must at least align to the buffer alignment")
429 }
430
431 if !alignment.is_ptr_aligned(subset.as_ptr()) {
432 vortex_panic!("slice_ref subset must be aligned to {:?}", alignment);
433 }
434
435 let subset_u8 =
436 unsafe { std::slice::from_raw_parts(subset.as_ptr().cast(), size_of_val(subset)) };
437
438 Self {
439 bytes: self.bytes.slice_ref(subset_u8),
440 length: subset.len(),
441 alignment,
442 _marker: Default::default(),
443 }
444 }
445
446 pub fn inner(&self) -> &Bytes {
448 debug_assert_eq!(
449 self.length * size_of::<T>(),
450 self.bytes.len(),
451 "Own length has to be the same as the underlying bytes length"
452 );
453 &self.bytes
454 }
455
456 pub fn into_inner(self) -> Bytes {
458 debug_assert_eq!(
459 self.length * size_of::<T>(),
460 self.bytes.len(),
461 "Own length has to be the same as the underlying bytes length"
462 );
463 self.bytes
464 }
465
466 pub fn into_byte_buffer(self) -> ByteBuffer {
468 ByteBuffer {
469 bytes: self.bytes,
470 length: self.length * size_of::<T>(),
471 alignment: self.alignment,
472 _marker: Default::default(),
473 }
474 }
475
476 pub fn try_into_mut(self) -> Result<BufferMut<T>, Self> {
478 self.bytes
479 .try_into_mut()
480 .map(|bytes| BufferMut {
481 bytes,
482 length: self.length,
483 alignment: self.alignment,
484 _marker: Default::default(),
485 })
486 .map_err(|bytes| Self {
487 bytes,
488 length: self.length,
489 alignment: self.alignment,
490 _marker: Default::default(),
491 })
492 }
493
494 pub fn into_mut(self) -> BufferMut<T> {
496 self.try_into_mut()
497 .unwrap_or_else(|buffer| BufferMut::<T>::copy_from_aligned(&buffer, buffer.alignment))
498 }
499
500 pub fn is_aligned(&self, alignment: Alignment) -> bool {
502 alignment.is_ptr_aligned(self.bytes.as_ptr())
503 }
504
505 pub fn aligned(mut self, alignment: Alignment) -> Self {
507 if alignment.is_ptr_aligned(self.as_ptr()) {
508 self.alignment = alignment;
509 self
510 } else {
511 #[cfg(feature = "warn-copy")]
512 {
513 let bt = std::backtrace::Backtrace::capture();
514 tracing::warn!(
515 "Buffer is not aligned to requested alignment {alignment}, copying: {bt}"
516 )
517 }
518 Self::copy_from_aligned(self, alignment)
519 }
520 }
521
522 pub fn ensure_aligned(mut self, alignment: Alignment) -> Self {
524 if alignment.is_ptr_aligned(self.as_ptr()) {
525 self.alignment = alignment;
526 self
527 } else {
528 vortex_panic!("Buffer is not aligned to requested alignment {}", alignment)
529 }
530 }
531}
532
533impl<T> Buffer<T> {
534 pub unsafe fn transmute<U>(self) -> Buffer<U> {
546 assert_eq!(size_of::<T>(), size_of::<U>(), "Buffer type size mismatch");
547 assert_eq!(
548 align_of::<T>(),
549 align_of::<U>(),
550 "Buffer type alignment mismatch"
551 );
552
553 Buffer {
554 bytes: self.bytes,
555 length: self.length,
556 alignment: self.alignment,
557 _marker: PhantomData,
558 }
559 }
560}
561
562pub struct Iter<'a, T> {
566 inner: std::slice::Iter<'a, T>,
567}
568
569impl<'a, T> Iterator for Iter<'a, T> {
570 type Item = &'a T;
571
572 #[inline]
573 fn next(&mut self) -> Option<Self::Item> {
574 self.inner.next()
575 }
576
577 #[inline]
578 fn size_hint(&self) -> (usize, Option<usize>) {
579 self.inner.size_hint()
580 }
581
582 #[inline]
583 fn count(self) -> usize {
584 self.inner.count()
585 }
586
587 #[inline]
588 fn last(self) -> Option<Self::Item> {
589 self.inner.last()
590 }
591
592 #[inline]
593 fn nth(&mut self, n: usize) -> Option<Self::Item> {
594 self.inner.nth(n)
595 }
596}
597
598impl<T> ExactSizeIterator for Iter<'_, T> {
599 #[inline]
600 fn len(&self) -> usize {
601 self.inner.len()
602 }
603}
604
605impl<T: Debug> Debug for Buffer<T> {
606 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
607 f.debug_struct(&format!("Buffer<{}>", type_name::<T>()))
608 .field("length", &self.length)
609 .field("alignment", &self.alignment)
610 .field("as_slice", &TruncatedDebug(self.as_slice()))
611 .finish()
612 }
613}
614
615impl<T> Deref for Buffer<T> {
616 type Target = [T];
617
618 #[inline]
619 fn deref(&self) -> &Self::Target {
620 self.as_slice()
621 }
622}
623
624impl<T> AsRef<[T]> for Buffer<T> {
625 #[inline]
626 fn as_ref(&self) -> &[T] {
627 self.as_slice()
628 }
629}
630
631impl<T> FromIterator<T> for Buffer<T> {
632 #[inline]
633 fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
634 BufferMut::from_iter(iter).freeze()
635 }
636}
637
638#[repr(transparent)]
640struct Wrapper<T>(Vec<T>);
641
642impl<T> AsRef<[u8]> for Wrapper<T> {
643 fn as_ref(&self) -> &[u8] {
644 let data = self.0.as_ptr().cast::<u8>();
645 let len = self.0.len() * size_of::<T>();
646 unsafe { std::slice::from_raw_parts(data, len) }
647 }
648}
649
650impl<T> From<Vec<T>> for Buffer<T>
651where
652 T: Send + 'static,
653{
654 fn from(value: Vec<T>) -> Self {
655 let original_len = value.len();
656 let wrapped_vec = Wrapper(value);
657
658 let bytes = Bytes::from_owner(wrapped_vec);
659
660 assert_eq!(bytes.as_ptr().align_offset(align_of::<T>()), 0);
661
662 Self {
663 bytes,
664 length: original_len,
665 alignment: Alignment::of::<T>(),
666 _marker: PhantomData,
667 }
668 }
669}
670
671impl From<Bytes> for ByteBuffer {
672 fn from(bytes: Bytes) -> Self {
673 let length = bytes.len();
674 Self {
675 bytes,
676 length,
677 alignment: Alignment::of::<u8>(),
678 _marker: Default::default(),
679 }
680 }
681}
682
683impl Buf for ByteBuffer {
684 #[inline]
685 fn remaining(&self) -> usize {
686 self.len()
687 }
688
689 #[inline]
690 fn chunk(&self) -> &[u8] {
691 self.as_slice()
692 }
693
694 #[inline]
695 fn advance(&mut self, cnt: usize) {
696 if !self.alignment.is_offset_aligned(cnt) {
697 vortex_panic!(
698 "Cannot advance buffer by {} items, resulting alignment is not {}",
699 cnt,
700 self.alignment
701 );
702 }
703 self.bytes.advance(cnt);
704 self.length -= cnt;
705 }
706}
707
708pub struct BufferIterator<T: Copy> {
710 _buffer: Buffer<T>,
712 ptr: *const T,
713 end: *const T,
714}
715
716impl<T: Copy> Iterator for BufferIterator<T> {
717 type Item = T;
718
719 #[inline]
720 fn next(&mut self) -> Option<Self::Item> {
721 if self.ptr == self.end {
722 None
723 } else {
724 let value = unsafe { self.ptr.read() };
726 self.ptr = unsafe { self.ptr.add(1) };
727 Some(value)
728 }
729 }
730
731 #[inline]
732 fn size_hint(&self) -> (usize, Option<usize>) {
733 let remaining = unsafe { self.end.offset_from(self.ptr) } as usize;
734 (remaining, Some(remaining))
735 }
736}
737
738impl<T: Copy> ExactSizeIterator for BufferIterator<T> {}
739
740impl<T: Copy> IntoIterator for Buffer<T> {
741 type Item = T;
742 type IntoIter = BufferIterator<T>;
743
744 #[inline]
745 fn into_iter(self) -> Self::IntoIter {
746 let ptr = self.as_slice().as_ptr();
747 let end = unsafe { ptr.add(self.len()) };
748 BufferIterator {
749 _buffer: self,
750 ptr,
751 end,
752 }
753 }
754}
755
756impl<T> From<BufferMut<T>> for Buffer<T> {
757 #[inline]
758 fn from(value: BufferMut<T>) -> Self {
759 value.freeze()
760 }
761}
762
763#[cfg(test)]
764mod test {
765 use bytes::Buf;
766
767 use crate::Alignment;
768 use crate::Buffer;
769 use crate::ByteBuffer;
770 use crate::buffer;
771
772 #[test]
773 fn align() {
774 let buf = buffer![0u8, 1, 2];
775 let aligned = buf.aligned(Alignment::new(32));
776 assert_eq!(aligned.alignment(), Alignment::new(32));
777 assert_eq!(aligned.as_slice(), &[0, 1, 2]);
778 }
779
780 #[test]
781 fn slice() {
782 let buf = buffer![0, 1, 2, 3, 4];
783 assert_eq!(buf.slice(1..3).as_slice(), &[1, 2]);
784 assert_eq!(buf.slice(1..=3).as_slice(), &[1, 2, 3]);
785 }
786
787 #[test]
788 fn slice_unaligned() {
789 let buf = buffer![0i32, 1, 2, 3, 4].into_byte_buffer();
790 let sliced = buf.slice_unaligned(1..2);
792 assert_eq!(sliced.len(), 1);
794 assert_eq!(sliced.as_slice(), &[0]);
797 }
798
799 #[test]
800 #[should_panic]
801 fn slice_bad_alignment() {
802 let buf = buffer![0i32, 1, 2, 3, 4].into_byte_buffer();
803 buf.slice(1..2);
805 }
806
807 #[test]
808 fn bytes_buf() {
809 let mut buf = ByteBuffer::copy_from("helloworld".as_bytes());
810 assert_eq!(buf.remaining(), 10);
811 assert_eq!(buf.chunk(), b"helloworld");
812
813 buf.advance(5);
814 assert_eq!(buf.remaining(), 5);
815 assert_eq!(buf.as_slice(), b"world");
816 assert_eq!(buf.chunk(), b"world");
817 }
818
819 #[test]
820 fn buffer_zeroed() {
821 const LEN: usize = 17;
822
823 let buf = Buffer::<u32>::zeroed(LEN);
824
825 assert!(buf.is_aligned(Alignment::of::<u32>()));
826 assert_eq!(buf.as_slice(), &[0; LEN]);
827 }
828
829 #[test]
830 fn buffer_zeroed_aligned() {
831 const LEN: usize = 17;
832 let alignment = Alignment::new(64);
833
834 let buf = Buffer::<u32>::zeroed_aligned(LEN, alignment);
835
836 assert!(buf.is_aligned(alignment));
837 assert_eq!(buf.as_slice(), &[0; LEN]);
838 }
839
840 #[test]
841 fn copy_from_over_aligns_to_default() {
842 let values = [1u32, 2, 3];
843 let buf = Buffer::<u32>::copy_from(values);
844
845 assert_eq!(buf.alignment(), Alignment::of::<u32>());
847 assert!(buf.is_aligned(Alignment::DEFAULT_ALIGNMENT));
849 assert_eq!(buf.as_slice(), &values);
850 }
851
852 #[test]
853 fn zeroed_over_aligns_to_default() {
854 const LEN: usize = 17;
855
856 let buf = Buffer::<u32>::zeroed(LEN);
857
858 assert_eq!(buf.alignment(), Alignment::of::<u32>());
859 assert!(buf.is_aligned(Alignment::DEFAULT_ALIGNMENT));
860 assert_eq!(buf.as_slice(), &[0; LEN]);
861 }
862
863 #[test]
864 fn from_vec() {
865 let vec = vec![1, 2, 3, 4, 5];
866 let buff = Buffer::from(vec.clone());
867 assert!(buff.is_aligned(Alignment::of::<i32>()));
868 assert_eq!(vec, buff.as_ref());
869 }
870
871 #[test]
872 fn empty_aligned_max_alignment() {
873 let buf = Buffer::<u8>::empty_aligned(Alignment::MAX);
875 assert!(buf.is_empty());
876 assert!(buf.is_aligned(Alignment::MAX));
877 }
878
879 #[test]
880 fn empty_slice_preserves_alignment() {
881 let buf = Buffer::<u64>::zeroed_aligned(8, Alignment::new(64));
882 let sliced = buf.slice(0..0);
883 assert!(sliced.is_empty());
884 assert_eq!(sliced.alignment(), Alignment::new(64));
885 assert!(sliced.is_aligned(Alignment::new(64)));
886 }
887
888 #[test]
889 fn empty_into_mut_preserves_alignment() {
890 let buf = Buffer::<u8>::empty_aligned(Alignment::new(64));
891 let buf_mut = buf.into_mut();
892 assert_eq!(buf_mut.alignment(), Alignment::new(64));
893 assert!(buf_mut.is_empty());
894 }
895
896 #[test]
897 fn test_slice_unaligned_end_pos() {
898 let data = vec![0u8; 2];
899 let aligned_buffer = Buffer::copy_from_aligned(&data, Alignment::new(8));
901 aligned_buffer.slice(0..1);
906 }
907
908 #[test]
909 fn test_empty_equality() {
910 let a = Buffer::<u16>::empty();
911 let b = Buffer::<u16>::empty();
912
913 assert_eq!(a, b);
914 }
915}