Skip to main content

mnemosyne_arena/scratch/
element.rs

1//! Element types admitted into scratch buffers: sealed plain-old-data
2//! scalars for which an all-zero bit pattern is a valid value.
3
4/// Default alignment for scratch buffers (64 bytes = one AVX-512 cache line).
5pub const DEFAULT_SCRATCH_ALIGN: usize = 64;
6
7mod sealed {
8    pub trait ScratchElementSealed {}
9
10    // Blanket impl: any Zeroable type has a valid all-zero bit pattern by
11    // definition. The sealed trait prevents non-Zeroable downstream impls.
12    #[cfg(feature = "bytemuck")]
13    impl<T: bytemuck::Zeroable + Copy + Send + Sync + 'static> ScratchElementSealed for T {}
14}
15
16/// Element types that the scratch pool can manage.
17///
18/// Implemented for `f32`, `f64`, `u8`, common integer types, and (with the
19/// `eunomia` feature) `eunomia::Complex<f32>` and `eunomia::Complex<f64>`.
20/// With the `bytemuck` feature, a blanket impl covers any
21/// `bytemuck::Zeroable + Copy + Send + Sync + 'static` type — including
22/// arbitrary repr(C) POD structs that derive `Zeroable`.
23///
24/// # Safety invariant
25///
26/// Every implementor must tolerate an all-zero bit pattern as a valid,
27/// non-trapping value. This is true for all floating-point scalars (±0.0),
28/// all integers (0), eunomia complex numbers, and any type that implements
29/// `bytemuck::Zeroable`. Types with validity invariants (e.g. non-null
30/// pointers, Rust references, enums with niche optimizations) must not
31/// implement this trait.
32pub trait ScratchElement: sealed::ScratchElementSealed + Copy + Send + Sync + 'static {
33    /// Alignment in bytes required for SIMD operations on this element type.
34    const ALIGN_BYTES: usize;
35}
36
37macro_rules! impl_scratch_element {
38    ($($t:ty),* $(,)?) => {
39        $(
40            impl sealed::ScratchElementSealed for $t {}
41            impl ScratchElement for $t {
42                const ALIGN_BYTES: usize = DEFAULT_SCRATCH_ALIGN;
43            }
44        )*
45    };
46}
47
48// Manual impls when the bytemuck blanket is not active.
49// With the `bytemuck` feature, the blanket `impl<T: Zeroable + ...>` below
50// covers all of these (plus every other Zeroable POD type).
51#[cfg(not(feature = "bytemuck"))]
52impl_scratch_element!(
53    f32, f64, bool, u8, u16, u32, u64, usize, i8, i16, i32, i64, isize
54);
55
56#[cfg(all(feature = "eunomia", not(feature = "bytemuck")))]
57impl sealed::ScratchElementSealed for eunomia::Complex<f32> {}
58#[cfg(all(feature = "eunomia", not(feature = "bytemuck")))]
59impl ScratchElement for eunomia::Complex<f32> {
60    const ALIGN_BYTES: usize = DEFAULT_SCRATCH_ALIGN;
61}
62
63#[cfg(all(feature = "eunomia", not(feature = "bytemuck")))]
64impl sealed::ScratchElementSealed for eunomia::Complex<f64> {}
65#[cfg(all(feature = "eunomia", not(feature = "bytemuck")))]
66impl ScratchElement for eunomia::Complex<f64> {
67    const ALIGN_BYTES: usize = DEFAULT_SCRATCH_ALIGN;
68}
69
70/// Blanket `ScratchElement` impl for arbitrary POD structs.
71///
72/// Enabled with `features = ["bytemuck"]`. Any type that implements
73/// `bytemuck::Zeroable` (i.e. has a valid all-zero representation) and
74/// satisfies `Copy + Send + Sync + 'static` can be used with
75/// [`AlignedVec`][super::aligned_vec::AlignedVec] and
76/// [`ScratchPool`][super::pool::ScratchPool].
77///
78/// The alignment for structs covered by this impl is the crate-wide
79/// `DEFAULT_SCRATCH_ALIGN` (64 bytes). If a type requires a different
80/// alignment, provide an explicit `impl ScratchElement` instead.
81#[cfg(feature = "bytemuck")]
82impl<T: bytemuck::Zeroable + Copy + Send + Sync + 'static> ScratchElement for T {
83    const ALIGN_BYTES: usize = DEFAULT_SCRATCH_ALIGN;
84}
85
86/// Default alignment constant for external consumers.
87#[inline]
88pub const fn default_align() -> usize {
89    DEFAULT_SCRATCH_ALIGN
90}