ps_alloc/
lib.rs

1#![allow(clippy::not_unsafe_ptr_arg_deref)]
2
3mod error;
4
5use std::alloc::Layout;
6
7pub use error::*;
8
9pub const HEADER_SIZE: usize = std::mem::size_of::<AllocationHeader>();
10pub const MARKER_FREE: [u8; 8] = *b"Fr33Mmry";
11pub const MARKER_USED: [u8; 8] = *b"U53dMmry";
12
13#[repr(align(16))]
14struct AllocationHeader {
15    marker: [u8; 8],
16    size: usize,
17}
18
19/// - A reasonably safe implementation of `alloc`.
20/// - Memory allocated by this function must be freed by this crate's `free`.
21/// - Caller guarantees `free` is called before the returned pointer goes out of scope.
22/// # Errors
23/// - `Err(ArithmeticError)` on integer overflow.
24/// - `Err(ImproperAlignment)` if the global allocator returns a misaligned pointer.
25/// - `Err(LayoutError)` if [`ALIGNMENT`] isn't a power of 2 or the computed size is not aligned.
26/// - `Err(OutOfMemory)` if `alloc()` returns a `nullptr`.
27#[allow(clippy::cast_ptr_alignment)]
28pub fn alloc(size: usize) -> Result<*mut u8, AllocationError> {
29    let size = size
30        .checked_add(HEADER_SIZE)
31        .ok_or(AllocationError::ArithmeticError)?
32        .checked_next_multiple_of(HEADER_SIZE)
33        .ok_or(AllocationError::ArithmeticError)?;
34
35    let layout = Layout::from_size_align(size, HEADER_SIZE)?;
36
37    let ptr = unsafe { std::alloc::alloc(layout) };
38
39    if ptr.is_null() {
40        return Err(AllocationError::OutOfMemory);
41    }
42
43    if 0 != (ptr as usize % HEADER_SIZE) {
44        unsafe { std::alloc::dealloc(ptr, layout) };
45
46        return Err(AllocationError::ImproperAlignment);
47    }
48
49    let header = unsafe { &mut *(ptr.cast::<AllocationHeader>()) };
50
51    header.marker = MARKER_USED;
52    header.size = size;
53
54    let ptr = unsafe { ptr.add(HEADER_SIZE) };
55
56    Ok(ptr)
57}
58
59/// - A reasonably safe implementation of `free`.
60/// - This function will free a pointer allocated by `alloc`.
61/// - Caller guarantees that the provided pointer was allocated by this crate's `alloc` function.
62/// - Providing `NULL` is safe and will return `Err(DeallocationError::NullPtr)`.
63/// - Providing any other pointer causes undefined behaviour.
64/// # Errors
65/// - Returns `Err(DeallocationError)` if a safety check fails.
66pub fn free<T>(ptr: *mut T) -> Result<(), DeallocationError> {
67    if ptr.is_null() {
68        return Err(DeallocationError::NullPtr);
69    }
70
71    if 0 != ptr as usize % HEADER_SIZE {
72        return Err(DeallocationError::ImproperAlignment);
73    }
74
75    #[allow(clippy::cast_ptr_alignment)]
76    let header_ptr = unsafe { ptr.cast::<u8>().sub(HEADER_SIZE).cast::<AllocationHeader>() };
77
78    if !header_ptr.is_aligned() {
79        return Err(DeallocationError::ImproperAlignment);
80    }
81
82    let header = unsafe { &mut *header_ptr };
83
84    if header.marker == MARKER_FREE {
85        return Err(DeallocationError::DoubleFree);
86    } else if header.marker != MARKER_USED {
87        return Err(DeallocationError::InvalidAllocation);
88    }
89
90    let layout = Layout::from_size_align(header.size, HEADER_SIZE)?;
91
92    header.marker = MARKER_FREE;
93
94    unsafe { std::alloc::dealloc(header_ptr.cast(), layout) };
95
96    Ok(())
97}
98
99/// Reallocates memory allocated by [`alloc`].
100/// # Errors
101/// - `AllocationError` if `alloc()` fails
102/// - `DeallocationError` if `free(ptr)` fails
103/// - `FreeFailedTwice` if `free(ptr)` fails and freeing the newly allocate pointer also fails
104/// - `ImproperAlignment` if `ptr` is not properly aligned
105/// - `InvalidPointer` if `ptr` was not allocated by [`alloc`] or is invalid
106/// - `UseAfterFree` if you try to `realloc` a freed pointer
107pub fn relloc(ptr: *mut u8, new_size: usize) -> Result<*mut u8, ReallocationError> {
108    if 0 == new_size {
109        if !ptr.is_null() {
110            free(ptr)?;
111        }
112
113        return Ok(std::ptr::null_mut());
114    }
115
116    if ptr.is_null() {
117        return Ok(alloc(new_size)?);
118    }
119
120    if 0 != ptr as usize % HEADER_SIZE {
121        return Err(ReallocationError::ImproperAlignment);
122    }
123
124    #[allow(clippy::cast_ptr_alignment)]
125    let header_ptr = unsafe { ptr.sub(HEADER_SIZE) }.cast::<AllocationHeader>();
126
127    if !header_ptr.is_aligned() {
128        return Err(ReallocationError::ImproperAlignment);
129    }
130
131    let header = unsafe { &*header_ptr };
132
133    if header.marker == MARKER_FREE {
134        return Err(ReallocationError::UseAfterFree);
135    } else if header.marker != MARKER_USED {
136        return Err(ReallocationError::InvalidPointer);
137    }
138
139    let new_ptr = alloc(new_size)?;
140
141    unsafe {
142        std::ptr::copy_nonoverlapping::<u8>(ptr, new_ptr, header.size.min(new_size));
143    }
144
145    match free(ptr) {
146        Ok(()) => Ok(new_ptr),
147        Err(err) => match free(new_ptr) {
148            Ok(()) => Err(err)?,
149            Err(err2) => Err(ReallocationError::FreeFailedTwice(err, err2)),
150        },
151    }
152}
ps_alloc/lib.rs

ps_alloc/
lib.rs
