alloc_cat 1.1.1

use core::fmt;
use core::sync::atomic::{AtomicUsize, Ordering};
use memmap2::MmapMut;
use spin::Once;

use super::{Heap, PAGE_SIZE};


#[derive(Debug)]
pub struct MmapHeap {
    pages: AtomicUsize,
    max_size: usize,
    region: Once<MmapMut>,
}

impl MmapHeap {
    pub const fn new(max_size: usize) -> MmapHeap {
        MmapHeap {
            pages: AtomicUsize::new(0),
            max_size,
            region: Once::new(),
        }
    }

    fn get_region(&self) -> &MmapMut {
        self.region.call_once(|| MmapMut::map_anon(self.max_size).unwrap())
    }

    fn get_pages(&self) -> usize {
        self.pages.load(Ordering::Relaxed)
    }

    fn increase_pages(&self, pages: usize) -> usize {
        self.pages.fetch_add(pages, Ordering::SeqCst)
    }

    fn decrease_pages(&self, pages: usize) -> usize {
        self.pages.fetch_sub(pages, Ordering::SeqCst)
    }

    pub fn reset(&self) {
        self.pages.store(1, Ordering::SeqCst);
    }
}

impl Heap for MmapHeap {
    fn heap_start(&self) -> usize {
        self.get_region().as_ptr() as usize
    }

    fn heap_end(&self) -> usize {
        self.heap_start() + self.get_pages() * PAGE_SIZE
    }

    fn heap_grow(&self, pages: usize) -> Option<*mut u8> {
        let prev_pages = self.increase_pages(pages);
        if (prev_pages + pages) * PAGE_SIZE <= self.max_size {
            Some((self.heap_start() + prev_pages * PAGE_SIZE) as *mut u8)
        } else {
            self.decrease_pages(pages);
            None
        }
    }
}

impl fmt::Display for MmapHeap {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "MmapHeap(pages={}, max_size={})", self.get_pages(), self.max_size)
    }
}