//! Global structures for the allocators.

use crate::arch::{
    paging::{PageTable, PageTableEntry, ASID, PAGE_TABLE_BITS, PAGE_TABLE_LEVELS},
    MAX_PAGE_SIZE_BITS, PAGE_SIZE, PAGE_SIZE_BITS,
};
use contracts::requires;
use core::ptr::NonNull;
use spin::mutex::FairMutex;
use vernos_alloc_buddy::BuddyAllocator;
use vernos_utils::BelieveMeSend;

/// The global instance of the physical page allocator.
static BUDDY_ALLOCATOR: FairMutex<
    Option<BuddyAllocator<PAGE_SIZE, PAGE_SIZE_BITS, { 1 + MAX_PAGE_SIZE_BITS - PAGE_SIZE_BITS }>>,
> = FairMutex::new(None);

/// The global kernel page table.
static KERNEL_PAGE_TABLE: FairMutex<BelieveMeSend<Option<NonNull<PageTable>>>> =
    FairMutex::new(BelieveMeSend(None));

/// Initializes the allocator and enables paging.
///
/// # Safety
///
/// - Paging must not have been enabled previously.
/// - The buddy allocator must be valid.
#[requires(BUDDY_ALLOCATOR.lock().is_none())]
#[ensures(BUDDY_ALLOCATOR.lock().is_some())]
#[requires(KERNEL_PAGE_TABLE.lock().is_none())]
#[ensures(KERNEL_PAGE_TABLE.lock().is_some())]
pub unsafe fn init(
    mut buddy_allocator: BuddyAllocator<
        'static,
        PAGE_SIZE,
        PAGE_SIZE_BITS,
        { 1 + MAX_PAGE_SIZE_BITS - PAGE_SIZE_BITS },
    >,
) {
    // Allocate a page to use (for now) as the global kernel page table. Later we'll actually
    // replace it with the hart0 initial stack's page, since we'll never free the root page of the
    // kernel page table, and we can't return that page to the buddy allocator anyway.
    let mut page_table = buddy_allocator
        .alloc_zeroed::<PageTable>()
        .expect("failed to allocate the kernel page table");

    // Create identity mappings for the lower half of memory.
    for (i, entry) in page_table
        .as_mut()
        .iter_mut()
        .enumerate()
        .take(1 << (PAGE_TABLE_BITS - 1))
    {
        let addr = (i as u64) << ((PAGE_TABLE_LEVELS - 1) * PAGE_TABLE_BITS + PAGE_SIZE_BITS);
        let mut pte = PageTableEntry::default();
        pte.set_valid(true)
            .set_rwx(true, true, true)
            .set_global(true)
            .set_addr(addr);
        *entry = pte;
    }

    // Set the page table as the current page table.
    page_table.as_mut().make_current(ASID::KERNEL);

    // Save the buddy allocator and kernel page table.
    *BUDDY_ALLOCATOR.lock() = Some(buddy_allocator);
    KERNEL_PAGE_TABLE.lock().0 = Some(page_table);

    // Print the page table after this.
    vernos_utils::dbg!(page_table.as_mut());
}