polars_utils/parma/raw/
key.rs

1/// A low-level trait for keys.
2///
3/// This is used to allow unsized types such as [`str`] for keys.
4pub trait Key {
5    /// The alignment necessary for the key. Must return a power of two.
6    fn align() -> usize;
7
8    /// The size of the key in bytes.
9    fn size(&self) -> usize;
10
11    /// Initialize the key in the given memory location.
12    ///
13    /// # Safety
14    /// The memory location must satisfy the specified size and alignment.
15    unsafe fn init(&self, ptr: *mut u8);
16
17    /// Get a reference to the key from the given memory location.
18    ///
19    /// # Safety
20    /// The pointer must be valid and initialized with [`Key::init`].
21    unsafe fn get<'a>(ptr: *const u8) -> &'a Self;
22
23    /// Drop the key in place.
24    ///
25    /// # Safety
26    /// The pointer must be valid and initialized with [`Key::init`].
27    unsafe fn drop_in_place(ptr: *mut u8);
28}
29
30impl<T: Copy> Key for [T] {
31    #[inline(always)]
32    fn align() -> usize {
33        align_of::<usize>().max(align_of::<T>())
34    }
35
36    #[inline(always)]
37    fn size(&self) -> usize {
38        size_of::<usize>().next_multiple_of(align_of::<T>()) + self.len()
39    }
40
41    #[inline(always)]
42    unsafe fn init(&self, ptr: *mut u8) {
43        unsafe {
44            let p_len = ptr.cast::<usize>();
45            let p_data = ptr.add(size_of::<usize>().next_multiple_of(align_of::<T>()));
46            let len = self.len();
47            p_len.write(len);
48            std::ptr::copy_nonoverlapping(self.as_ptr(), p_data.cast(), len)
49        }
50    }
51
52    #[inline(always)]
53    unsafe fn get<'a>(ptr: *const u8) -> &'a Self {
54        unsafe {
55            let p_len = ptr.cast::<usize>();
56            let p_data = ptr.add(size_of::<usize>().next_multiple_of(align_of::<T>()));
57            let len = p_len.read();
58            core::slice::from_raw_parts(p_data.cast(), len)
59        }
60    }
61
62    #[inline(always)]
63    unsafe fn drop_in_place(_ptr: *mut u8) {}
64}
65
66impl Key for str {
67    #[inline(always)]
68    fn align() -> usize {
69        <[u8] as Key>::align()
70    }
71
72    #[inline(always)]
73    fn size(&self) -> usize {
74        <[u8] as Key>::size(self.as_bytes())
75    }
76
77    #[inline(always)]
78    unsafe fn init(&self, ptr: *mut u8) {
79        unsafe { <[u8] as Key>::init(self.as_bytes(), ptr) }
80    }
81
82    #[inline(always)]
83    unsafe fn get<'a>(ptr: *const u8) -> &'a Self {
84        unsafe { core::str::from_utf8_unchecked(<[u8] as Key>::get(ptr)) }
85    }
86
87    #[inline(always)]
88    unsafe fn drop_in_place(_ptr: *mut u8) {}
89}
90
91impl<T: Clone + Sized> Key for T {
92    fn align() -> usize {
93        align_of::<T>()
94    }
95
96    fn size(&self) -> usize {
97        size_of::<T>()
98    }
99
100    #[inline(always)]
101    unsafe fn init(&self, ptr: *mut u8) {
102        unsafe {
103            ptr.cast::<T>().write(self.clone());
104        }
105    }
106
107    #[inline(always)]
108    unsafe fn get<'a>(ptr: *const u8) -> &'a Self {
109        unsafe { &*ptr.cast::<T>() }
110    }
111
112    #[inline(always)]
113    unsafe fn drop_in_place(ptr: *mut u8) {
114        unsafe { ptr.cast::<T>().drop_in_place() }
115    }
116}
polars_utils/parma/raw/key.rs

polars_utils/parma/raw/
key.rs
