vortex_io/
io_buf.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Provides types that can be used by I/O frameworks to work with byte buffer-shaped data.
5
6use std::ops::Range;
7
8use bytes::Bytes;
9use vortex_buffer::{Buffer, ConstByteBuffer};
10
11/// Trait for types that can provide a readonly byte buffer interface to I/O frameworks.
12///
13/// # Safety
14/// The type must support contiguous raw memory access via pointer, such as `Vec` or `[u8]`.
15pub unsafe trait IoBuf: Unpin + 'static {
16    /// Returns a raw pointer to the vector’s buffer.
17    fn read_ptr(&self) -> *const u8;
18
19    /// Number of initialized bytes.
20    fn bytes_init(&self) -> usize;
21
22    /// Access the buffer as a byte slice
23    fn as_slice(&self) -> &[u8];
24
25    /// Access the buffer as a byte slice with begin and end indices
26    #[inline]
27    fn slice_owned(self, range: Range<usize>) -> OwnedSlice<Self>
28    where
29        Self: Sized,
30    {
31        OwnedSlice {
32            buf: self,
33            begin: range.start,
34            end: range.end,
35        }
36    }
37}
38
39/// An owned view into a contiguous sequence of bytes.
40pub struct OwnedSlice<T> {
41    buf: T,
42    begin: usize,
43    end: usize,
44}
45
46impl<T> OwnedSlice<T> {
47    /// Unwrap the slice into its underlying type.
48    pub fn into_inner(self) -> T {
49        self.buf
50    }
51}
52
53unsafe impl IoBuf for &'static [u8] {
54    #[inline]
55    fn read_ptr(&self) -> *const u8 {
56        self.as_ptr()
57    }
58
59    #[inline]
60    fn bytes_init(&self) -> usize {
61        self.len()
62    }
63
64    #[inline]
65    fn as_slice(&self) -> &[u8] {
66        self
67    }
68}
69
70unsafe impl<const N: usize> IoBuf for [u8; N] {
71    #[inline]
72    fn read_ptr(&self) -> *const u8 {
73        self.as_ptr()
74    }
75
76    #[inline]
77    fn bytes_init(&self) -> usize {
78        N
79    }
80
81    #[inline]
82    fn as_slice(&self) -> &[u8] {
83        self.as_ref()
84    }
85}
86
87unsafe impl IoBuf for Vec<u8> {
88    #[inline]
89    fn read_ptr(&self) -> *const u8 {
90        self.as_ptr()
91    }
92
93    #[inline]
94    fn bytes_init(&self) -> usize {
95        self.len()
96    }
97
98    #[inline]
99    fn as_slice(&self) -> &[u8] {
100        self.as_ref()
101    }
102}
103
104unsafe impl<T: IoBuf> IoBuf for OwnedSlice<T> {
105    #[inline]
106    fn read_ptr(&self) -> *const u8 {
107        unsafe { self.buf.read_ptr().add(self.begin) }
108    }
109
110    #[inline]
111    fn bytes_init(&self) -> usize {
112        self.end - self.begin
113    }
114
115    #[inline]
116    fn as_slice(&self) -> &[u8] {
117        unsafe { std::slice::from_raw_parts(self.read_ptr(), self.bytes_init()) }
118    }
119}
120
121unsafe impl IoBuf for Bytes {
122    fn read_ptr(&self) -> *const u8 {
123        self.as_ptr()
124    }
125
126    fn bytes_init(&self) -> usize {
127        self.len()
128    }
129
130    fn as_slice(&self) -> &[u8] {
131        self.as_ref()
132    }
133}
134
135unsafe impl<const A: usize> IoBuf for ConstByteBuffer<A> {
136    fn read_ptr(&self) -> *const u8 {
137        self.as_ptr()
138    }
139
140    fn bytes_init(&self) -> usize {
141        self.len()
142    }
143
144    fn as_slice(&self) -> &[u8] {
145        self.as_ref()
146    }
147}
148
149unsafe impl<T: Unpin + 'static> IoBuf for Buffer<T> {
150    fn read_ptr(&self) -> *const u8 {
151        self.as_ptr().cast()
152    }
153
154    fn bytes_init(&self) -> usize {
155        self.len() * size_of::<T>()
156    }
157
158    fn as_slice(&self) -> &[u8] {
159        unsafe { std::slice::from_raw_parts(self.read_ptr(), self.bytes_init()) }
160    }
161}