alloc_align/
lib.rs

1#![feature(allocator_api)]
2#![cfg_attr(test, feature(pointer_is_aligned_to))]
3
4use std::{
5    alloc::{AllocError, Allocator, Global, Layout},
6    ptr::NonNull,
7};
8
9/// Allocator that forces all allocations to be aligned
10/// to at least `ALIGN` bytes.
11pub struct Align<const ALIGN: usize, A = Global> {
12    inner: A,
13}
14
15impl<const ALIGN: usize> Align<ALIGN, Global> {
16    pub const fn new() -> Self {
17        Self::wrap(Global)
18    }
19}
20
21impl<const ALIGN: usize, A> Align<ALIGN, A> {
22    pub const fn wrap(inner: A) -> Self {
23        Self { inner }
24    }
25}
26
27impl<const ALIGN: usize, A> Default for Align<ALIGN, A>
28where
29    A: Default,
30{
31    fn default() -> Self {
32        Self::wrap(A::default())
33    }
34}
35
36impl<const ALIGN: usize, A> Align<ALIGN, A> {
37    fn align_layout(layout: Layout) -> Layout {
38        layout.align_to(ALIGN).unwrap()
39    }
40}
41
42unsafe impl<const ALIGN: usize, A> Allocator for Align<ALIGN, A>
43where
44    A: Allocator,
45{
46    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
47        self.inner.allocate(Self::align_layout(layout))
48    }
49
50    fn allocate_zeroed(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
51        self.inner.allocate_zeroed(Self::align_layout(layout))
52    }
53
54    unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
55        unsafe { self.inner.deallocate(ptr, Self::align_layout(layout)) }
56    }
57
58    unsafe fn grow(
59        &self,
60        ptr: NonNull<u8>,
61        old_layout: Layout,
62        new_layout: Layout,
63    ) -> Result<NonNull<[u8]>, AllocError> {
64        unsafe {
65            self.inner.grow(
66                ptr,
67                Self::align_layout(old_layout),
68                Self::align_layout(new_layout),
69            )
70        }
71    }
72
73    unsafe fn grow_zeroed(
74        &self,
75        ptr: NonNull<u8>,
76        old_layout: Layout,
77        new_layout: Layout,
78    ) -> Result<NonNull<[u8]>, AllocError> {
79        unsafe {
80            self.inner.grow_zeroed(
81                ptr,
82                Self::align_layout(old_layout),
83                Self::align_layout(new_layout),
84            )
85        }
86    }
87
88    unsafe fn shrink(
89        &self,
90        ptr: NonNull<u8>,
91        old_layout: Layout,
92        new_layout: Layout,
93    ) -> Result<NonNull<[u8]>, AllocError> {
94        unsafe {
95            self.inner.shrink(
96                ptr,
97                Self::align_layout(old_layout),
98                Self::align_layout(new_layout),
99            )
100        }
101    }
102}
103
104#[cfg(test)]
105mod tests {
106    use super::*;
107
108    #[test]
109    fn aligned_vec() {
110        let mut vec = Vec::<u32, Align<4096, Global>>::new_in(Align::default());
111        vec.push(1);
112        assert!(vec.as_ptr().is_aligned_to(4096));
113    }
114}
alloc_align/lib.rs

alloc_align/
lib.rs
