// heap.c - user-mode page-backed heap
#include <stddef.h>
#include <stdint.h>
#include <string.h>
#include <intrin.h>
#include <ntlite/syscall.h>

#define USER_HEAP_PAGE_SIZE 4096u
#define USER_HEAP_MIN_PAGES 16u
#define USER_HEAP_ALIGN     16u
#define USER_HEAP_MAGIC     0x48424C4Bu /* KLBH */

/*
 * Keep the mask pointer-sized. A 32-bit unsigned mask truncates high VA bits.
 */
#define ALIGN_UP(v, a) (((v) + ((uintptr_t)(a) - 1)) & ~((uintptr_t)(a) - 1))

typedef struct HEAP_BLOCK {
    uint32_t Magic;
    uint32_t Free;
    size_t Size;
    struct HEAP_BLOCK* Prev;
    struct HEAP_BLOCK* Next;
} HEAP_BLOCK;

typedef struct HEAP_SEGMENT {
    struct HEAP_SEGMENT* Prev;
    struct HEAP_SEGMENT* Next;
    size_t PageCount;
    size_t Size;
    HEAP_BLOCK* First;
} HEAP_SEGMENT;

static volatile long gHeapLock;
static HEAP_SEGMENT* gSegments;

static void HeapLock(void) {
    while (_InterlockedCompareExchange(&gHeapLock, 1, 0) != 0)
        _mm_pause();
    _ReadWriteBarrier();
}

static void HeapUnlock(void) {
    _ReadWriteBarrier();
    _InterlockedExchange(&gHeapLock, 0);
}

static size_t AlignSize(size_t size) {
    return ALIGN_UP(size ? size : 1, USER_HEAP_ALIGN);
}

static void* BlockPayload(HEAP_BLOCK* block) {
    return (void*)((uint8_t*)block + sizeof(HEAP_BLOCK));
}

static HEAP_BLOCK* PayloadBlock(void* ptr) {
    return (HEAP_BLOCK*)((uint8_t*)ptr - sizeof(HEAP_BLOCK));
}

static HEAP_SEGMENT* SegmentForPointer(void* ptr) {
    uintptr_t p = (uintptr_t)ptr;
    for (HEAP_SEGMENT* seg = gSegments; seg; seg = seg->Next) {
        uintptr_t base = (uintptr_t)seg;
        uintptr_t end = base + seg->Size;
        if (p > base && p < end)
            return seg;
    }
    return (HEAP_SEGMENT*)0;
}

static void SplitBlock(HEAP_BLOCK* block, size_t size) {
    size_t aligned = AlignSize(size);
    if (block->Size < aligned + sizeof(HEAP_BLOCK) + USER_HEAP_ALIGN)
        return;

    HEAP_BLOCK* tail = (HEAP_BLOCK*)((uint8_t*)BlockPayload(block) + aligned);
    tail->Magic = USER_HEAP_MAGIC;
    tail->Free = 1;
    tail->Size = block->Size - aligned - sizeof(HEAP_BLOCK);
    tail->Prev = block;
    tail->Next = block->Next;
    if (tail->Next) tail->Next->Prev = tail;

    block->Size = aligned;
    block->Next = tail;
}

static void CoalesceBlock(HEAP_BLOCK* block) {
    if (block->Next && block->Next->Magic == USER_HEAP_MAGIC && block->Next->Free) {
        HEAP_BLOCK* next = block->Next;
        block->Size += sizeof(HEAP_BLOCK) + next->Size;
        block->Next = next->Next;
        if (block->Next) block->Next->Prev = block;
    }

    if (block->Prev && block->Prev->Magic == USER_HEAP_MAGIC && block->Prev->Free) {
        HEAP_BLOCK* prev = block->Prev;
        prev->Size += sizeof(HEAP_BLOCK) + block->Size;
        prev->Next = block->Next;
        if (prev->Next) prev->Next->Prev = prev;
    }
}

static HEAP_SEGMENT* AddSegment(size_t payloadSize) {
    size_t need = sizeof(HEAP_SEGMENT) + sizeof(HEAP_BLOCK) + AlignSize(payloadSize);
    size_t pages = (need + USER_HEAP_PAGE_SIZE - 1) / USER_HEAP_PAGE_SIZE;
    if (pages < USER_HEAP_MIN_PAGES) pages = USER_HEAP_MIN_PAGES;

    void* base = 0;
    uint32_t s = SysVmAlloc(0, pages, SYS_VM_READWRITE, SYS_VM_FLAG_NONE, &base);
    if (s != SYS_STATUS_SUCCESS || !base)
        return (HEAP_SEGMENT*)0;

    HEAP_SEGMENT* seg = (HEAP_SEGMENT*)base;
    seg->Prev = 0;
    seg->Next = gSegments;
    if (gSegments) gSegments->Prev = seg;
    gSegments = seg;
    seg->PageCount = pages;
    seg->Size = pages * USER_HEAP_PAGE_SIZE;

    uintptr_t blockAddr = ALIGN_UP((uintptr_t)(seg + 1), USER_HEAP_ALIGN);
    HEAP_BLOCK* block = (HEAP_BLOCK*)blockAddr;
    block->Magic = USER_HEAP_MAGIC;
    block->Free = 1;
    block->Size = seg->Size - (blockAddr - (uintptr_t)seg) - sizeof(HEAP_BLOCK);
    block->Prev = 0;
    block->Next = 0;
    seg->First = block;
    return seg;
}

static void RemoveSegmentIfEmpty(HEAP_SEGMENT* seg) {
    if (!seg || !seg->First || !seg->First->Free || seg->First->Prev || seg->First->Next)
        return;

    if (seg->Prev) seg->Prev->Next = seg->Next;
    else gSegments = seg->Next;
    if (seg->Next) seg->Next->Prev = seg->Prev;

    SysVmFree(seg);
}

void* _malloc(size_t size) {
    if (!size) return 0;
    size = AlignSize(size);

    HeapLock();
    for (;;) {
        for (HEAP_SEGMENT* seg = gSegments; seg; seg = seg->Next) {
            for (HEAP_BLOCK* block = seg->First; block; block = block->Next) {
                if (block->Magic == USER_HEAP_MAGIC && block->Free && block->Size >= size) {
                    SplitBlock(block, size);
                    block->Free = 0;
                    void* out = BlockPayload(block);
                    HeapUnlock();
                    return out;
                }
            }
        }

        if (!AddSegment(size)) {
            HeapUnlock();
            return 0;
        }
    }
}

void _free(void* ptr) {
    if (!ptr) return;

    HeapLock();
    HEAP_SEGMENT* seg = SegmentForPointer(ptr);
    if (!seg) {
        HeapUnlock();
        return;
    }

    HEAP_BLOCK* block = PayloadBlock(ptr);
    if (block->Magic != USER_HEAP_MAGIC || block->Free) {
        HeapUnlock();
        return;
    }

    block->Free = 1;
    CoalesceBlock(block);
    RemoveSegmentIfEmpty(seg);
    HeapUnlock();
}

void* _calloc(size_t count, size_t size) {
    if (count && size > ((size_t)-1) / count)
        return 0;
    size_t total = count * size;
    void* ptr = _malloc(total);
    if (ptr) memset(ptr, 0, total);
    return ptr;
}

void* _realloc(void* ptr, size_t newSize) {
    if (!ptr) return _malloc(newSize);
    if (!newSize) {
        _free(ptr);
        return 0;
    }

    HeapLock();
    HEAP_BLOCK* block = PayloadBlock(ptr);
    if (block->Magic != USER_HEAP_MAGIC || block->Free) {
        HeapUnlock();
        return 0;
    }
    size_t oldSize = block->Size;
    if (oldSize >= newSize) {
        SplitBlock(block, newSize);
        HeapUnlock();
        return ptr;
    }
    HeapUnlock();

    void* newPtr = _malloc(newSize);
    if (!newPtr) return 0;
    memcpy(newPtr, ptr, oldSize);
    _free(ptr);
    return newPtr;
}