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

#![feature(allocator_api)]
#![feature(const_fn)]
#![feature(try_trait)]
#![no_std]

use core::alloc::{GlobalAlloc, Layout};
use core::ops::Try;
use core::ptr;
use uefi::memory::MemoryType;
use uefi::system::SystemTable;

static mut UEFI: *mut SystemTable = 0 as *mut SystemTable;

pub unsafe fn init(table: &'static mut SystemTable) {
    UEFI = table;
}

fn get_uefi() -> Option<&'static mut SystemTable> {
    unsafe {
        if UEFI as usize == 0 {
            None
        } else {
            Some(&mut *UEFI)
        }
    }
}

pub struct Allocator;

unsafe impl GlobalAlloc for Allocator {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        if let Some(ref mut uefi) = get_uefi() {
            let mut ptr = 0;
            if let Err(_) = (uefi.BootServices.AllocatePool)(MemoryType::EfiLoaderData, layout.size(), &mut ptr).into_result() {
                ptr::null_mut()
            } else {
                ptr as *mut u8
            }
        } else {
            panic!("__rust_allocate: uefi not initialized");
        }
    }

    unsafe fn dealloc(&self, ptr: *mut u8, _layout: Layout) {
        if let Some(ref mut uefi) = get_uefi() {
            let _ = (uefi.BootServices.FreePool)(ptr as usize);
        } else {
            panic!("__rust_deallocate: uefi not initialized");
        }
    }
}