narrow/
buffer.rs

1//! Traits for memory buffers.
2
3use crate::{FixedSize, Index, Length};
4use std::{marker::PhantomData, mem, rc::Rc, slice, sync::Arc};
5
6/// A memory buffer type constructor for Arrow data.
7///
8/// The generic associated type constructor [`Self::Buffer`] defines the
9/// [`Buffer`] type that stores [`FixedSize`] items.
10///
11// note
12// Arrow buffers are like Rust slices with "primitive" item types.
13// Another way to implement the buffer trait: a subtrait of Borrow<[T]> and then
14// implement Buffer<T> for all U: Borrow<[T] where T: FixedSize, however,the approach here is a little
15// bit more elaborate to also support buffer types that don't implement Borrow<[T]>.
16pub trait BufferType {
17    /// A [`Buffer`] type for [`FixedSize`] items of type `T`.
18    type Buffer<T: FixedSize>: Buffer<T>;
19}
20
21/// An immutable reference to a buffer.
22///
23/// This can be used to provide immutable access to an internal buffer.
24pub trait BufferRef<T: FixedSize> {
25    /// The [Buffer] type.
26    type Buffer: Buffer<T>;
27
28    /// Returns an immutable reference to a buffer.
29    fn buffer_ref(&self) -> &Self::Buffer;
30}
31
32/// A mutable reference to a buffer.
33///
34/// This can be used to provide mutable access to an internal buffer.
35pub trait BufferRefMut<T: FixedSize> {
36    /// The [`BufferMut`] type.
37    type BufferMut: BufferMut<T>;
38
39    /// Returns a mutable reference to a buffer.
40    fn buffer_ref_mut(&mut self) -> &mut Self::BufferMut;
41}
42
43/// A contiguous immutable memory buffer for Arrow data.
44pub trait Buffer<T: FixedSize>: Index + Length {
45    /// Extracts a slice containing the entire buffer.
46    fn as_slice(&self) -> &[T];
47
48    /// Returns the contents of the entire buffer as a byte slice.
49    fn as_bytes(&self) -> &[u8] {
50        // Safety:
51        // - The pointer returned by slice::as_ptr (via Borrow) points to slice::len()
52        //   consecutive properly initialized values of type T, with size_of::<T> bytes
53        //   per element.
54        unsafe {
55            slice::from_raw_parts(
56                self.as_slice().as_ptr().cast(),
57                mem::size_of_val(self.as_slice()),
58            )
59        }
60    }
61}
62
63/// A contiguous mutable memory buffer for Arrow data.
64pub trait BufferMut<T: FixedSize>: Buffer<T> {
65    /// Extracts a mutable slice containing the entire buffer.
66    fn as_mut_slice(&mut self) -> &mut [T];
67
68    /// Returns the contents of the entire buffer as a mutable byte slice.
69    fn as_mut_bytes(&mut self) -> &mut [u8] {
70        // Safety:
71        // - The pointer returned by slice::as_mut_ptr (via Borrow) points to slice::len()
72        //   consecutive properly initialized values of type T, with size_of::<T> bytes
73        //   per element.
74        unsafe {
75            slice::from_raw_parts_mut(
76                self.as_mut_slice().as_mut_ptr().cast(),
77                mem::size_of_val(self.as_slice()),
78            )
79        }
80    }
81}
82
83/// A [`BufferType`] for a single item.
84#[derive(Clone, Copy, Debug)]
85pub struct SingleBuffer;
86
87impl BufferType for SingleBuffer {
88    type Buffer<T: FixedSize> = <ArrayBuffer<1> as BufferType>::Buffer<T>;
89}
90
91/// A [`BufferType`] implementation for array.
92///
93/// Stores items `T` in `[T; N]`.
94#[derive(Clone, Copy, Debug)]
95pub struct ArrayBuffer<const N: usize>;
96
97impl<const N: usize> BufferType for ArrayBuffer<N> {
98    type Buffer<T: FixedSize> = [T; N];
99}
100
101impl<T: FixedSize, const N: usize> Buffer<T> for [T; N] {
102    fn as_slice(&self) -> &[T] {
103        self.as_slice()
104    }
105}
106
107impl<T: FixedSize, const N: usize> BufferMut<T> for [T; N] {
108    fn as_mut_slice(&mut self) -> &mut [T] {
109        self.as_mut_slice()
110    }
111}
112
113/// A [`BufferType`] implementation for array in array.
114///
115/// Stores items `T` in `[[T; M]; N]`.
116#[derive(Clone, Copy, Debug)]
117pub struct ArrayArrayBuffer<const M: usize, const N: usize>;
118
119impl<const M: usize, const N: usize> BufferType for ArrayArrayBuffer<M, N> {
120    type Buffer<T: FixedSize> = [[T; M]; N];
121}
122
123impl<T: FixedSize, const M: usize, const N: usize> Buffer<T> for [[T; M]; N] {
124    fn as_slice(&self) -> &[T] {
125        // self.flatten() is nightly
126        // SAFETY: `[T]` is layout-identical to `[[T; M]; N]`
127        unsafe { std::slice::from_raw_parts(self.as_ptr().cast(), M * N) }
128    }
129}
130
131impl<T: FixedSize, const M: usize, const N: usize> BufferMut<T> for [[T; M]; N] {
132    fn as_mut_slice(&mut self) -> &mut [T] {
133        // self.flatten() is nightly
134        // SAFETY: `[T]` is layout-identical to `[[T; M]; N]`
135        unsafe { std::slice::from_raw_parts_mut(self.as_mut_ptr().cast(), M * N) }
136    }
137}
138
139/// A [`BufferType`] implementation for slice.
140///
141/// Stores items `T` in `&[T]`.
142#[derive(Clone, Copy, Debug)]
143pub struct SliceBuffer<'a>(PhantomData<&'a ()>);
144
145impl<'a> BufferType for SliceBuffer<'a> {
146    type Buffer<T: FixedSize> = &'a [T];
147}
148
149impl<T: FixedSize> Buffer<T> for &[T] {
150    fn as_slice(&self) -> &[T] {
151        self
152    }
153}
154
155/// A [`BufferType`] implementation for mutable slice.
156///
157/// Stores items `T` in `&mut [T]`.
158#[derive(Clone, Copy, Debug)]
159pub struct SliceMutBuffer<'a>(PhantomData<&'a ()>);
160
161impl<'a> BufferType for SliceMutBuffer<'a> {
162    type Buffer<T: FixedSize> = &'a mut [T];
163}
164
165impl<T: FixedSize> Buffer<T> for &mut [T] {
166    fn as_slice(&self) -> &[T] {
167        self
168    }
169}
170
171impl<T: FixedSize> BufferMut<T> for &mut [T] {
172    fn as_mut_slice(&mut self) -> &mut [T] {
173        self
174    }
175}
176
177/// A [`BufferType`] implementation for slice with array items.
178///
179/// Stores items `T` in `&[[T; N]]`.
180#[derive(Clone, Copy, Debug)]
181pub struct SliceArrayBuffer<'a, const N: usize>(PhantomData<&'a ()>);
182
183impl<'a, const N: usize> BufferType for SliceArrayBuffer<'a, N> {
184    type Buffer<T: FixedSize> = &'a [[T; N]];
185}
186
187impl<T: FixedSize, const N: usize> Buffer<T> for &[[T; N]] {
188    fn as_slice(&self) -> &[T] {
189        // self.flatten() is nightly
190        // SAFETY: `[T]` is layout-identical to `[T; N]`
191        unsafe { std::slice::from_raw_parts(self.as_ptr().cast(), <[[T; N]]>::len(self) * N) }
192    }
193}
194
195/// A [`BufferType`] implementation for mutable slice with array items.
196///
197/// Stores items `T` in `&mut [[T; N]]`.
198#[derive(Clone, Copy, Debug)]
199pub struct SliceArrayMutBuffer<'a, const N: usize>(PhantomData<&'a ()>);
200
201impl<'a, const N: usize> BufferType for SliceArrayMutBuffer<'a, N> {
202    type Buffer<T: FixedSize> = &'a mut [[T; N]];
203}
204
205impl<T: FixedSize, const N: usize> Buffer<T> for &mut [[T; N]] {
206    fn as_slice(&self) -> &[T] {
207        // self.flatten() is nightly
208        // SAFETY: `[T]` is layout-identical to `[T; N]`
209        unsafe { std::slice::from_raw_parts(self.as_ptr().cast(), <[[T; N]]>::len(self) * N) }
210    }
211}
212
213impl<T: FixedSize, const N: usize> BufferMut<T> for &mut [[T; N]] {
214    fn as_mut_slice(&mut self) -> &mut [T] {
215        // self.flatten() is nightly
216        // SAFETY: `[T]` is layout-identical to `[T; N]`
217        unsafe {
218            std::slice::from_raw_parts_mut(self.as_mut_ptr().cast(), <[[T; N]]>::len(self) * N)
219        }
220    }
221}
222
223/// A [`BufferType`] implementation for [`Vec`].
224///
225/// Stores items `T` in `Vec<T>`.
226#[derive(Clone, Copy, Debug)]
227pub struct VecBuffer;
228
229impl BufferType for VecBuffer {
230    type Buffer<T: FixedSize> = Vec<T>;
231}
232
233impl<T: FixedSize> Buffer<T> for Vec<T> {
234    fn as_slice(&self) -> &[T] {
235        self.as_slice()
236    }
237}
238
239impl<T: FixedSize> BufferMut<T> for Vec<T> {
240    fn as_mut_slice(&mut self) -> &mut [T] {
241        self.as_mut_slice()
242    }
243}
244
245/// A [`BufferType`] implementation for [`Vec`] with array items.
246///
247/// Stores items `T` in `Vec<[T;N]>`.
248#[derive(Clone, Copy, Debug)]
249pub struct VecArrayBuffer<const N: usize>;
250
251impl<const N: usize> BufferType for VecArrayBuffer<N> {
252    type Buffer<T: FixedSize> = Vec<[T; N]>;
253}
254
255impl<T: FixedSize, const N: usize> Buffer<T> for Vec<[T; N]> {
256    fn as_slice(&self) -> &[T] {
257        // self.flatten() is nightly
258        // SAFETY: `[T]` is layout-identical to `[T; N]`
259        unsafe { std::slice::from_raw_parts(self.as_ptr().cast(), Vec::<[T; N]>::len(self) * N) }
260    }
261}
262
263impl<T: FixedSize, const N: usize> BufferMut<T> for Vec<[T; N]> {
264    fn as_mut_slice(&mut self) -> &mut [T] {
265        // self.flatten() is nightly
266        // SAFETY: `[T]` is layout-identical to `[T; N]`
267        unsafe {
268            std::slice::from_raw_parts_mut(self.as_mut_ptr().cast(), Vec::<[T; N]>::len(self) * N)
269        }
270    }
271}
272
273/// A [`BufferType`] implementation for [`Box`].
274///
275/// Stores items `T` in `Box<[T]>`.
276#[derive(Clone, Copy, Debug)]
277pub struct BoxBuffer;
278
279impl BufferType for BoxBuffer {
280    type Buffer<T: FixedSize> = Box<[T]>;
281}
282
283impl<T: FixedSize> Buffer<T> for Box<[T]> {
284    fn as_slice(&self) -> &[T] {
285        <&[T]>::from(self)
286    }
287}
288
289impl<T: FixedSize> BufferMut<T> for Box<[T]> {
290    fn as_mut_slice(&mut self) -> &mut [T] {
291        <&mut [T]>::from(self)
292    }
293}
294
295/// A [`BufferType`] implementation for [`Arc`].
296///
297/// Stores items `T` in `Arc<[T]>`.
298#[derive(Clone, Copy, Debug)]
299pub struct ArcBuffer;
300
301impl BufferType for ArcBuffer {
302    type Buffer<T: FixedSize> = Arc<[T]>;
303}
304
305impl<T: FixedSize> Buffer<T> for Arc<[T]> {
306    fn as_slice(&self) -> &[T] {
307        <&[T]>::from(self)
308    }
309}
310
311impl<T: FixedSize> BufferMut<T> for Arc<[T]> {
312    fn as_mut_slice(&mut self) -> &mut [T] {
313        match Arc::get_mut(self) {
314            Some(slice) => slice,
315            None => panic!("not safe to mutate shared value"),
316        }
317    }
318}
319
320/// A [`BufferType`] implementation for [`Rc`].
321///
322/// Stores items `T` in `Rc<[T]>`.
323#[derive(Clone, Copy, Debug)]
324pub struct RcBuffer;
325
326impl BufferType for RcBuffer {
327    type Buffer<T: FixedSize> = Rc<[T]>;
328}
329
330impl<T: FixedSize> Buffer<T> for Rc<[T]> {
331    fn as_slice(&self) -> &[T] {
332        <&[T]>::from(self)
333    }
334}
335
336impl<T: FixedSize> BufferMut<T> for Rc<[T]> {
337    fn as_mut_slice(&mut self) -> &mut [T] {
338        match Rc::get_mut(self) {
339            Some(slice) => slice,
340            None => panic!("not safe to mutate shared value"),
341        }
342    }
343}
344
345#[cfg(test)]
346mod tests {
347    use super::*;
348
349    #[test]
350    fn single() {
351        let mut single: <SingleBuffer as BufferType>::Buffer<u16> = [1234];
352        assert_eq!(single.as_bytes(), [210, 4]);
353        single.as_mut_bytes()[1] = 0;
354        assert_eq!(single.as_bytes(), [210, 0]);
355        single.as_mut_slice()[0] = 1234;
356        assert_eq!(single, [1234]);
357    }
358
359    #[test]
360    fn array() {
361        let mut array: <ArrayBuffer<4> as BufferType>::Buffer<u16> = [1, 2, 3, 4];
362        assert_eq!(
363            <_ as Buffer<u16>>::as_bytes(&array),
364            &[1, 0, 2, 0, 3, 0, 4, 0]
365        );
366        <_ as BufferMut<u16>>::as_mut_bytes(&mut array)[1] = 1;
367        assert_eq!(<_ as Buffer<u16>>::as_bytes(&array)[..2], [1, 1]);
368        array.as_mut_slice()[0] = 1;
369        assert_eq!(array, [1, 2, 3, 4]);
370    }
371
372    #[test]
373    fn array_array() {
374        let mut array_array: <ArrayArrayBuffer<2, 4> as BufferType>::Buffer<u8> =
375            [[1, 2], [3, 4], [1, 2], [3, 4]];
376        assert_eq!(
377            <_ as Buffer<u8>>::as_bytes(&array_array),
378            &[1, 2, 3, 4, 1, 2, 3, 4]
379        );
380        <_ as BufferMut<u8>>::as_mut_bytes(&mut array_array)[1] = 1;
381        assert_eq!(
382            <_ as Buffer<u8>>::as_slice(&array_array),
383            [1, 1, 3, 4, 1, 2, 3, 4]
384        );
385    }
386
387    #[test]
388    fn slice() {
389        let slice: <SliceBuffer as BufferType>::Buffer<u16> = &[1234, 4321];
390        assert_eq!(slice.as_bytes(), &[210, 4, 225, 16]);
391        let mut slice_mut: <SliceMutBuffer as BufferType>::Buffer<u16> = &mut [4321, 1234];
392        BufferMut::as_mut_slice(&mut slice_mut)[0] = 1234;
393        BufferMut::as_mut_slice(&mut slice_mut)[1] = 4321;
394        assert_eq!(slice, slice_mut);
395    }
396
397    #[test]
398    fn slice_array() {
399        let slice_array: <SliceArrayBuffer<2> as BufferType>::Buffer<u32> = &[[1, 2], [3, 4]];
400        assert_eq!(<_ as Buffer<u32>>::as_slice(&slice_array), [1, 2, 3, 4]);
401        let mut slice_array_mut: <SliceArrayMutBuffer<3> as BufferType>::Buffer<u8> =
402            &mut [[1, 2, 3], [4, 5, 6]];
403        BufferMut::as_mut_slice(&mut slice_array_mut)[0] = 0;
404        assert_eq!(
405            <_ as Buffer<u8>>::as_bytes(&slice_array_mut),
406            &[0, 2, 3, 4, 5, 6]
407        );
408    }
409}
narrow/buffer.rs

narrow/
buffer.rs
