#![feature(const_fn)]
#![feature(alloc, allocator_api)]
#![no_std]
#[cfg(test)]
#[macro_use]
extern crate std;
#[cfg(feature = "use_spin")]
extern crate spin;
extern crate alloc;
use alloc::alloc::{Alloc, AllocErr, Layout};
use core::alloc::GlobalAlloc;
use core::cmp::{max, min};
use core::fmt;
use core::mem::size_of;
#[cfg(feature = "use_spin")]
use core::ops::Deref;
use core::ptr::NonNull;
#[cfg(feature = "use_spin")]
use spin::Mutex;
pub mod linked_list;
#[cfg(test)]
mod test;
pub struct Heap {
free_list: [linked_list::LinkedList; 32],
user: usize,
allocated: usize,
total: usize,
}
impl Heap {
pub const fn new() -> Self {
Heap {
free_list: [linked_list::LinkedList::new(); 32],
user: 0,
allocated: 0,
total: 0,
}
}
pub const fn empty() -> Self {
Self::new()
}
pub unsafe fn add_to_heap(&mut self, start: usize, end: usize) {
assert!(start <= end);
let mut total = 0;
let mut current_start = start;
while current_start + size_of::<usize>() <= end {
let lowbit = current_start & (!current_start + 1);
let size = min(lowbit, prev_power_of_two(end - current_start));
total += size;
self.free_list[size.trailing_zeros() as usize].push(current_start as *mut usize);
current_start += size;
}
self.total += total;
}
pub unsafe fn init(&mut self, start: usize, size: usize) {
self.add_to_heap(start, start + size);
}
pub fn alloc(&mut self, layout: Layout) -> Result<NonNull<u8>, AllocErr> {
let size = max(
layout.size().next_power_of_two(),
max(layout.align(), size_of::<usize>()),
);
let class = size.trailing_zeros() as usize;
for i in class..self.free_list.len() {
if !self.free_list[i].is_empty() {
for j in (class + 1..i + 1).rev() {
if let Some(block) = self.free_list[j].pop() {
unsafe {
self.free_list[j - 1].push((block as usize + (1 << (j - 1))) as *mut usize);
self.free_list[j - 1].push(block);
}
} else {
return Err(AllocErr {});
}
}
let result = NonNull::new(self.free_list[class]
.pop()
.expect("current block should have free space now")
as *mut u8);
if let Some(result) = result {
self.user += layout.size();
self.allocated += size;
return Ok(result);
} else {
return Err(AllocErr {});
}
}
}
Err(AllocErr {})
}
pub fn dealloc(&mut self, ptr: NonNull<u8>, layout: Layout) {
let size = max(
layout.size().next_power_of_two(),
max(layout.align(), size_of::<usize>()),
);
let class = size.trailing_zeros() as usize;
unsafe {
self.free_list[class].push(ptr.as_ptr() as *mut usize);
let mut current_ptr = ptr.as_ptr() as usize;
let mut current_class = class;
while current_class < self.free_list.len() {
let buddy = current_ptr ^ (1 << current_class);
let mut flag = false;
for block in self.free_list[current_class].iter_mut() {
if block.value() as usize == buddy {
block.pop();
flag = true;
break;
}
}
if flag {
self.free_list[current_class].pop();
current_ptr = min(current_ptr, buddy);
current_class += 1;
self.free_list[current_class].push(current_ptr as *mut usize);
} else {
break;
}
}
}
self.user -= layout.size();
self.allocated -= size;
}
}
impl fmt::Debug for Heap {
fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
fmt.debug_struct("Heap")
.field("user", &self.user)
.field("allocated", &self.allocated)
.field("total", &self.total)
.finish()
}
}
unsafe impl Alloc for Heap {
unsafe fn alloc(&mut self, layout: Layout) -> Result<NonNull<u8>, AllocErr> {
self.alloc(layout)
}
unsafe fn dealloc(&mut self, ptr: NonNull<u8>, layout: Layout) {
self.dealloc(ptr, layout)
}
}
#[cfg(feature = "use_spin")]
pub struct LockedHeap(Mutex<Heap>);
#[cfg(feature = "use_spin")]
impl LockedHeap {
pub const fn new() -> LockedHeap {
LockedHeap(Mutex::new(Heap::new()))
}
pub const fn empty() -> LockedHeap {
LockedHeap(Mutex::new(Heap::new()))
}
}
#[cfg(feature = "use_spin")]
impl Deref for LockedHeap {
type Target = Mutex<Heap>;
fn deref(&self) -> &Mutex<Heap> {
&self.0
}
}
#[cfg(feature = "use_spin")]
unsafe impl GlobalAlloc for LockedHeap {
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
self.0
.lock()
.alloc(layout)
.ok()
.map_or(0 as *mut u8, |allocation| allocation.as_ptr())
}
unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
self.0
.lock()
.dealloc(NonNull::new_unchecked(ptr), layout)
}
}
fn prev_power_of_two(num: usize) -> usize {
1 << (8 * (size_of::<usize>()) - num.leading_zeros() as usize - 1)
}