ps_alloc/
lib.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
mod error;

use std::alloc::Layout;

pub use error::*;

pub const ALIGNMENT: usize = std::mem::size_of::<Allocation>();
pub const MARKER_FREE: [u8; 8] = *b"Fr33Mmry";
pub const MARKER_USED: [u8; 8] = *b"U53dMmry";

#[repr(align(16))]
struct Allocation {
    marker: [u8; 8],
    size: usize,
}

/// - A reasonably safe implementation of `alloc`.
/// - Memory allocated by this function must be freed by this crate's `free`.
/// - Caller guarantees `free` is called before the returned pointer goes out of scope.
pub fn alloc(size: usize) -> Result<*mut u8, AllocationError> {
    use AllocationError::*;

    let size = size
        .div_ceil(ALIGNMENT)
        .checked_add(1)
        .ok_or(ArithmeticError)?
        .checked_mul(ALIGNMENT)
        .ok_or(ArithmeticError)?;

    let layout = Layout::from_size_align(size, ALIGNMENT)?;

    let ptr = unsafe { std::alloc::alloc(layout) };

    if ptr.is_null() {
        Err(OutOfMemory)?
    }

    let allocation = unsafe { &mut *(ptr.cast::<Allocation>()) };

    allocation.marker = MARKER_USED;
    allocation.size = size;

    let ptr = unsafe { ptr.add(ALIGNMENT) };

    Ok(ptr)
}

/// - A reasonably safe implementation of `free`.
/// - This function will free a pointer allocated by `alloc`.
/// - Caller guarantees that the provided pointer was allocated by this crate's `alloc` function.
/// - Providing `NULL` is safe and will return `Err(DeallocationError::NullPtr)`.
/// - Providing any other pointer is undefined behaviour.
pub fn free<T>(ptr: *mut T) -> Result<(), DeallocationError> {
    use DeallocationError::*;

    if ptr.is_null() {
        Err(NullPtr)?
    }

    let ptr = ptr.cast::<Allocation>();

    if !ptr.is_aligned() {
        Err(ImproperAlignment)?
    }

    let ptr = unsafe { ptr.sub(1) };
    let allocation = unsafe { &mut *ptr };

    if allocation.marker == MARKER_FREE {
        Err(DoubleFree)?
    } else if allocation.marker != MARKER_USED {
        Err(InvalidAllocation)?
    }

    let layout = Layout::from_size_align(allocation.size, ALIGNMENT)?;

    allocation.marker = MARKER_FREE;

    unsafe { std::alloc::dealloc(ptr.cast(), layout) }

    Ok(())
}