//! The kernel heap: dynamic allocation for the kernel.
//!
//! Where the frame allocator ([pmm]) hands out fixed 4 KiB physical frames, the
//! heap hands out arbitrary byte-sized blocks from a virtual region, growing on
//! demand by mapping fresh frames into it (arch.mapPage) — the first real user of
//! the VMM (see docs/paging.md).
//!
//! The algorithm is a first-fit free list: an address-ordered singly linked list
//! of free blocks, split on allocation and coalesced with neighbours on free. It
//! is exposed as a std.mem.Allocator, so the kernel can use std containers.
//!
//! Not yet concurrency-safe: it assumes a single caller and no allocation from
//! interrupt handlers (ours don't). A lock comes with threads/SMP.

const std = @import("std");
const danos = @import("danos");
const arch = @import("arch");
const pmm = @import("pmm.zig");

const page_size = danos.page_size;

/// Virtual base of the heap: the start of the higher half, which is unmapped and
/// well clear of the identity-mapped low half. (Canonical on x86_64; an arch that
/// splits the address space differently would choose its own.)
const heap_base: usize = 0xFFFF_8000_0000_0000;
/// Cap on heap growth for now.
const heap_max: usize = 64 * 1024 * 1024;

/// A block header, placed at the start of every block. While the block is free it
/// also links into the free list via `next`.
const Block = extern struct {
    size: usize, // total block size in bytes, including this header; a multiple of 16
    next: ?*Block, // free-list link (only meaningful while free)
};

const header_size = @sizeOf(Block); // 16
const min_block = header_size + 16; // smallest block worth splitting off

var free_list: ?*Block = null;
var heap_end: usize = heap_base; // [heap_base, heap_end) is currently mapped

fn alignUp(value: usize, alignment: usize) usize {
    return (value + alignment - 1) & ~(alignment - 1);
}

fn payloadOf(block: *Block) [*]u8 {
    return @ptrFromInt(@intFromPtr(block) + header_size);
}

/// Bring the heap up with an initial mapped region.
pub fn init() void {
    free_list = null;
    heap_end = heap_base;
    _ = grow(page_size);
}

/// Map more pages onto the end of the heap and add them as a free block. Returns
/// false if out of heap virtual space or out of physical frames.
fn grow(min_bytes: usize) bool {
    const start = heap_end;
    const bytes = alignUp(min_bytes, page_size);
    if (start + bytes > heap_base + heap_max) return false;

    var virt = start;
    while (virt < start + bytes) : (virt += page_size) {
        const frame = pmm.alloc() orelse return false;
        arch.mapPage(virt, frame, true);
    }
    heap_end = start + bytes;

    const block: *Block = @ptrFromInt(start);
    block.size = bytes;
    insertFree(block); // coalesces with the previous tail block if adjacent
    return true;
}

/// Insert a block into the address-ordered free list, coalescing with the
/// physically adjacent free blocks on either side.
fn insertFree(block: *Block) void {
    var prev: ?*Block = null;
    var cur = free_list;
    while (cur) |c| : (cur = c.next) {
        if (@intFromPtr(c) > @intFromPtr(block)) break;
        prev = c;
    }

    block.next = cur;
    if (prev) |p| p.next = block else free_list = block;

    // Merge forward into `cur` if they're contiguous.
    if (cur) |c| {
        if (@intFromPtr(block) + block.size == @intFromPtr(c)) {
            block.size += c.size;
            block.next = c.next;
        }
    }
    // Merge `prev` forward into `block` if they're contiguous.
    if (prev) |p| {
        if (@intFromPtr(p) + p.size == @intFromPtr(block)) {
            p.size += block.size;
            p.next = block.next;
        }
    }
}

/// Allocate `len` bytes (16-byte aligned), or null if out of memory.
fn rawAlloc(len: usize) ?[*]u8 {
    const need = alignUp(header_size + len, 16);

    var attempts: u32 = 0;
    while (attempts < 2) : (attempts += 1) {
        var prev: ?*Block = null;
        var cur = free_list;
        while (cur) |block| : ({
            prev = block;
            cur = block.next;
        }) {
            if (block.size < need) continue;

            if (block.size >= need + min_block) {
                // Split: carve `need` off the front, leave the rest free.
                const rest: *Block = @ptrFromInt(@intFromPtr(block) + need);
                rest.size = block.size - need;
                rest.next = block.next;
                if (prev) |p| p.next = rest else free_list = rest;
                block.size = need;
            } else {
                // Take the whole block.
                if (prev) |p| p.next = block.next else free_list = block.next;
            }
            return payloadOf(block);
        }

        // Nothing fit: grow and try once more.
        if (!grow(need)) return null;
    }
    return null;
}

fn rawFree(ptr: [*]u8) void {
    const block: *Block = @ptrFromInt(@intFromPtr(ptr) - header_size);
    insertFree(block);
}

// --- std.mem.Allocator interface -----------------------------------------

pub fn allocator() std.mem.Allocator {
    return .{ .ptr = undefined, .vtable = &vtable };
}

const vtable = std.mem.Allocator.VTable{
    .alloc = allocImpl,
    .resize = resizeImpl,
    .remap = remapImpl,
    .free = freeImpl,
};

fn allocImpl(_: *anyopaque, len: usize, alignment: std.mem.Alignment, _: usize) ?[*]u8 {
    // Blocks are 16-byte aligned; larger alignments aren't supported yet.
    if (alignment.toByteUnits() > 16) return null;
    return rawAlloc(len);
}

fn resizeImpl(_: *anyopaque, _: []u8, _: std.mem.Alignment, _: usize, _: usize) bool {
    return false; // no in-place resize; the caller reallocates
}

fn remapImpl(_: *anyopaque, _: []u8, _: std.mem.Alignment, _: usize, _: usize) ?[*]u8 {
    return null;
}

fn freeImpl(_: *anyopaque, memory: []u8, _: std.mem.Alignment, _: usize) void {
    rawFree(memory.ptr);
}