RustForge

rust-kernel

Rustbelt: a hobby x86_64 kernel written in Rust

rust-kernel / src / allocator / bump.rs
2427 bytes ยท raw
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
//! Bump allocator: hand out memory by moving a pointer forward ("bumping").
//! Super fast and tiny, but memory is only reclaimed when *every*
//! allocation has been freed. Not active; kept for learning.

use super::{Locked, align_up};
use alloc::alloc::{GlobalAlloc, Layout};
use core::ptr;

pub struct BumpAllocator {
    heap_start: usize,
    heap_end: usize,
    next: usize,        // start of the unused part of the heap
    allocations: usize, // number of live allocations
}

impl BumpAllocator {
    pub const fn new() -> Self {
        BumpAllocator { heap_start: 0, heap_end: 0, next: 0, allocations: 0 }
    }

    /// Unsafe: the range must be mapped and unused; call only once.
    pub unsafe fn init(&mut self, heap_start: usize, heap_size: usize) {
        self.heap_start = heap_start;
        self.heap_end = heap_start.saturating_add(heap_size);
        self.next = heap_start;
    }
}

// The heap lock is taken with interrupts off: code that runs with
// interrupts off (the scheduler, wait queues, ...) may allocate, and if it
// found the lock held by a thread that was preempted mid-allocation, it
// would spin forever -- that thread can't run again until interrupts are on.
unsafe impl GlobalAlloc for Locked<BumpAllocator> {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        x86_64::instructions::interrupts::without_interrupts(|| unsafe { self.alloc_inner(layout) })
    }
    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        x86_64::instructions::interrupts::without_interrupts(|| unsafe { self.dealloc_inner(ptr, layout) })
    }
}

impl Locked<BumpAllocator> {
    unsafe fn alloc_inner(&self, layout: Layout) -> *mut u8 {
        let mut bump = self.lock();
        let alloc_start = align_up(bump.next, layout.align());
        let alloc_end = match alloc_start.checked_add(layout.size()) {
            Some(end) => end,
            None => return ptr::null_mut(),
        };
        if alloc_end > bump.heap_end {
            ptr::null_mut() // out of memory
        } else {
            bump.next = alloc_end;
            bump.allocations += 1;
            alloc_start as *mut u8
        }
    }

    unsafe fn dealloc_inner(&self, _ptr: *mut u8, _layout: Layout) {
        let mut bump = self.lock();
        bump.allocations -= 1;
        if bump.allocations == 0 {
            bump.next = bump.heap_start; // everything freed: start over
        }
    }
}