zing/src/mem/phys.zig

//! Physical memory management implementation.

const std = @import("std");

const kernel = @import("../kernel.zig");
const log = @import("../debug.zig").log;
const mem = @import("../mem.zig");
const vmm = @import("vmm.zig");
const sync = @import("../sync.zig");

const Range = @import("../util/range.zig").Range;
const Spinlock = sync.Spinlock;

/// Represents a single region of physical memory (reserved or available).
pub const MemoryRegion = struct {
    /// Name string, used to represent where the memory comes from.
    name: []const u8,
    /// Page frame number range of the region.
    range: Range(u64),
};

const Bitmap = struct {
    data: []u64,

    const Self = @This();

    pub const empty: Self = .{ .data = &.{} };

    fn get_bit(self: *Self, index: usize) u1 {
        const word_index = index / 64;
        const bit_index = index % 64;
        const masked = self.data[word_index] & (@as(u64, 1) << @intCast(bit_index));
        return if (masked == 0) 0 else 1;
    }

    fn set_bit(self: *Self, index: usize) void {
        const word_index = index / 64;
        const bit_index = index % 64;
        self.data[word_index] |= (@as(u64, 1) << @intCast(bit_index));
    }

    fn clear_bit(self: *Self, index: usize) void {
        const word_index = index / 64;
        const bit_index = index % 64;
        self.data[word_index] &= ~(@as(u64, 1) << @intCast(bit_index));
    }
};

const PhysicalMemoryManager = struct {
    memory_start: u64,
    last_free: usize,
    len: usize,

    /// Each bit represents a page, there can be more u64s than needed
    usage_bitmap: Bitmap,
    page_refcounters: []u32,

    const empty: @This() = .{
        .memory_start = 0,
        .last_free = 0,
        .len = 0,
        .usage_bitmap = .empty,
        .page_refcounters = &.{},
    };

    fn alloc_page(self: *@This()) ?mem.PhysicalAddress {
        for (self.last_free..self.len) |i| {
            if (!self.is_page_used(i)) {
                self.page_refcounters[i] += 1;
                self.set_page_used(i);
                self.last_free = (i + 1) % self.len;
                return .{ .raw = self.memory_start + i * vmm.PAGE_SIZE };
            }
        }
        for (0..self.last_free) |i| {
            if (!self.is_page_used(i)) {
                self.page_refcounters[i] += 1;
                self.set_page_used(i);
                self.last_free = (i + 1) % self.len;
                return .{ .raw = self.memory_start + i * vmm.PAGE_SIZE };
            }
        }
        return null;
    }

    fn alloc_pages(self: *@This(), count: usize) ?mem.PhysicalAddress {
        if (self.last_free + count < self.len) {
            if (self.alloc_from(self.last_free, self.len, count)) |p| {
                return p;
            }
        }
        return self.alloc_from(0, self.last_free, count);
    }

    fn alloc_from(self: *@This(), start: usize, end: usize, count: usize) ?mem.PhysicalAddress {
        for (start..end) |i| {
            const taken = taken: {
                for (0..count) |j|
                    if (self.is_page_used(i + j))
                        break :taken true;
                break :taken false;
            };

            if (!taken) {
                for (0..count) |j| {
                    self.page_refcounters[i + j] = 1;
                    self.set_page_used(i + j);
                }
                return .{ .raw = self.memory_start + i * vmm.PAGE_SIZE };
            }
        }

        return null;
    }

    fn valid_index(self: *@This(), page: mem.PhysicalAddress) usize {
        if (page.raw < self.memory_start) {
            log.panic("free_page: invalid page 0x{x}: outside of the allocation range", .{page.raw});
        }
        const index = (page.raw - self.memory_start) / vmm.PAGE_SIZE;
        if (index >= self.len) {
            log.panic("free_page: invalid page 0x{x}: outside of the allocation range", .{page.raw});
        }
        return index;
    }

    fn free_page(self: *@This(), page: mem.PhysicalAddress) void {
        const index = self.valid_index(page);
        if (!self.is_page_used(index)) {
            log.panic("free_page: double free of page 0x{x} detected", .{page.raw});
        }
        self.page_refcounters[index] -= 1;
        if (self.page_refcounters[index] == 0) {
            self.clear_page_used(index);
            self.last_free = index;
        }
    }

    fn is_page_used(self: *@This(), index: usize) bool {
        return self.usage_bitmap.get_bit(index) == 1;
    }

    fn set_page_used(self: *@This(), index: usize) void {
        self.usage_bitmap.set_bit(index);
    }

    fn clear_page_used(self: *@This(), index: usize) void {
        self.usage_bitmap.clear_bit(index);
    }
};

var g_memory_regions: std.BoundedArray(MemoryRegion, 16) = .{};
var g_reserved_regions: std.BoundedArray(MemoryRegion, 16) = .{};
var g_physical_memory_lock = Spinlock{};
var g_physical_memory = PhysicalMemoryManager.empty;

/// Adds an available memory region to the list.
///
/// `base` and `size` are in bytes. Regions are page-aligned "inwards", meaning the function will
/// only add the range of full pages of the specified region. If a combination is provided that
/// does not yield any full 4KiB pages (e.g. `base=0x1234, size=0x123`), it is ignored.
///
/// # Note
///
/// Only meaningful to call before calling `init()`.
pub fn add_memory_region(name: []const u8, base: u64, size: u64) void {
    log.info("Memory: '{s}', base 0x{x}, size 0x{x}", .{ name, base, size });
    const start = vmm.L3.align_up(base) / vmm.L3.SIZE;
    const len = vmm.L3.align_down(base + size) / vmm.L3.SIZE - start;
    if (len > 0) {
        g_memory_regions.append(.{ .name = name, .range = .{ .start = start, .len = len } }) //
        catch @panic("memory regions overflow");
    }
}

/// Adds an reserved memory region to the list.
///
/// `base` and `size` are in bytes. Regions are page-aligned "outwards", meaning that the
/// reservation extends to any pages affected by the specified region.
///
/// # Note
///
/// Only meaningful to call before calling `init()`.
pub fn add_reserved_region(name: []const u8, base: u64, size: u64) void {
    log.info("Reserved: '{s}', base 0x{x}, size 0x{x}", .{ name, base, size });
    const start = base / vmm.L3.SIZE;
    const len = vmm.L3.align_up(base + size) / vmm.L3.SIZE - start;
    if (len > 0) {
        g_reserved_regions.append(.{ .name = name, .range = .{ .start = start, .len = len } }) //
        catch @panic("reserved regions overflow");
    }
}

fn is_reserved_in(page_index: u64) ?*const MemoryRegion {
    for (0..g_reserved_regions.len) |i| {
        const region = &g_reserved_regions.buffer[i];
        if (page_index >= region.range.start and page_index < region.range.end()) {
            return region;
        }
    }
    return null;
}

fn alloc_from_region(region: *const MemoryRegion, reason: []const u8, page_count: usize) ?u64 {
    var offset = @as(u64, 0);
    while (offset < region.range.len) {
        var taken: ?*const MemoryRegion = null;
        for (0..page_count) |i| {
            if (is_reserved_in(region.range.start + offset + i)) |resv| {
                taken = resv;
                break;
            }
        }

        if (taken) |resv| {
            offset = (resv.range.start + resv.range.len) - region.range.start;
            continue;
        }

        const base = (region.range.start + offset) * vmm.L3.SIZE;
        add_reserved_region(reason, base, page_count * vmm.PAGE_SIZE);
        return base;
    }
    return null;
}

/// Allocates a slice of type `T` that spans the `page_count` pages.
fn alloc_slice_pages(comptime T: type, reason: []const u8, page_count: usize) []T {
    for (g_memory_regions.constSlice()) |region| {
        if (alloc_from_region(&region, reason, page_count)) |physAddress| {
            const vaddr = (mem.PhysicalAddress{ .raw = physAddress }).virtualize();
            const len = (page_count * vmm.PAGE_SIZE) / @sizeOf(T);
            const ptr: [*]T = @ptrFromInt(vaddr);
            const slice: []T = ptr[0..len];
            return slice;
        }
    }
    @panic("Failed to allocate a slice");
}

/// Allocates a slice of type `T` that has at least `min_len` items, allocates in 4KiB pages.
/// The items are zeroed out.
fn alloc_slice(comptime T: type, reason: []const u8, min_len: usize) []T {
    const min_alloc_bytes = min_len * @sizeOf(T);
    // Round up to make sure we have enough space for the data
    const needed_pages = vmm.L3.page_count(min_alloc_bytes);
    const slice = alloc_slice_pages(T, reason, needed_pages);
    const slice_as_bytes = std.mem.sliceAsBytes(slice);
    @memset(slice_as_bytes, 0);
    return slice;
}

/// Allocates a bitmap that has at least `bits_required` total bits.
/// It can have more since we allocate 4KiB pages and the backing type is []u64
fn alloc_bitmap(bits_required: usize) Bitmap {
    // Round up to the upper u64 that has at least `pages` bits
    const bitmap_entries = (bits_required + 63) / 64;
    const bitmap = alloc_slice(u64, "bitmap", bitmap_entries);
    return .{ .data = bitmap };
}
// TODO: combine refcounters and bitmap allocation into a single chunk
fn alloc_refcounters(count: usize) []u32 {
    const refcounters = alloc_slice(u32, "refcounters", count);
    @memset(refcounters, std.math.maxInt(u32));
    return refcounters;
}

/// Initializes the physical memory management.
///
/// # Notes
///
/// Calls to `add***Region()` functions have no meaning past this point, so all the memory
/// present in the system, along with memory reservations, should be added **prior** to this point.
pub fn init() void {
    var memory_start: u64 = std.math.maxInt(u64);
    var memory_end: u64 = std.math.minInt(u64);

    for (g_memory_regions.constSlice()) |region| {
        log.info("Region: {}..{}", .{ region.range.start, region.range.end() });
        if (region.range.start < memory_start) {
            memory_start = region.range.start;
        }
        if (region.range.end() > memory_end) {
            memory_end = region.range.end();
        }
    }

    const memory_pages = memory_end - memory_start; // == bitmap bits required
    var bitmap = alloc_bitmap(memory_pages);
    const refcounters = alloc_refcounters(memory_pages);

    var available_pages: usize = 0;

    for (g_memory_regions.constSlice()) |region| {
        const offset = region.range.start - memory_start;
        for (0..region.range.len) |i| {
            refcounters[offset + i] = 0;
            available_pages += 1;
        }
    }
    for (g_reserved_regions.constSlice()) |region| {
        const offset = region.range.start - memory_start;
        for (0..region.range.len) |i| {
            if (offset + i >= memory_pages) {
                break;
            }
            refcounters[offset + i] = std.math.maxInt(u32);
            bitmap.set_bit(offset + i);

            available_pages -= 1;
        }
    }

    var size_fmt: [64]u8 = undefined;
    const size_fmt_str = mem.format_size(&size_fmt, available_pages * vmm.PAGE_SIZE);
    log.info(
        "Available memory: {s}, bitmap {*}, refcounts {*}",
        .{ size_fmt_str, bitmap.data, refcounters },
    );

    g_physical_memory.len = memory_pages;
    g_physical_memory.memory_start = memory_start * vmm.L3.SIZE;
    g_physical_memory.usage_bitmap = bitmap;
    g_physical_memory.page_refcounters = refcounters;
}

fn trace_allocation(count: usize, page: ?mem.PhysicalAddress) void {
    if (page) |p| {
        log.debug("alloc {} = 0x{x}", .{ count, p.raw });
    } else {
        log.debug("alloc {} = FAIL", .{count});
    }
}

fn trace_free(page: mem.PhysicalAddress) void {
    log.debug("free 0x{x}", .{page.raw});
}

/// Allocates a single 4KiB physical memory page.
pub fn alloc_page() ?mem.PhysicalAddress {
    const guard = g_physical_memory_lock.lock_irqsave();
    defer guard.release();
    const page = g_physical_memory.alloc_page();
    if (comptime kernel.TRACE_PHYSICAL_ALLOCATOR) {
        trace_allocation(1, page);
    }
    return page;
}

/// Allocates a set of `count` contiguous 4KiB pages.
pub fn alloc_pages(count: usize) ?mem.PhysicalAddress {
    const guard = g_physical_memory_lock.lock_irqsave();
    defer guard.release();
    const pages = g_physical_memory.alloc_pages(count);
    if (comptime kernel.TRACE_PHYSICAL_ALLOCATOR) {
        trace_allocation(count, pages);
    }
    return pages;
}

/// Deallocates a single page of physical memory.
///
/// # Panics
///
/// * If the `page` does not represent a valid managed page.
/// * If the `page` was not previously allocated.
/// * If a double free is detected.
pub fn free_page(page: mem.PhysicalAddress) void {
    if (comptime kernel.TRACE_PHYSICAL_ALLOCATOR) {
        trace_free(page);
    }
    const guard = g_physical_memory_lock.lock_irqsave();
    defer guard.release();
    g_physical_memory.free_page(page);
}