Skip to content

Commit 341039a

Browse files
committed
Optimize buffer ops
Signed-off-by: Adam Gutglick <adam@spiraldb.com>
1 parent 4b6c382 commit 341039a

3 files changed

Lines changed: 90 additions & 18 deletions

File tree

vortex-buffer/src/alignment.rs

Lines changed: 29 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,9 @@ impl Alignment {
5454
Self::new(align_of::<T>())
5555
}
5656

57+
/// The largest valid alignment: the greatest power of 2 that fits into a `u16`.
58+
pub const MAX: Alignment = Alignment::new(1 << 15);
59+
5760
/// Check if `self` alignment is a "larger" than `other` alignment.
5861
///
5962
/// ## Example
@@ -67,10 +70,32 @@ impl Alignment {
6770
/// assert!(!b.is_aligned_to(a));
6871
/// ```
6972
#[inline]
70-
pub fn is_aligned_to(&self, other: Alignment) -> bool {
71-
// Since we know alignments are powers of 2, we can compare them by checking if the number
72-
// of trailing zeros in the binary representation of the alignment is greater or equal.
73-
self.0.trailing_zeros() >= other.0.trailing_zeros()
73+
pub const fn is_aligned_to(&self, other: Alignment) -> bool {
74+
// Since both alignments are powers of 2, divisibility is equivalent to ordering.
75+
self.0 >= other.0
76+
}
77+
78+
/// Check if the given byte offset (or length) is a multiple of this alignment.
79+
///
80+
/// ## Example
81+
///
82+
/// ```
83+
/// use vortex_buffer::Alignment;
84+
///
85+
/// let a = Alignment::new(4);
86+
/// assert!(a.is_offset_aligned(8));
87+
/// assert!(!a.is_offset_aligned(2));
88+
/// ```
89+
#[inline]
90+
pub const fn is_offset_aligned(&self, offset: usize) -> bool {
91+
// Alignment is always a power of 2, so a mask test is equivalent to `offset % self == 0`.
92+
offset & (self.0 - 1) == 0
93+
}
94+
95+
/// Check if the given pointer is aligned to this alignment.
96+
#[inline]
97+
pub fn is_ptr_aligned<T>(&self, ptr: *const T) -> bool {
98+
self.is_offset_aligned(ptr.addr())
7499
}
75100

76101
/// Returns the log2 of the alignment.

vortex-buffer/src/buffer.rs

Lines changed: 58 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -32,10 +32,17 @@ pub struct Buffer<T> {
3232
pub(crate) _marker: PhantomData<T>,
3333
}
3434

35+
/// Zero-length backing memory for empty buffers, aligned to [`Alignment::MAX`] so it satisfies
36+
/// any valid alignment without allocating.
37+
#[repr(align(32768))]
38+
struct EmptyBacking([u8; 0]);
39+
40+
static EMPTY_BACKING: EmptyBacking = EmptyBacking([]);
41+
3542
impl<T> Default for Buffer<T> {
3643
fn default() -> Self {
3744
Self {
38-
bytes: Default::default(),
45+
bytes: Bytes::from_static(&EMPTY_BACKING.0),
3946
length: 0,
4047
alignment: Alignment::of::<T>(),
4148
_marker: PhantomData,
@@ -101,12 +108,27 @@ impl<T> Buffer<T> {
101108

102109
/// Create a new empty `ByteBuffer` with the provided alignment.
103110
pub fn empty() -> Self {
104-
BufferMut::empty().freeze()
111+
Self::empty_aligned(Alignment::of::<T>())
105112
}
106113

107114
/// Create a new empty `ByteBuffer` with the provided alignment.
115+
///
116+
/// This does not allocate: empty buffers are backed by a shared static allocation that is
117+
/// aligned to [`Alignment::MAX`].
108118
pub fn empty_aligned(alignment: Alignment) -> Self {
109-
BufferMut::empty_aligned(alignment).freeze()
119+
if !alignment.is_aligned_to(Alignment::of::<T>()) {
120+
vortex_panic!(
121+
"Alignment {} must align to the scalar type's alignment {}",
122+
alignment,
123+
Alignment::of::<T>(),
124+
);
125+
}
126+
Self {
127+
bytes: Bytes::from_static(&EMPTY_BACKING.0),
128+
length: 0,
129+
alignment,
130+
_marker: PhantomData,
131+
}
110132
}
111133

112134
/// Create a new full `ByteBuffer` with the given value.
@@ -152,7 +174,7 @@ impl<T> Buffer<T> {
152174
Alignment::of::<T>(),
153175
);
154176
}
155-
if bytes.as_ptr().align_offset(*alignment) != 0 {
177+
if !alignment.is_ptr_aligned(bytes.as_ptr()) {
156178
vortex_panic!(
157179
"Bytes alignment must align to the requested alignment {}",
158180
alignment,
@@ -320,7 +342,7 @@ impl<T> Buffer<T> {
320342
let begin_byte = begin * size_of::<T>();
321343
let end_byte = end * size_of::<T>();
322344

323-
if !begin_byte.is_multiple_of(*alignment) {
345+
if !alignment.is_offset_aligned(begin_byte) {
324346
vortex_panic!(
325347
"range start must be aligned to {alignment:?}, byte {}",
326348
begin_byte
@@ -369,7 +391,7 @@ impl<T> Buffer<T> {
369391
vortex_panic!("slice_ref subset alignment must at least align to the buffer alignment")
370392
}
371393

372-
if subset.as_ptr().align_offset(*alignment) != 0 {
394+
if !alignment.is_ptr_aligned(subset.as_ptr()) {
373395
vortex_panic!("slice_ref subset must be aligned to {:?}", alignment);
374396
}
375397

@@ -435,17 +457,17 @@ impl<T> Buffer<T> {
435457
/// Convert self into `BufferMut<T>`, cloning the data if there are multiple strong references.
436458
pub fn into_mut(self) -> BufferMut<T> {
437459
self.try_into_mut()
438-
.unwrap_or_else(|buffer| BufferMut::<T>::copy_from(&buffer))
460+
.unwrap_or_else(|buffer| BufferMut::<T>::copy_from_aligned(&buffer, buffer.alignment))
439461
}
440462

441463
/// Returns whether a `Buffer<T>` is aligned to the given alignment.
442464
pub fn is_aligned(&self, alignment: Alignment) -> bool {
443-
self.bytes.as_ptr().align_offset(*alignment) == 0
465+
alignment.is_ptr_aligned(self.bytes.as_ptr())
444466
}
445467

446468
/// Return a `Buffer<T>` with the given alignment. Where possible, this will be zero-copy.
447469
pub fn aligned(mut self, alignment: Alignment) -> Self {
448-
if self.as_ptr().align_offset(*alignment) == 0 {
470+
if alignment.is_ptr_aligned(self.as_ptr()) {
449471
self.alignment = alignment;
450472
self
451473
} else {
@@ -462,7 +484,7 @@ impl<T> Buffer<T> {
462484

463485
/// Return a `Buffer<T>` with the given alignment. Panics if the buffer is not aligned.
464486
pub fn ensure_aligned(mut self, alignment: Alignment) -> Self {
465-
if self.as_ptr().align_offset(*alignment) == 0 {
487+
if alignment.is_ptr_aligned(self.as_ptr()) {
466488
self.alignment = alignment;
467489
self
468490
} else {
@@ -634,7 +656,7 @@ impl Buf for ByteBuffer {
634656

635657
#[inline]
636658
fn advance(&mut self, cnt: usize) {
637-
if !cnt.is_multiple_of(*self.alignment) {
659+
if !self.alignment.is_offset_aligned(cnt) {
638660
vortex_panic!(
639661
"Cannot advance buffer by {} items, resulting alignment is not {}",
640662
cnt,
@@ -786,6 +808,31 @@ mod test {
786808
assert_eq!(vec, buff.as_ref());
787809
}
788810

811+
#[test]
812+
fn empty_aligned_max_alignment() {
813+
// Empty buffers are backed by a static and must satisfy any valid alignment.
814+
let buf = Buffer::<u8>::empty_aligned(Alignment::MAX);
815+
assert!(buf.is_empty());
816+
assert!(buf.is_aligned(Alignment::MAX));
817+
}
818+
819+
#[test]
820+
fn empty_slice_preserves_alignment() {
821+
let buf = Buffer::<u64>::zeroed_aligned(8, Alignment::new(64));
822+
let sliced = buf.slice(0..0);
823+
assert!(sliced.is_empty());
824+
assert_eq!(sliced.alignment(), Alignment::new(64));
825+
assert!(sliced.is_aligned(Alignment::new(64)));
826+
}
827+
828+
#[test]
829+
fn empty_into_mut_preserves_alignment() {
830+
let buf = Buffer::<u8>::empty_aligned(Alignment::new(64));
831+
let buf_mut = buf.into_mut();
832+
assert_eq!(buf_mut.alignment(), Alignment::new(64));
833+
assert!(buf_mut.is_empty());
834+
}
835+
789836
#[test]
790837
fn test_slice_unaligned_end_pos() {
791838
let data = vec![0u8; 2];

vortex-buffer/src/buffer_mut.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -358,7 +358,7 @@ impl<T> BufferMut<T> {
358358
}
359359

360360
let bytes_at = at * size_of::<T>();
361-
if !bytes_at.is_multiple_of(*self.alignment) {
361+
if !self.alignment.is_offset_aligned(bytes_at) {
362362
vortex_panic!(
363363
"Cannot split buffer at {}, resulting alignment is not {}",
364364
at,
@@ -742,7 +742,7 @@ impl Buf for ByteBufferMut {
742742
}
743743

744744
fn advance(&mut self, cnt: usize) {
745-
if !cnt.is_multiple_of(*self.alignment) {
745+
if !self.alignment.is_offset_aligned(cnt) {
746746
vortex_panic!(
747747
"Cannot advance buffer by {} items, resulting alignment is not {}",
748748
cnt,
@@ -765,7 +765,7 @@ unsafe impl BufMut for ByteBufferMut {
765765

766766
#[inline]
767767
unsafe fn advance_mut(&mut self, cnt: usize) {
768-
if !cnt.is_multiple_of(*self.alignment) {
768+
if !self.alignment.is_offset_aligned(cnt) {
769769
vortex_panic!(
770770
"Cannot advance buffer by {} items, resulting alignment is not {}",
771771
cnt,

0 commit comments

Comments
 (0)