vortex_buffer/
bytes.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use bytes::Buf;
5use vortex_error::VortexExpect;
6
7use crate::Alignment;
8use crate::ByteBuffer;
9use crate::ConstBuffer;
10use crate::ConstByteBuffer;
11
12/// An extension to the [`Buf`] trait that provides a function `copy_to_aligned` similar to
13/// `copy_to_bytes` that allows for zero-copy aligned reads where possible.
14pub trait AlignedBuf: Buf {
15    /// Copy the next `len` bytes from the buffer into a new buffer with the given alignment.
16    /// This will be zero-copy wherever possible.
17    ///
18    /// The [`Buf`] trait has a specialized `copy_to_bytes` function that allows the implementation
19    /// of `Buf` for `Bytes` and `BytesMut` to return bytes with zero-copy.
20    ///
21    /// This function provides similar functionality for `ByteBuffer`.
22    ///
23    /// TODO(ngates): what should this do the alignment of the current buffer? We have to advance
24    ///  it by len..
25    fn copy_to_aligned(&mut self, len: usize, alignment: Alignment) -> ByteBuffer {
26        // The default implementation uses copy_to_bytes, and then tries to align.
27        // When the underlying `copy_to_bytes` is zero-copy, this may perform one copy to align
28        // the bytes. When the underlying `copy_to_bytes` is not zero-copy, this may perform two
29        // copies.
30        // The only way to fix this would be to invert the implementation so `copy_to_bytes`
31        // invokes `copy_to_aligned` with an alignment of 1. But we cannot override this in the
32        // default trait.
33        // In practice, we tend to only call this function on `ByteBuffer: AlignedBuf`, and
34        // therefore we have a maximum of one copy, so I'm not too worried about it.
35        ByteBuffer::from(self.copy_to_bytes(len)).aligned(alignment)
36    }
37
38    /// See [`AlignedBuf::copy_to_aligned`].
39    fn copy_to_const_aligned<const A: usize>(&mut self, len: usize) -> ConstByteBuffer<A> {
40        // The default implementation uses copy_to_bytes, and then returns a ByteBuffer with
41        // alignment of 1. This will be zero-copy if the underlying `copy_to_bytes` is zero-copy.
42        ConstBuffer::try_from(self.copy_to_aligned(len, Alignment::new(A)))
43            .vortex_expect("we just aligned the buffer")
44    }
45}
46
47impl<B: Buf> AlignedBuf for B {}
vortex_buffer/bytes.rs

vortex_buffer/
bytes.rs
