yume_pdq/
alignment.rs

1/*
2 * Copyright (c) 2025 Yumechi <yume@yumechi.jp>
3 *
4 * Created on Wednesday, March 26, 2025
5 * Author: Yumechi <yume@yumechi.jp>
6 *
7 * SPDX-License-Identifier: Apache-2.0
8 *
9 * Licensed under the Apache License, Version 2.0 (the "License");
10 * you may not use this file except in compliance with the License.
11 * You may obtain a copy of the License at
12 *
13 * http://www.apache.org/licenses/LICENSE-2.0
14 *
15 * Unless required by applicable law or agreed to in writing, software
16 * distributed under the License is distributed on an "AS IS" BASIS,
17 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
18 * See the License for the specific language governing permissions and
19 * limitations under the License.
20 */
21
22use core::ops::{Deref, DerefMut};
23
24use const_default::ConstDefault;
25use generic_array::{
26    ArrayLength, GenericArray,
27    typenum::{B0, B1, IsGreater, PowerOfTwo, U0, U1, U8, U32, U64, UInt, Unsigned},
28};
29use zeroize::DefaultIsZeroes;
30
31/// A trait for a trivial aligner that has no function except setting the alignment and transparently holding a value of type `T`.
32///
33/// # Safety
34///
35/// Implementor must ensure their memory layout is valid.
36pub unsafe trait AlignerTo<T>: Deref<Target = T> + DerefMut<Target = T> + From<T> {
37    /// The alignment of the aligner.
38    type Alignment: Unsigned + PowerOfTwo + IsGreater<U0, Output = B1>;
39    /// The type of the aligner.
40    type Output;
41
42    /// Create a `core::alloc::Layout` for the aligner.
43    ///
44    /// # Panics
45    ///
46    /// Panics if the request memory size is rejected by the allocator API.
47    fn create_layout() -> core::alloc::Layout;
48}
49
50unsafe impl<T> AlignerTo<T> for Align1<T> {
51    type Alignment = U1;
52    type Output = Align1<T>;
53
54    fn create_layout() -> core::alloc::Layout {
55        core::alloc::Layout::from_size_align(core::mem::size_of::<T>(), 1)
56            .expect("request memory size is too large for allocation")
57    }
58}
59
60unsafe impl<T> AlignerTo<T> for Align8<T> {
61    type Alignment = U8;
62    type Output = Align8<T>;
63
64    fn create_layout() -> core::alloc::Layout {
65        core::alloc::Layout::from_size_align(core::mem::size_of::<T>(), 8)
66            .expect("request memory size is too large for allocation")
67    }
68}
69
70unsafe impl<T> AlignerTo<T> for Align32<T> {
71    type Alignment = U32;
72    type Output = Align32<T>;
73    fn create_layout() -> core::alloc::Layout {
74        core::alloc::Layout::from_size_align(core::mem::size_of::<T>(), 32)
75            .expect("request memory size is too large for allocation")
76    }
77}
78
79unsafe impl<T> AlignerTo<T> for Align64<T> {
80    type Alignment = U64;
81    type Output = Align64<T>;
82
83    fn create_layout() -> core::alloc::Layout {
84        core::alloc::Layout::from_size_align(core::mem::size_of::<T>(), 64)
85            .expect("request memory size is too large for allocation")
86    }
87}
88
89/// Align the item to 1 byte (a dummy aligner).
90#[derive(Debug, Clone, Copy)]
91#[repr(transparent)]
92pub struct Align1<T>(pub T);
93
94impl<T: DefaultIsZeroes + Copy> DefaultIsZeroes for Align1<T> {}
95
96impl<T: Default> Default for Align1<T> {
97    fn default() -> Self {
98        Self(Default::default())
99    }
100}
101
102impl<T> Deref for Align1<T> {
103    type Target = T;
104
105    fn deref(&self) -> &Self::Target {
106        &self.0
107    }
108}
109
110impl<T> DerefMut for Align1<T> {
111    fn deref_mut(&mut self) -> &mut Self::Target {
112        &mut self.0
113    }
114}
115
116impl<T> From<T> for Align1<T> {
117    fn from(value: T) -> Self {
118        Self(value)
119    }
120}
121
122impl<T> Align1<T> {
123    /// Create a new `Align1<T>` from a `T`.
124    pub const fn new(value: T) -> Self {
125        Self(value)
126    }
127
128    /// Convert a pointer to a `Align1<T>` to a reference.
129    ///
130    /// # Safety
131    ///
132    /// The pointer was not checked to be aligned to 1 byte, so it is the caller's responsibility to ensure it is.
133    pub const unsafe fn from_raw_unchecked(ptr: *const T) -> *const Self {
134        ptr.cast::<Self>()
135    }
136
137    /// Convert a reference to a `T` to a reference to a `Align1<T>`.
138    ///
139    /// # Panics
140    ///
141    pub fn from_raw(input: &T) -> &Self {
142        unsafe { &*(input as *const T).cast::<Self>() }
143    }
144
145    /// Convert a pointer to a `Align1<T>` to a mutable reference.
146    ///
147    /// # Safety
148    ///
149    /// The pointer was not checked to be aligned to 1 byte, so it is the caller's responsibility to ensure it is.
150    pub const unsafe fn from_raw_mut_unchecked(ptr: *mut T) -> *mut Self {
151        ptr.cast::<Self>()
152    }
153
154    /// Convert a reference to a `T` to a mutable reference to a `Align1<T>`.
155    ///
156    /// # Panics
157    ///
158    /// Panics if the reference is not aligned to 1 byte.
159    pub fn from_raw_mut(input: &mut T) -> &mut Self {
160        unsafe { &mut *(input as *mut T).cast::<Self>() }
161    }
162}
163
164#[repr(align(8))]
165#[derive(Debug, Clone, Copy)]
166/// Align the item to 8 bytes.
167pub struct Align8<T>(pub T);
168
169impl<T: DefaultIsZeroes + Copy> DefaultIsZeroes for Align8<T> {}
170
171impl<T: Default> Default for Align8<T> {
172    fn default() -> Self {
173        Self(Default::default())
174    }
175}
176
177impl<T> Deref for Align8<T> {
178    type Target = T;
179
180    fn deref(&self) -> &Self::Target {
181        &self.0
182    }
183}
184
185impl<T> DerefMut for Align8<T> {
186    fn deref_mut(&mut self) -> &mut Self::Target {
187        &mut self.0
188    }
189}
190
191impl<T> From<T> for Align8<T> {
192    fn from(value: T) -> Self {
193        Self(value)
194    }
195}
196
197type Times8<T> = UInt<UInt<UInt<T, B0>, B0>, B0>;
198
199type Times32<T> = UInt<UInt<Times8<T>, B0>, B0>;
200
201type Times64<T> = UInt<Times32<T>, B0>;
202
203impl<T, L: ArrayLength> Align8<GenericArray<T, Times8<L>>> {
204    /// Convert a `GenericArray<Align8<GenericArray<T, L>>, U>` to a `Align8<GenericArray<GenericArray<T, L>, U>>`, if L is a multiple of 8.
205    #[must_use]
206    pub const fn lift<U: ArrayLength>(
207        input: GenericArray<Self, U>,
208    ) -> Align8<GenericArray<GenericArray<T, Times8<L>>, U>> {
209        // SAFETY: if L is a multiple of 8, then GenericArray<T, L> is either a ZST or at least a multiple of 8 bytes, guaranteeing a zero-padding layout
210        unsafe { generic_array::const_transmute(input) }
211    }
212
213    /// Convert a `Box<GenericArray<Align8<GenericArray<T, L>>, U>>` to a `Box<Align8<GenericArray<GenericArray<T, L>, U>>>`, if L is a multiple of 8.
214    #[cfg(feature = "alloc")]
215    #[must_use]
216    pub const fn lift_boxed<U: ArrayLength>(
217        input: alloc::boxed::Box<GenericArray<Self, U>>,
218    ) -> alloc::boxed::Box<Align8<GenericArray<GenericArray<T, Times8<L>>, U>>> {
219        // SAFETY: if L is a multiple of 8, then GenericArray<T, L> is either a ZST or at least a multiple of 8 bytes, guaranteeing a zero-padding layout
220        unsafe { generic_array::const_transmute(input) }
221    }
222}
223
224impl<T> Align8<T> {
225    /// Create a new `Align8<T>` from a `T`.
226    pub const fn new(value: T) -> Self {
227        Self(value)
228    }
229
230    /// Convert a pointer to a `Align8<T>` to a reference.
231    ///
232    /// # Safety
233    ///
234    /// The pointer was not checked to be aligned to 8 bytes, so it is the caller's responsibility to ensure it is.
235    pub const unsafe fn from_raw_unchecked(ptr: *const T) -> *const Self {
236        ptr.cast::<Self>()
237    }
238
239    /// Convert a reference to a `T` to a reference to a `Align8<T>`.
240    ///
241    /// # Panics
242    ///
243    /// Panics if the reference is not aligned to 8 bytes.  
244    pub fn from_raw(input: &T) -> &Self {
245        let ptr = input as *const T;
246        assert_eq!(ptr.align_offset(8), 0, "pointer is not aligned to 8 bytes");
247        unsafe { &*(ptr.cast::<Self>()) }
248    }
249
250    /// Convert a pointer to a `Align8<T>` to a mutable reference.
251    ///
252    /// # Safety
253    ///
254    /// The pointer was not checked to be aligned to 8 bytes, so it is the caller's responsibility to ensure it is.
255    pub const unsafe fn from_raw_mut_unchecked(ptr: *mut T) -> *mut Self {
256        ptr.cast::<Self>()
257    }
258
259    /// Convert a reference to a `T` to a mutable reference to a `Align8<T>`.
260    ///
261    /// # Panics
262    ///
263    /// Panics if the reference is not aligned to 8 bytes.
264    pub fn from_raw_mut(input: &mut T) -> &mut Self {
265        let ptr = input as *mut T;
266        assert_eq!(ptr.align_offset(8), 0, "pointer is not aligned to 8 bytes");
267        unsafe { &mut *(ptr.cast::<Self>()) }
268    }
269}
270
271#[repr(align(32))]
272#[derive(Debug, Clone, Copy)]
273/// Align the item to 32 bytes.
274pub struct Align32<T>(pub T);
275
276impl<T: DefaultIsZeroes + Copy> DefaultIsZeroes for Align32<T> {}
277
278impl<T: Default> Default for Align32<T> {
279    fn default() -> Self {
280        Self(Default::default())
281    }
282}
283
284impl<T> Deref for Align32<T> {
285    type Target = T;
286
287    fn deref(&self) -> &Self::Target {
288        &self.0
289    }
290}
291
292impl<T> DerefMut for Align32<T> {
293    fn deref_mut(&mut self) -> &mut Self::Target {
294        &mut self.0
295    }
296}
297
298impl<T> From<T> for Align32<T> {
299    fn from(value: T) -> Self {
300        Self(value)
301    }
302}
303
304impl<T, L: ArrayLength> Align32<GenericArray<T, Times32<L>>> {
305    /// Convert a `GenericArray<Align32<GenericArray<T, L>>, U>` to a `Align32<GenericArray<GenericArray<T, L>, U>>`, if L is a multiple of 32.
306    #[must_use]
307    pub const fn lift<U: ArrayLength>(
308        input: GenericArray<Self, U>,
309    ) -> Align32<GenericArray<GenericArray<T, Times32<L>>, U>> {
310        // SAFETY: if L is a multiple of 32, then GenericArray<T, L> is either a ZST or at least a multiple of 32 bytes, guaranteeing a zero-padding layout
311        unsafe { generic_array::const_transmute(input) }
312    }
313
314    /// Convert a `Box<GenericArray<Align32<GenericArray<T, L>>, U>>` to a `Box<Align32<GenericArray<GenericArray<T, L>, U>>>`, if L is a multiple of 32.
315    #[cfg(feature = "alloc")]
316    #[must_use]
317    pub const fn lift_boxed<U: ArrayLength>(
318        input: alloc::boxed::Box<GenericArray<Self, U>>,
319    ) -> alloc::boxed::Box<Align32<GenericArray<GenericArray<T, Times32<L>>, U>>> {
320        // SAFETY: if L is a multiple of 32, then GenericArray<T, L> is either a ZST or at least a multiple of 32 bytes, guaranteeing a zero-padding layout
321        unsafe { generic_array::const_transmute(input) }
322    }
323}
324
325impl<T> Align32<T> {
326    /// Create a new `Align32<T>` from a `T`.
327    pub const fn new(value: T) -> Self {
328        Self(value)
329    }
330
331    /// Convert a pointer to a `Align32<T>` to a reference.
332    ///
333    /// # Safety
334    ///
335    /// The pointer was not checked to be aligned to 32 bytes, so it is the caller's responsibility to ensure it is.
336    pub const unsafe fn from_raw_unchecked(ptr: *const T) -> *const Self {
337        ptr.cast::<Self>()
338    }
339
340    /// Convert a reference to a `T` to a reference to a `Align32<T>`.
341    ///
342    /// # Panics
343    ///
344    /// Panics if the reference is not aligned to 32 bytes.
345    pub fn from_raw(input: &T) -> &Self {
346        let ptr = input as *const T;
347        assert_eq!(
348            ptr.align_offset(32),
349            0,
350            "pointer is not aligned to 32 bytes"
351        );
352        unsafe { &*(ptr.cast::<Self>()) }
353    }
354
355    /// Convert a pointer to a `Align32<T>` to a mutable reference.
356    ///
357    /// # Safety
358    ///
359    /// The pointer was not checked to be aligned to 32 bytes, so it is the caller's responsibility to ensure it is.
360    pub const unsafe fn from_raw_mut_unchecked(ptr: *mut T) -> *mut Self {
361        ptr.cast::<Self>()
362    }
363
364    /// Convert a reference to a `T` to a mutable reference to a `Align32<T>`.
365    ///
366    /// # Panics
367    ///
368    /// Panics if the reference is not aligned to 32 bytes.
369    pub fn from_raw_mut(input: &mut T) -> &mut Self {
370        let ptr = input as *mut T;
371        assert_eq!(
372            ptr.align_offset(32),
373            0,
374            "pointer is not aligned to 32 bytes"
375        );
376        unsafe { &mut *(ptr.cast::<Self>()) }
377    }
378}
379
380#[repr(align(64))]
381/// Align the item to 64 bytes.
382#[derive(Debug, Clone, Copy)]
383pub struct Align64<T>(pub T);
384
385impl<T: DefaultIsZeroes + Copy> DefaultIsZeroes for Align64<T> {}
386
387impl<T: Default> Default for Align64<T> {
388    fn default() -> Self {
389        Self(Default::default())
390    }
391}
392
393impl<T> Deref for Align64<T> {
394    type Target = T;
395
396    fn deref(&self) -> &Self::Target {
397        &self.0
398    }
399}
400
401impl<T> DerefMut for Align64<T> {
402    fn deref_mut(&mut self) -> &mut Self::Target {
403        &mut self.0
404    }
405}
406
407impl<T> From<T> for Align64<T> {
408    fn from(value: T) -> Self {
409        Self(value)
410    }
411}
412
413impl<T> Align64<T> {
414    /// Create a new `Align64<T>` from a `T`.
415    pub const fn new(value: T) -> Self {
416        Self(value)
417    }
418
419    /// Convert a pointer to a `Align64<T>` to a reference.
420    ///
421    /// # Safety
422    ///
423    /// The pointer was not checked to be aligned to 64 bytes, so it is the caller's responsibility to ensure it is.
424    pub const unsafe fn from_raw_unchecked(ptr: *const T) -> *const Self {
425        ptr.cast::<Self>()
426    }
427
428    /// Convert a reference to a `T` to a reference to a `Align64<T>`.
429    ///
430    /// # Panics
431    ///
432    /// Panics if the reference is not aligned to 64 bytes.
433    pub fn from_raw(input: &T) -> &Self {
434        let ptr = input as *const T;
435        assert_eq!(
436            ptr.align_offset(64),
437            0,
438            "pointer is not aligned to 64 bytes"
439        );
440        unsafe { &*(ptr.cast::<Self>()) }
441    }
442
443    /// Convert a pointer to a `Align64<T>` to a mutable reference.
444    ///
445    /// # Safety
446    ///
447    /// The pointer was not checked to be aligned to 64 bytes, so it is the caller's responsibility to ensure it is.
448    pub const unsafe fn from_raw_mut_unchecked(ptr: *mut T) -> *mut Self {
449        ptr.cast::<Self>()
450    }
451
452    /// Convert a reference to a `T` to a mutable reference to a `Align64<T>`.
453    ///
454    /// # Panics
455    ///
456    /// Panics if the reference is not aligned to 64 bytes.
457    pub fn from_raw_mut(input: &mut T) -> &mut Self {
458        let ptr = input as *mut T;
459        assert_eq!(
460            ptr.align_offset(64),
461            0,
462            "pointer is not aligned to 64 bytes"
463        );
464        unsafe { &mut *(ptr.cast::<Self>()) }
465    }
466}
467
468impl<T, L: ArrayLength> Align64<GenericArray<T, Times64<L>>> {
469    /// Convert a `GenericArray<Align64<GenericArray<T, L>>, U>` to a `Align64<GenericArray<GenericArray<T, L>, U>>`, if L is a multiple of 64.
470    #[must_use]
471    pub const fn lift<U: ArrayLength>(
472        input: GenericArray<Self, U>,
473    ) -> Align64<GenericArray<GenericArray<T, Times64<L>>, U>> {
474        // SAFETY: if L is a multiple of 64, then GenericArray<T, L> is either a ZST or at least a multiple of 64 bytes, guaranteeing a zero-padding layout
475        unsafe { generic_array::const_transmute(input) }
476    }
477
478    /// Convert a `Box<GenericArray<Align64<GenericArray<T, L>>, U>>` to a `Box<Align64<GenericArray<GenericArray<T, L>, U>>>`, if L is a multiple of 64.
479    #[cfg(feature = "alloc")]
480    #[must_use]
481    pub const fn lift_boxed<U: ArrayLength>(
482        input: alloc::boxed::Box<GenericArray<Self, U>>,
483    ) -> alloc::boxed::Box<Align64<GenericArray<GenericArray<T, Times64<L>>, U>>> {
484        // SAFETY: if L is a multiple of 64, then GenericArray<T, L> is either a ZST or at least a multiple of 64 bytes, guaranteeing a zero-padding layout
485        unsafe { generic_array::const_transmute(input) }
486    }
487}
488
489#[repr(C)]
490/// A double-ended padded array.
491pub struct DefaultPaddedArray<E, L: ArrayLength, P: ArrayLength> {
492    _pad0: GenericArray<E, P>,
493    inner: GenericArray<E, L>,
494    _pad1: GenericArray<E, P>,
495}
496
497impl<E, L: ArrayLength, P: ArrayLength> DefaultPaddedArray<E, L, P> {
498    /// Get a pointer to the inner array.
499    pub fn as_ptr(&self) -> *const E {
500        self.inner.as_ptr()
501    }
502
503    /// Get a mutable pointer to the inner array.
504    pub fn as_mut_ptr(&mut self) -> *mut E {
505        self.inner.as_mut_ptr()
506    }
507}
508
509impl<E: ConstDefault, L: ArrayLength, P: ArrayLength> AsRef<GenericArray<E, L>>
510    for DefaultPaddedArray<E, L, P>
511{
512    fn as_ref(&self) -> &GenericArray<E, L> {
513        &self.inner
514    }
515}
516
517impl<E: ConstDefault, L: ArrayLength, P: ArrayLength> AsMut<GenericArray<E, L>>
518    for DefaultPaddedArray<E, L, P>
519{
520    fn as_mut(&mut self) -> &mut GenericArray<E, L> {
521        &mut self.inner
522    }
523}
524
525impl<E: ConstDefault, L: ArrayLength, P: ArrayLength> DefaultPaddedArray<E, L, P>
526where
527    <P as ArrayLength>::ArrayType<E>: ConstDefault,
528{
529    /// Create a new `DefaultPaddedArray` with the given inner array.
530    pub const fn new(inner: GenericArray<E, L>) -> Self {
531        Self {
532            _pad0: GenericArray::const_default(),
533            inner,
534            _pad1: GenericArray::const_default(),
535        }
536    }
537}
538
539impl<E: ConstDefault, L: ArrayLength, P: ArrayLength> Deref for DefaultPaddedArray<E, L, P> {
540    type Target = GenericArray<E, L>;
541
542    fn deref(&self) -> &Self::Target {
543        &self.inner
544    }
545}
546
547impl<E: ConstDefault, L: ArrayLength, P: ArrayLength> DerefMut for DefaultPaddedArray<E, L, P> {
548    fn deref_mut(&mut self) -> &mut Self::Target {
549        &mut self.inner
550    }
551}
552
553impl<E: Default, L: ArrayLength, P: ArrayLength> Default for DefaultPaddedArray<E, L, P> {
554    fn default() -> Self {
555        Self {
556            _pad0: GenericArray::default(),
557            inner: GenericArray::default(),
558            _pad1: GenericArray::default(),
559        }
560    }
561}
562
563#[cfg(feature = "alloc")]
564/// Allocate a 1D generic array on the heap and zeroize it, without creating intermediates on the stack.
565pub fn calloc_generic_array_1d<
566    T: DefaultIsZeroes + Copy,
567    A: AlignerTo<GenericArray<T, L>>,
568    L: ArrayLength,
569>() -> alloc::boxed::Box<<A as AlignerTo<GenericArray<T, L>>>::Output> {
570    // SAFETY:
571    // - T must be non-drop (constrained by Copy)
572    // - Default must be the natural zero value for T (constrained by DefaultIsZeroes)
573
574    use core::mem::MaybeUninit;
575    unsafe {
576        let template = T::default();
577        let layout = A::create_layout();
578        let memory = alloc::alloc::alloc(layout);
579        let mut_slice = core::slice::from_raw_parts_mut(memory.cast::<MaybeUninit<T>>(), L::USIZE);
580        mut_slice.iter_mut().for_each(|item| {
581            item.as_mut_ptr().write(template);
582        });
583        alloc::boxed::Box::from_raw(memory.cast())
584    }
585}
586
587#[cfg(feature = "alloc")]
588/// Allocate a 2D generic array on the heap and zeroize it, without creating intermediates on the stack.
589pub fn calloc_generic_array_2d<
590    T: DefaultIsZeroes + Copy,
591    A: AlignerTo<GenericArray<GenericArray<T, M>, L>>,
592    L: ArrayLength,
593    M: ArrayLength,
594>() -> alloc::boxed::Box<<A as AlignerTo<GenericArray<GenericArray<T, M>, L>>>::Output> {
595    // SAFETY:
596    // - T must be non-drop (constrained by Copy)
597    // - Default must be the natural zero value for T (constrained by DefaultIsZeroes)
598
599    use core::mem::MaybeUninit;
600    unsafe {
601        let template = T::default();
602        let layout = A::create_layout();
603        let memory = alloc::alloc::alloc(layout);
604        let mut_slice =
605            core::slice::from_raw_parts_mut(memory.cast::<MaybeUninit<T>>(), L::USIZE * M::USIZE);
606
607        mut_slice.iter_mut().for_each(|item| {
608            item.as_mut_ptr().write(template);
609        });
610        alloc::boxed::Box::from_raw(memory.cast())
611    }
612}
613
614#[cfg(feature = "alloc")]
615/// Allocate a 3D generic array on the heap and zeroize it, without creating intermediates on the stack.
616pub fn calloc_generic_array_3d<
617    T: DefaultIsZeroes + Copy,
618    A: AlignerTo<GenericArray<GenericArray<GenericArray<T, M>, N>, L>>,
619    L: ArrayLength,
620    M: ArrayLength,
621    N: ArrayLength,
622>()
623-> alloc::boxed::Box<<A as AlignerTo<GenericArray<GenericArray<GenericArray<T, M>, N>, L>>>::Output>
624{
625    // SAFETY:
626    // - T must be non-drop (constrained by Copy)
627    // - Default must be the natural zero value for T (constrained by DefaultIsZeroes)
628
629    use core::mem::MaybeUninit;
630
631    unsafe {
632        let template = T::default();
633        let layout = A::create_layout();
634        let memory = alloc::alloc::alloc(layout);
635        let mut_slice = core::slice::from_raw_parts_mut(
636            memory.cast::<MaybeUninit<T>>(),
637            L::USIZE * M::USIZE * N::USIZE,
638        );
639        mut_slice.iter_mut().for_each(|item| {
640            item.as_mut_ptr().write(template);
641        });
642        alloc::boxed::Box::from_raw(memory.cast())
643    }
644}
yume_pdq/alignment.rs

yume_pdq/
alignment.rs
