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
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125

//! A fast lock-free Allocator
//!
//! # Internal design
//! ## Thread-Local Caches
//! Each thread has a small Cache of ready to use allocations, which help with performance
//! in most cases as they dont need any extra synchronization between threads.
//!
//! ## Heap
//! The Heap is the central shared entity, which actually manages the underlying allocations
//! as well as the needed synchronization between different threads.
//!
//! # References
//! * [Paper - 'LRMalloc: a Modern and Competitive Lock-Free Dynamic Memory Allocator'](https://vecpar2018.ncc.unesp.br/wp-content/uploads/2018/09/VECPAR_2018_paper_27.pdf)

use std::{
    alloc::{handle_alloc_error, GlobalAlloc},
    cell::RefCell,
};

mod util;

mod cache;
mod size_classes;
use cache::Cache;
mod heap;
use heap::Heap;
mod pagemap;
use pagemap::PageMap;

mod descriptor;

static PAGEMAP: PageMap = PageMap::new();

thread_local! {
    static CACHE: RefCell<Cache> = RefCell::new(Cache::new());
}

/// The actual Allocator Struct, which can be used for allocating and freeing memory
#[derive(Debug)]
pub struct Allocator {
    heap: Heap,
}

impl Allocator {
    /// Creates a new Instance of the Allocator
    ///
    /// # Note
    /// All Instances of the Allocator share some amount of Data, so they are currently not
    /// independant of each other.
    /// You should only create a single Instance for use as the Global-Allocator of your program
    pub const fn new() -> Self {
        Self { heap: Heap::new() }
    }

    /// Allocates Memory for the given Layout using this allocator
    ///
    /// # Safety
    /// The caller needs to ensure that the given Memory Layout is valid
    pub unsafe fn allocate(&self, layout: std::alloc::Layout) -> *mut u8 {
        let size_class = match size_classes::get_size_class_index(layout.size()) {
            Some(s) => s,
            None => {
                return self.heap.allocate_large(layout, &PAGEMAP);
            }
        };

        CACHE.with(|raw| {
            let mut cache = match raw.try_borrow_mut() {
                Ok(r) => r,
                Err(_) => {
                    handle_alloc_error(layout);
                }
            };

            if let Some(ptr) = cache.try_alloc(size_class) {
                return ptr;
            }

            self.heap.fill_cache(&mut cache, size_class, &PAGEMAP);
            cache.try_alloc(size_class).expect("We just filled the Cache with new Blocks, so there should at least be one available block to use for the Allocation")
        })
    }

    /// Deallocates the Memory for the given Ptr with the given Layout
    ///
    /// # Safety
    /// The Caller needs to ensure that the given Ptr was given out by this allocator and the given
    /// Layout matches the layout that was used when obtaining the Ptr
    pub unsafe fn deallocate(&self, ptr: *mut u8, layout: std::alloc::Layout) {
        let desc_ptr = match PAGEMAP.load_descriptor(ptr) {
            Some(ptr) => ptr,
            None => {
                panic!("PTR was not allocated with this allocator");
            }
        };
        let desc = unsafe { &*desc_ptr };

        let size_class = match desc.size_class() {
            Some(s) => s,
            None => {
                self.heap.free_large(ptr, layout, &PAGEMAP);
                return;
            }
        };

        CACHE.with(|raw| {
            let mut cache = raw.borrow_mut();

            if cache.add_block(size_class, ptr).is_err() {
                self.heap.flush_cache(&mut cache, size_class, &PAGEMAP);
                cache.add_block(size_class, ptr).unwrap();
            };
        });
    }
}

unsafe impl GlobalAlloc for Allocator {
    unsafe fn alloc(&self, layout: std::alloc::Layout) -> *mut u8 {
        unsafe { self.allocate(layout) }
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: std::alloc::Layout) {
        unsafe { self.deallocate(ptr, layout) }
    }
}