polars_arrow/types/
aligned_bytes.rs

1use bytemuck::{Pod, Zeroable};
2
3use super::{days_ms, f16, i256, months_days_ns};
4use crate::array::View;
5
6/// Define that a type has the same byte alignment and size as `B`.
7///
8/// # Safety
9///
10/// This is safe to implement if both types have the same alignment and size.
11pub unsafe trait AlignedBytesCast<B: AlignedBytes>: Pod {}
12
13/// A representation of a type as raw bytes with the same alignment as the original type.
14pub trait AlignedBytes: Pod + Zeroable + Copy + Default + Eq {
15    const ALIGNMENT: usize;
16    const SIZE: usize;
17    const SIZE_ALIGNMENT_PAIR: PrimitiveSizeAlignmentPair;
18
19    type Unaligned: AsRef<[u8]>
20        + AsMut<[u8]>
21        + std::ops::Index<usize, Output = u8>
22        + std::ops::IndexMut<usize, Output = u8>
23        + for<'a> TryFrom<&'a [u8]>
24        + std::fmt::Debug
25        + Default
26        + IntoIterator<Item = u8>
27        + Pod;
28
29    fn to_unaligned(&self) -> Self::Unaligned;
30    fn from_unaligned(unaligned: Self::Unaligned) -> Self;
31
32    /// Safely cast a mutable reference to a [`Vec`] of `T` to a mutable reference of `Self`.
33    fn cast_vec_ref_mut<T: AlignedBytesCast<Self>>(vec: &mut Vec<T>) -> &mut Vec<Self> {
34        if cfg!(debug_assertions) {
35            assert_eq!(size_of::<T>(), size_of::<Self>());
36            assert_eq!(align_of::<T>(), align_of::<Self>());
37        }
38
39        // SAFETY: SameBytes guarantees that T:
40        // 1. has the same size
41        // 2. has the same alignment
42        // 3. is Pod (therefore has no life-time issues)
43        unsafe { std::mem::transmute(vec) }
44    }
45}
46
47macro_rules! impl_aligned_bytes {
48    (
49        $(($name:ident, $size:literal, $alignment:literal, $sap:ident, [$($eq_type:ty),*]),)+
50    ) => {
51        $(
52        /// Bytes with a size and alignment.
53        ///
54        /// This is used to reduce the monomorphizations for routines that solely rely on the size
55        /// and alignment of types.
56        #[derive(Debug, Copy, Clone, PartialEq, Eq, Default, Pod, Zeroable)]
57        #[repr(C, align($alignment))]
58        pub struct $name([u8; $size]);
59
60        impl AlignedBytes for $name {
61            const ALIGNMENT: usize = $alignment;
62            const SIZE: usize = $size;
63            const SIZE_ALIGNMENT_PAIR: PrimitiveSizeAlignmentPair = PrimitiveSizeAlignmentPair::$sap;
64
65            type Unaligned = [u8; $size];
66
67            #[inline(always)]
68            fn to_unaligned(&self) -> Self::Unaligned {
69                self.0
70            }
71            #[inline(always)]
72            fn from_unaligned(unaligned: Self::Unaligned) -> Self {
73                Self(unaligned)
74            }
75        }
76
77        impl AsRef<[u8; $size]> for $name {
78            #[inline(always)]
79            fn as_ref(&self) -> &[u8; $size] {
80                &self.0
81            }
82        }
83
84        $(
85        impl From<$eq_type> for $name {
86            #[inline(always)]
87            fn from(value: $eq_type) -> Self {
88                bytemuck::must_cast(value)
89            }
90        }
91        impl From<$name> for $eq_type {
92            #[inline(always)]
93            fn from(value: $name) -> Self {
94                bytemuck::must_cast(value)
95            }
96        }
97        unsafe impl AlignedBytesCast<$name> for $eq_type {}
98        )*
99        )+
100    }
101}
102
103#[derive(Clone, Copy)]
104pub enum PrimitiveSizeAlignmentPair {
105    S1A1,
106    S2A2,
107    S4A4,
108    S8A4,
109    S8A8,
110    S12A4,
111    S16A4,
112    S16A8,
113    S16A16,
114    S32A16,
115}
116
117impl PrimitiveSizeAlignmentPair {
118    pub const fn size(self) -> usize {
119        match self {
120            Self::S1A1 => 1,
121            Self::S2A2 => 2,
122            Self::S4A4 => 4,
123            Self::S8A4 | Self::S8A8 => 8,
124            Self::S12A4 => 12,
125            Self::S16A4 | Self::S16A8 | Self::S16A16 => 16,
126            Self::S32A16 => 32,
127        }
128    }
129
130    pub const fn alignment(self) -> usize {
131        match self {
132            Self::S1A1 => 1,
133            Self::S2A2 => 2,
134            Self::S4A4 | Self::S8A4 | Self::S12A4 | Self::S16A4 => 4,
135            Self::S8A8 | Self::S16A8 => 8,
136            Self::S16A16 | Self::S32A16 => 16,
137        }
138    }
139}
140
141impl_aligned_bytes! {
142    (Bytes1Alignment1, 1, 1, S1A1, [u8, i8]),
143    (Bytes2Alignment2, 2, 2, S2A2, [u16, i16, f16]),
144    (Bytes4Alignment4, 4, 4, S4A4, [u32, i32, f32]),
145    (Bytes8Alignment8, 8, 8, S8A8, [u64, i64, f64]),
146    (Bytes8Alignment4, 8, 4, S8A4, [days_ms]),
147    (Bytes12Alignment4, 12, 4, S12A4, [[u32; 3]]),
148    (Bytes16Alignment4, 16, 4, S16A4, [View]),
149    (Bytes16Alignment8, 16, 8, S16A8, [months_days_ns]),
150    (Bytes16Alignment16, 16, 16, S16A16, [u128, i128]),
151    (Bytes32Alignment16, 32, 16, S32A16, [i256]),
152}
polars_arrow/types/aligned_bytes.rs

polars_arrow/types/
aligned_bytes.rs
