#include "kernel/lib.h"
#include "kernel/kalloc.h"
#include "kernel/spinlock.h"

extern uint32_t _heap_start;
extern uint32_t _heap_end;

union Page
{
    union Page* next;
    char dummy[PAGE_SIZE];
};
typedef union Page Page;
static_assert(sizeof(Page) == PAGE_SIZE);

static inline void* align_page_down(void* p)
{
    return (void*)((uint32_t)p & ~(PAGE_SIZE - 1));
}

static inline void* align_page_up(void* p)
{
    return align_page_down((char*)p + PAGE_SIZE - 1);
}

static inline bool is_page_aligned(void* p)
{
    return ((uint32_t)p & (PAGE_SIZE - 1)) == 0;
}

static Page* g_page_start;
static Page* g_page_end;
static Page* g_free_list;
static Spinlock g_lock;

void kalloc_init()
{
    spinlock_init(&g_lock);
    
    g_free_list  = NULL;
    g_page_start = (Page*)align_page_up(&_heap_start);
    g_page_end   = (Page*)align_page_down(&_heap_end);

    for (Page* p = g_page_end - 1; p >= g_page_start; --p)
    {
        p->next = g_free_list;
        g_free_list = p;
    }
}

void* kalloc()
{
    spinlock_acquire(&g_lock);
    Page* page = g_free_list;
    if (!page)
    {
        panic("kalloc: Out of memory");
    }
    g_free_list = page->next;
    spinlock_release(&g_lock);
    return page;
}

void* kzalloc()
{
    void* page = kalloc();
    uint32_t* p = (uint32_t*)page;
    for (uint32_t i = 0; i < PAGE_SIZE / 4; ++i)
    {
        *p++ = 0U;
    }
    return page;
}

void kfree(void* ptr)
{
    if (ptr < (void*)g_page_start || ptr >= (void*)g_page_end || !is_page_aligned(ptr))
    {
        panic("kfree: Invalid page");
    }
    
    spinlock_acquire(&g_lock);
    Page* page = (Page*)ptr;
    page->next = g_free_list;
    g_free_list = page;
    spinlock_release(&g_lock);
}