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 BufferMut::from_trusted_len_iter(iter).freeze()
240 }
241
242 pub fn map_each_in_place<R, F>(self, mut f: F) -> BufferMut<R>
244 where
245 T: Copy,
246 F: FnMut(T) -> R,
247 {
248 match self.try_into_mut() {
249 Ok(mut_buf) => mut_buf.map_each_in_place(f),
250 Err(buf) => {
251 let len = buf.len();
252 let mut out_buf = BufferMut::with_capacity(len);
253 out_buf
254 .spare_capacity_mut()
255 .iter_mut()
256 .zip(buf)
257 .for_each(|(out, in_)| {
258 out.write(f(in_));
259 });
260 unsafe { out_buf.set_len(len) }
262 out_buf
263 }
264 }
265 }
266
267 pub fn clear(&mut self) {
269 self.bytes.clear();
270 self.length = 0;
271 }
272
273 #[inline(always)]
275 pub fn len(&self) -> usize {
276 self.length
277 }
278
279 #[inline(always)]
281 pub fn is_empty(&self) -> bool {
282 self.length == 0
283 }
284
285 #[inline(always)]
287 pub fn alignment(&self) -> Alignment {
288 self.alignment
289 }
290
291 #[inline(always)]
293 pub fn as_slice(&self) -> &[T] {
294 unsafe { std::slice::from_raw_parts(self.bytes.as_ptr().cast(), self.length) }
296 }
297
298 #[inline(always)]
300 pub fn as_bytes(&self) -> &[u8] {
301 self.bytes.as_ref()
302 }
303
304 pub fn iter(&self) -> Iter<'_, T> {
306 Iter {
307 inner: self.as_slice().iter(),
308 }
309 }
310
311 #[inline(always)]
318 pub fn slice(&self, range: impl RangeBounds<usize>) -> Self {
319 self.slice_with_alignment(range, self.alignment)
320 }
321
322 #[inline(always)]
329 pub fn slice_unaligned(&self, range: impl RangeBounds<usize>) -> Self {
330 self.slice_with_alignment(range, Alignment::of::<u8>())
331 }
332
333 pub fn slice_with_alignment(
341 &self,
342 range: impl RangeBounds<usize>,
343 alignment: Alignment,
344 ) -> Self {
345 let len = self.len();
346 let begin = match range.start_bound() {
347 Bound::Included(&n) => n,
348 Bound::Excluded(&n) => n.checked_add(1).vortex_expect("out of range"),
349 Bound::Unbounded => 0,
350 };
351 let end = match range.end_bound() {
352 Bound::Included(&n) => n.checked_add(1).vortex_expect("out of range"),
353 Bound::Excluded(&n) => n,
354 Bound::Unbounded => len,
355 };
356
357 if begin > end {
358 vortex_panic!(
359 "range start must not be greater than end: {:?} <= {:?}",
360 begin,
361 end
362 );
363 }
364 if end > len {
365 vortex_panic!("range end out of bounds: {:?} > {:?}", end, len);
366 }
367
368 if end == begin {
369 return Self::empty_aligned(alignment);
372 }
373
374 let begin_byte = begin * size_of::<T>();
375 let end_byte = end * size_of::<T>();
376
377 if !alignment.is_offset_aligned(begin_byte) {
378 vortex_panic!(
379 "range start must be aligned to {alignment:?}, byte {}",
380 begin_byte
381 );
382 }
383 if !alignment.is_aligned_to(Alignment::of::<T>()) {
384 vortex_panic!("Slice alignment must at least align to type T")
385 }
386
387 Self {
388 bytes: self.bytes.slice(begin_byte..end_byte),
389 length: end - begin,
390 alignment,
391 _marker: Default::default(),
392 }
393 }
394
395 #[inline(always)]
404 pub fn slice_ref(&self, subset: &[T]) -> Self {
405 self.slice_ref_with_alignment(subset, Alignment::of::<T>())
406 }
407
408 pub fn slice_ref_with_alignment(&self, subset: &[T], alignment: Alignment) -> Self {
418 if !alignment.is_aligned_to(Alignment::of::<T>()) {
419 vortex_panic!("slice_ref alignment must at least align to type T")
420 }
421
422 if !self.alignment.is_aligned_to(alignment) {
423 vortex_panic!("slice_ref subset alignment must at least align to the buffer alignment")
424 }
425
426 if !alignment.is_ptr_aligned(subset.as_ptr()) {
427 vortex_panic!("slice_ref subset must be aligned to {:?}", alignment);
428 }
429
430 let subset_u8 =
431 unsafe { std::slice::from_raw_parts(subset.as_ptr().cast(), size_of_val(subset)) };
432
433 Self {
434 bytes: self.bytes.slice_ref(subset_u8),
435 length: subset.len(),
436 alignment,
437 _marker: Default::default(),
438 }
439 }
440
441 pub fn inner(&self) -> &Bytes {
443 debug_assert_eq!(
444 self.length * size_of::<T>(),
445 self.bytes.len(),
446 "Own length has to be the same as the underlying bytes length"
447 );
448 &self.bytes
449 }
450
451 pub fn into_inner(self) -> Bytes {
453 debug_assert_eq!(
454 self.length * size_of::<T>(),
455 self.bytes.len(),
456 "Own length has to be the same as the underlying bytes length"
457 );
458 self.bytes
459 }
460
461 pub fn into_byte_buffer(self) -> ByteBuffer {
463 ByteBuffer {
464 bytes: self.bytes,
465 length: self.length * size_of::<T>(),
466 alignment: self.alignment,
467 _marker: Default::default(),
468 }
469 }
470
471 pub fn try_into_mut(self) -> Result<BufferMut<T>, Self> {
473 self.bytes
474 .try_into_mut()
475 .map(|bytes| BufferMut {
476 bytes,
477 length: self.length,
478 alignment: self.alignment,
479 _marker: Default::default(),
480 })
481 .map_err(|bytes| Self {
482 bytes,
483 length: self.length,
484 alignment: self.alignment,
485 _marker: Default::default(),
486 })
487 }
488
489 pub fn into_mut(self) -> BufferMut<T> {
491 self.try_into_mut()
492 .unwrap_or_else(|buffer| BufferMut::<T>::copy_from_aligned(&buffer, buffer.alignment))
493 }
494
495 pub fn is_aligned(&self, alignment: Alignment) -> bool {
497 alignment.is_ptr_aligned(self.bytes.as_ptr())
498 }
499
500 pub fn aligned(mut self, alignment: Alignment) -> Self {
502 if alignment.is_ptr_aligned(self.as_ptr()) {
503 self.alignment = alignment;
504 self
505 } else {
506 #[cfg(feature = "warn-copy")]
507 {
508 let bt = std::backtrace::Backtrace::capture();
509 tracing::warn!(
510 "Buffer is not aligned to requested alignment {alignment}, copying: {bt}"
511 )
512 }
513 Self::copy_from_aligned(self, alignment)
514 }
515 }
516
517 pub fn ensure_aligned(mut self, alignment: Alignment) -> Self {
519 if alignment.is_ptr_aligned(self.as_ptr()) {
520 self.alignment = alignment;
521 self
522 } else {
523 vortex_panic!("Buffer is not aligned to requested alignment {}", alignment)
524 }
525 }
526}
527
528impl<T> Buffer<T> {
529 pub unsafe fn transmute<U>(self) -> Buffer<U> {
541 assert_eq!(size_of::<T>(), size_of::<U>(), "Buffer type size mismatch");
542 assert_eq!(
543 align_of::<T>(),
544 align_of::<U>(),
545 "Buffer type alignment mismatch"
546 );
547
548 Buffer {
549 bytes: self.bytes,
550 length: self.length,
551 alignment: self.alignment,
552 _marker: PhantomData,
553 }
554 }
555}
556
557pub struct Iter<'a, T> {
561 inner: std::slice::Iter<'a, T>,
562}
563
564impl<'a, T> Iterator for Iter<'a, T> {
565 type Item = &'a T;
566
567 #[inline]
568 fn next(&mut self) -> Option<Self::Item> {
569 self.inner.next()
570 }
571
572 #[inline]
573 fn size_hint(&self) -> (usize, Option<usize>) {
574 self.inner.size_hint()
575 }
576
577 #[inline]
578 fn count(self) -> usize {
579 self.inner.count()
580 }
581
582 #[inline]
583 fn last(self) -> Option<Self::Item> {
584 self.inner.last()
585 }
586
587 #[inline]
588 fn nth(&mut self, n: usize) -> Option<Self::Item> {
589 self.inner.nth(n)
590 }
591}
592
593impl<T> ExactSizeIterator for Iter<'_, T> {
594 #[inline]
595 fn len(&self) -> usize {
596 self.inner.len()
597 }
598}
599
600impl<T: Debug> Debug for Buffer<T> {
601 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
602 f.debug_struct(&format!("Buffer<{}>", type_name::<T>()))
603 .field("length", &self.length)
604 .field("alignment", &self.alignment)
605 .field("as_slice", &TruncatedDebug(self.as_slice()))
606 .finish()
607 }
608}
609
610impl<T> Deref for Buffer<T> {
611 type Target = [T];
612
613 #[inline]
614 fn deref(&self) -> &Self::Target {
615 self.as_slice()
616 }
617}
618
619impl<T> AsRef<[T]> for Buffer<T> {
620 #[inline]
621 fn as_ref(&self) -> &[T] {
622 self.as_slice()
623 }
624}
625
626impl<T> FromIterator<T> for Buffer<T> {
627 #[inline]
628 fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
629 BufferMut::from_iter(iter).freeze()
630 }
631}
632
633#[repr(transparent)]
635struct Wrapper<T>(Vec<T>);
636
637impl<T> AsRef<[u8]> for Wrapper<T> {
638 fn as_ref(&self) -> &[u8] {
639 let data = self.0.as_ptr().cast::<u8>();
640 let len = self.0.len() * size_of::<T>();
641 unsafe { std::slice::from_raw_parts(data, len) }
642 }
643}
644
645impl<T> From<Vec<T>> for Buffer<T>
646where
647 T: Send + 'static,
648{
649 fn from(value: Vec<T>) -> Self {
650 let original_len = value.len();
651 let wrapped_vec = Wrapper(value);
652
653 let bytes = Bytes::from_owner(wrapped_vec);
654
655 assert_eq!(bytes.as_ptr().align_offset(align_of::<T>()), 0);
656
657 Self {
658 bytes,
659 length: original_len,
660 alignment: Alignment::of::<T>(),
661 _marker: PhantomData,
662 }
663 }
664}
665
666impl From<Bytes> for ByteBuffer {
667 fn from(bytes: Bytes) -> Self {
668 let length = bytes.len();
669 Self {
670 bytes,
671 length,
672 alignment: Alignment::of::<u8>(),
673 _marker: Default::default(),
674 }
675 }
676}
677
678impl Buf for ByteBuffer {
679 #[inline]
680 fn remaining(&self) -> usize {
681 self.len()
682 }
683
684 #[inline]
685 fn chunk(&self) -> &[u8] {
686 self.as_slice()
687 }
688
689 #[inline]
690 fn advance(&mut self, cnt: usize) {
691 if !self.alignment.is_offset_aligned(cnt) {
692 vortex_panic!(
693 "Cannot advance buffer by {} items, resulting alignment is not {}",
694 cnt,
695 self.alignment
696 );
697 }
698 self.bytes.advance(cnt);
699 self.length -= cnt;
700 }
701}
702
703pub struct BufferIterator<T: Copy> {
705 _buffer: Buffer<T>,
707 ptr: *const T,
708 end: *const T,
709}
710
711unsafe impl<T: Copy + Send> Send for BufferIterator<T> {}
714unsafe impl<T: Copy + Sync> Sync for BufferIterator<T> {}
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 buffer_iterator_send_sync() {
782 fn assert_send_sync<T: Send + Sync>(_: &T) {}
783
784 let mut iter = buffer![0i32, 1, 2, 3].into_iter();
785 assert_send_sync(&iter);
786 iter.next();
787 let remaining: Vec<i32> = std::thread::spawn(move || iter.collect()).join().unwrap();
788 assert_eq!(remaining, vec![1, 2, 3]);
789 }
790
791 #[test]
792 fn slice() {
793 let buf = buffer![0, 1, 2, 3, 4];
794 assert_eq!(buf.slice(1..3).as_slice(), &[1, 2]);
795 assert_eq!(buf.slice(1..=3).as_slice(), &[1, 2, 3]);
796 }
797
798 #[test]
799 fn slice_unaligned() {
800 let buf = buffer![0i32, 1, 2, 3, 4].into_byte_buffer();
801 let sliced = buf.slice_unaligned(1..2);
803 assert_eq!(sliced.len(), 1);
805 assert_eq!(sliced.as_slice(), &[0]);
808 }
809
810 #[test]
811 #[should_panic]
812 fn slice_bad_alignment() {
813 let buf = buffer![0i32, 1, 2, 3, 4].into_byte_buffer();
814 buf.slice(1..2);
816 }
817
818 #[test]
819 fn bytes_buf() {
820 let mut buf = ByteBuffer::copy_from("helloworld".as_bytes());
821 assert_eq!(buf.remaining(), 10);
822 assert_eq!(buf.chunk(), b"helloworld");
823
824 buf.advance(5);
825 assert_eq!(buf.remaining(), 5);
826 assert_eq!(buf.as_slice(), b"world");
827 assert_eq!(buf.chunk(), b"world");
828 }
829
830 #[test]
831 fn buffer_zeroed() {
832 const LEN: usize = 17;
833
834 let buf = Buffer::<u32>::zeroed(LEN);
835
836 assert!(buf.is_aligned(Alignment::of::<u32>()));
837 assert_eq!(buf.as_slice(), &[0; LEN]);
838 }
839
840 #[test]
841 fn buffer_zeroed_aligned() {
842 const LEN: usize = 17;
843 let alignment = Alignment::new(64);
844
845 let buf = Buffer::<u32>::zeroed_aligned(LEN, alignment);
846
847 assert!(buf.is_aligned(alignment));
848 assert_eq!(buf.as_slice(), &[0; LEN]);
849 }
850
851 #[test]
852 fn copy_from_over_aligns_to_default() {
853 let values = [1u32, 2, 3];
854 let buf = Buffer::<u32>::copy_from(values);
855
856 assert_eq!(buf.alignment(), Alignment::of::<u32>());
858 assert!(buf.is_aligned(Alignment::DEFAULT_ALIGNMENT));
860 assert_eq!(buf.as_slice(), &values);
861 }
862
863 #[test]
864 fn zeroed_over_aligns_to_default() {
865 const LEN: usize = 17;
866
867 let buf = Buffer::<u32>::zeroed(LEN);
868
869 assert_eq!(buf.alignment(), Alignment::of::<u32>());
870 assert!(buf.is_aligned(Alignment::DEFAULT_ALIGNMENT));
871 assert_eq!(buf.as_slice(), &[0; LEN]);
872 }
873
874 #[test]
875 fn from_vec() {
876 let vec = vec![1, 2, 3, 4, 5];
877 let buff = Buffer::from(vec.clone());
878 assert!(buff.is_aligned(Alignment::of::<i32>()));
879 assert_eq!(vec, buff.as_ref());
880 }
881
882 #[test]
883 fn empty_aligned_max_alignment() {
884 let buf = Buffer::<u8>::empty_aligned(Alignment::MAX);
886 assert!(buf.is_empty());
887 assert!(buf.is_aligned(Alignment::MAX));
888 }
889
890 #[test]
891 fn empty_slice_preserves_alignment() {
892 let buf = Buffer::<u64>::zeroed_aligned(8, Alignment::new(64));
893 let sliced = buf.slice(0..0);
894 assert!(sliced.is_empty());
895 assert_eq!(sliced.alignment(), Alignment::new(64));
896 assert!(sliced.is_aligned(Alignment::new(64)));
897 }
898
899 #[test]
900 fn empty_into_mut_preserves_alignment() {
901 let buf = Buffer::<u8>::empty_aligned(Alignment::new(64));
902 let buf_mut = buf.into_mut();
903 assert_eq!(buf_mut.alignment(), Alignment::new(64));
904 assert!(buf_mut.is_empty());
905 }
906
907 #[test]
908 fn test_slice_unaligned_end_pos() {
909 let data = vec![0u8; 2];
910 let aligned_buffer = Buffer::copy_from_aligned(&data, Alignment::new(8));
912 aligned_buffer.slice(0..1);
917 }
918
919 #[test]
920 fn test_empty_equality() {
921 let a = Buffer::<u16>::empty();
922 let b = Buffer::<u16>::empty();
923
924 assert_eq!(a, b);
925 }
926}