mm: implement a basic virtual memory manager

src/mem/vmalloc.zig

@@ -0,0 +1,426 @@
+const std = @import("std");
+
+const Range = @import("../util/range.zig").Range;
+
+const Allocator = std.mem.Allocator;
+
+/// Describes a single virtual memory range.
+///
+/// Used by `VirtualMemoryAllocator` to track allocated/used regions.
+pub const VirtualMemoryRange = struct {
+    range: Range(u64),
+
+    prev: ?*VirtualMemoryRange = null,
+    next: ?*VirtualMemoryRange = null,
+};
+
+/// Virtual memory allocator implementation.
+pub const VirtualMemoryAllocator = struct {
+    gpa: Allocator,
+    head: ?*VirtualMemoryRange = null,
+    outer_range: Range(u64),
+
+    /// One of errors returned by the allocation logic + underlying allocator error.
+    pub const Error = error{ already_exists, invalid_region, cannot_fit } || Allocator.Error;
+
+    /// An iterator over VM regions being freed.
+    pub const FreeIterator = struct {
+        range: Range(u64),
+        vma: *VirtualMemoryAllocator,
+        current: ?*VirtualMemoryRange,
+
+        fn next(self: *@This()) Error!?Range(u64) {
+            while (self.current) |n| {
+                if (n.range.intersect(&self.range)) |xs| {
+                    if (xs.start == n.range.start) {
+                        if (xs.end() == n.range.end()) {
+                            // Whole range encompassed by requested range
+                            // Unlink the node
+                            if (n.next) |nn| {
+                                nn.prev = n.prev;
+                            }
+                            if (n.prev) |np| {
+                                np.next = n.next;
+                            } else {
+                                self.vma.head = n.next;
+                            }
+                            // Free it
+                            self.current = n.next;
+                            self.vma.gpa.destroy(n);
+
+                            return xs;
+                        }
+
+                        // Remove space from the start
+                        n.range.start += xs.len;
+                        n.range.len -= xs.len;
+                        // Does not touch the end, so can be sure this is the last node
+                        self.current = null;
+                        return xs;
+                    } else if (xs.end() == n.range.end()) {
+                        n.range.len -= xs.len;
+                        // Continue, there might be a following node affected
+                        self.current = n.next;
+                        return xs;
+                    } else {
+                        // Insert a new node after the current one
+                        const new_node = try self.vma.gpa.create(VirtualMemoryRange);
+                        new_node.* = VirtualMemoryRange {
+                            .range = .{ .start = xs.end(), .len = n.range.end() - xs.end() },
+                            .prev = n,
+                            .next = n.next,
+                        };
+                        n.range.len = xs.start - n.range.start;
+                        if (n.next) |nn| {
+                            nn.prev = new_node;
+                        }
+                        n.next = new_node;
+                        // Requested region is fully encompassed by this one, so no intersections
+                        // will follow
+                        self.current = null;
+                        return xs;
+                    }
+                } else {
+                    // No intersect
+                    self.current = n.next;
+                }
+            }
+            return null;
+        }
+    };
+
+    /// Creates a new instance of a virtual memory allocator.
+    pub fn init(gpa: Allocator, outer_range: Range(u64)) @This() {
+        return .{
+            .outer_range = outer_range,
+            .gpa = gpa,
+        };
+    }
+
+    /// Allocates a free region of virtual memory of requested (`pfn_count`) size.
+    ///
+    /// # Errors
+    ///
+    /// * `cannot_fit` - if no free space found to fit the requested allocation.
+    /// * Underlying allocator error - if allocation of a new node fails.
+    pub fn allocate(self: *@This(), pfn_count: u64) Error!u64 {
+        // Try to fit before first entry
+        const gap_before_first = if (self.head) |n| (n.range.start - self.outer_range.start) else self.outer_range.len;
+
+        if (gap_before_first >= pfn_count) {
+            var new_node = try self.gpa.create(VirtualMemoryRange);
+
+            new_node.range = .{ .start = self.outer_range.start, .len = pfn_count };
+            new_node.next = self.head;
+            new_node.prev = null;
+
+            if (self.head) |n| {
+                n.prev = new_node;
+            }
+
+            self.head = new_node;
+
+            return self.outer_range.start;
+        }
+
+        // If cannot fit before first entry, find an entry to fit after
+        var node = self.head;
+        while (node) |n| {
+            const gap =
+                if (n.next) |nn|
+                    // Gap between this and next
+                    (nn.range.start - n.range.end())
+                else
+                    // Gap between this and the end
+                    (self.outer_range.end() - n.range.end());
+
+            if (gap >= pfn_count) {
+                // Insert after this
+                const result = n.range.end();
+                var new_node = try self.gpa.create(VirtualMemoryRange);
+                new_node.prev = n;
+                new_node.next = n.next;
+                new_node.range = .{ .start = result, .len = pfn_count };
+                if (n.next) |nn| {
+                    nn.prev = new_node;
+                }
+                n.next = new_node;
+                return result;
+            }
+
+            node = n.next;
+        }
+
+        return error.cannot_fit;
+    }
+
+    /// Inserts a reservation into the VM allocator.
+    ///
+    /// # Errors
+    ///
+    /// * `already_exists` - if the requested range intersects existing ranges.
+    /// * Underlying allocator error - if allocation of a new node fails.
+    pub fn insert(self: *@This(), region: Range(u64)) Error!void {
+        // Validate that the range does not escape the outer range
+        if (region.start < self.outer_range.start or region.end() > self.outer_range.end()) {
+            return error.invalid_region;
+        }
+
+        // Find the last node which is before the region supposed to be inserted
+        var node = self.head;
+        var insert_after: ?*VirtualMemoryRange = null;
+        while (node) |n| {
+            if (n.range.intersect(&region) != null) {
+                return error.already_exists;
+            }
+
+            if (n.range.end() <= region.start) {
+                insert_after = n;
+            }
+
+            node = n.next;
+        }
+
+        var new_node = try self.gpa.create(VirtualMemoryRange);
+
+        new_node.range = region;
+
+        if (insert_after) |ia| {
+            new_node.prev = ia;
+            new_node.next = ia.next;
+
+            if (ia.next) |ian| {
+                ian.prev = new_node;
+            }
+            ia.next = new_node;
+        } else {
+            new_node.next = null;
+            new_node.prev = null;
+
+            self.head = new_node;
+        }
+    }
+
+    /// Deallocates (shrinks/truncates) regions intersecting the requested range.
+    pub fn free(self: *@This(), start_pfn: u64, pfn_count: u64) FreeIterator {
+        const range = Range(u64) { .start = start_pfn, .len = pfn_count };
+        return FreeIterator {
+            .current = self.head,
+            .vma = self,
+            .range = range,
+        };
+    }
+};
+
+test "Inserted entries in vmalloc are properly ordered" {
+    var vma = VirtualMemoryAllocator.init(std.testing.allocator, .{ .start = 0x1000, .len = 0x2000 });
+    defer {
+        while (vma.head) |n| {
+            vma.head = n.next;
+            std.testing.allocator.destroy(n);
+        }
+    }
+    try vma.insert(.{ .start = 0x1200, .len = 0x200 });
+    {
+        const n0 = vma.head.?;
+        try std.testing.expectEqual(0x1200, n0.range.start);
+        try std.testing.expectEqual(0x200, n0.range.len);
+        try std.testing.expectEqual(null, n0.next);
+        try std.testing.expectEqual(null, n0.prev);
+    }
+    try vma.insert(.{ .start = 0x2000, .len = 0x200 });
+    {
+        const n0 = vma.head.?;
+        try std.testing.expectEqual(0x1200, n0.range.start);
+        try std.testing.expectEqual(0x200, n0.range.len);
+        try std.testing.expectEqual(null, n0.prev);
+        const n1 = n0.next.?;
+        try std.testing.expectEqual(0x2000, n1.range.start);
+        try std.testing.expectEqual(0x200, n1.range.len);
+        try std.testing.expectEqual(n0, n1.prev);
+        try std.testing.expectEqual(null, n1.next);
+    }
+    try vma.insert(.{ .start = 0x1400, .len = 0x200 });
+    {
+        const n0 = vma.head.?;
+        try std.testing.expectEqual(0x1200, n0.range.start);
+        try std.testing.expectEqual(0x200, n0.range.len);
+        try std.testing.expectEqual(null, n0.prev);
+        const n1 = n0.next.?;
+        try std.testing.expectEqual(0x1400, n1.range.start);
+        try std.testing.expectEqual(0x200, n1.range.len);
+        try std.testing.expectEqual(n0, n1.prev);
+        const n2 = n1.next.?;
+        try std.testing.expectEqual(0x2000, n2.range.start);
+        try std.testing.expectEqual(0x200, n2.range.len);
+        try std.testing.expectEqual(n1, n2.prev);
+        try std.testing.expectEqual(null, n2.next);
+    }
+}
+
+test "Overlapping insertions are denied" {
+    var vma = VirtualMemoryAllocator.init(std.testing.allocator, .{ .start = 0x1000, .len = 0x1000 });
+    defer {
+        while (vma.head) |n| {
+            vma.head = n.next;
+            std.testing.allocator.destroy(n);
+        }
+    }
+    try vma.insert(.{ .start = 0x1200, .len = 0x200 });
+    try std.testing.expectError(error.already_exists, vma.insert(.{ .start = 0x1100, .len = 0x200 }));
+    try std.testing.expectError(error.already_exists, vma.insert(.{ .start = 0x1300, .len = 0x200 }));
+    try std.testing.expectError(error.already_exists, vma.insert(.{ .start = 0x1100, .len = 0x400 }));
+}
+
+test "Insertions outside of bounds are denied" {
+    var vma = VirtualMemoryAllocator.init(std.testing.allocator, .{ .start = 0x1000, .len = 0x1000 });
+    // As above...
+    try std.testing.expectError(error.invalid_region, vma.insert(.{ .start = 0x2200, .len = 0x200 }));
+    // ... so below
+    try std.testing.expectError(error.invalid_region, vma.insert(.{ .start = 0x200, .len = 0x200 }));
+    // Crosses from below
+    try std.testing.expectError(error.invalid_region, vma.insert(.{ .start = 0x200, .len = 0x1000 }));
+    // Crosses into above
+    try std.testing.expectError(error.invalid_region, vma.insert(.{ .start = 0x1200, .len = 0x1000 }));
+    // Encompasses whole
+    try std.testing.expectError(error.invalid_region, vma.insert(.{ .start = 0x200, .len = 0x2000 }));
+}
+
+test "Allocations from vmalloc" {
+    var vma = VirtualMemoryAllocator.init(std.testing.allocator, .{ .start = 0x1000, .len = 0x1000 });
+    defer {
+        while (vma.head) |n| {
+            vma.head = n.next;
+            std.testing.allocator.destroy(n);
+        }
+    }
+    try vma.insert(.{ .start = 0x1200, .len = 0x200 });
+    try std.testing.expectEqual(0x1000, try vma.allocate(0x100));
+    try std.testing.expectEqual(0x1400, try vma.allocate(0x400));
+    try std.testing.expectEqual(0x1100, try vma.allocate(0x100));
+}
+
+test "vmalloc free" {
+    var vma = VirtualMemoryAllocator.init(std.testing.allocator, .{ .start = 0x1000, .len = 0x1000 });
+
+    try vma.insert(.{ .start = 0x1200, .len = 0x800 });
+    try vma.insert(.{ .start = 0x1A00, .len = 0x400 });
+
+    // Remove nothing
+    {
+        var free_it = vma.free(0x1000, 0x200);
+        try std.testing.expectEqual(null, free_it.next());
+    }
+
+    // Remove a chunk in the middle of a node
+    {
+        var free_it = vma.free(0x1400, 0x400);
+        const r0 = (try free_it.next()).?;
+        try std.testing.expectEqual(0x1400, r0.start);
+        try std.testing.expectEqual(0x400, r0.len);
+        try std.testing.expectEqual(null, free_it.next());
+
+        const n0 = vma.head.?;
+        try std.testing.expectEqual(0x1200, n0.range.start);
+        try std.testing.expectEqual(0x200, n0.range.len);
+        const n1 = n0.next.?;
+        try std.testing.expectEqual(0x1800, n1.range.start);
+        try std.testing.expectEqual(0x200, n1.range.len);
+        try std.testing.expectEqual(n0, n1.prev);
+        const n2 = n1.next.?;
+        try std.testing.expectEqual(0x1A00, n2.range.start);
+        try std.testing.expectEqual(0x400, n2.range.len);
+        try std.testing.expectEqual(n1, n2.prev);
+        try std.testing.expectEqual(null, n2.next);
+    }
+
+    // Remove from the start
+    {
+        var free_it = vma.free(0x1200, 0x100);
+        const r0 = (try free_it.next()).?;
+        try std.testing.expectEqual(0x1200, r0.start);
+        try std.testing.expectEqual(0x100, r0.len);
+        try std.testing.expectEqual(null, free_it.next());
+
+        const n0 = vma.head.?;
+        try std.testing.expectEqual(0x1300, n0.range.start);
+        try std.testing.expectEqual(0x100, n0.range.len);
+        const n1 = n0.next.?;
+        try std.testing.expectEqual(0x1800, n1.range.start);
+        try std.testing.expectEqual(0x200, n1.range.len);
+        try std.testing.expectEqual(n0, n1.prev);
+        const n2 = n1.next.?;
+        try std.testing.expectEqual(0x1A00, n2.range.start);
+        try std.testing.expectEqual(0x400, n2.range.len);
+        try std.testing.expectEqual(n1, n2.prev);
+        try std.testing.expectEqual(null, n2.next);
+    }
+
+    // Remove from the end
+    {
+        var free_it = vma.free(0x1900, 0x100);
+        const r0 = (try free_it.next()).?;
+        try std.testing.expectEqual(0x1900, r0.start);
+        try std.testing.expectEqual(0x100, r0.len);
+        try std.testing.expectEqual(null, free_it.next());
+
+        const n0 = vma.head.?;
+        try std.testing.expectEqual(0x1300, n0.range.start);
+        try std.testing.expectEqual(0x100, n0.range.len);
+        const n1 = n0.next.?;
+        try std.testing.expectEqual(0x1800, n1.range.start);
+        try std.testing.expectEqual(0x100, n1.range.len);
+        try std.testing.expectEqual(n0, n1.prev);
+        const n2 = n1.next.?;
+        try std.testing.expectEqual(0x1A00, n2.range.start);
+        try std.testing.expectEqual(0x400, n2.range.len);
+        try std.testing.expectEqual(n1, n2.prev);
+        try std.testing.expectEqual(null, n2.next);
+    }
+
+    // Remove single full
+    {
+        var free_it = vma.free(0x1000, 0x600);
+        const r0 = (try free_it.next()).?;
+        try std.testing.expectEqual(0x1300, r0.start);
+        try std.testing.expectEqual(0x100, r0.len);
+        try std.testing.expectEqual(null, free_it.next());
+
+        const n0 = vma.head.?;
+        try std.testing.expectEqual(0x1800, n0.range.start);
+        try std.testing.expectEqual(0x100, n0.range.len);
+        const n1 = n0.next.?;
+        try std.testing.expectEqual(0x1A00, n1.range.start);
+        try std.testing.expectEqual(0x400, n1.range.len);
+        try std.testing.expectEqual(n0, n1.prev);
+        try std.testing.expectEqual(null, n1.next);
+    }
+
+    // Remove one full + one partial
+    {
+        var free_it = vma.free(0x1600, 0x600);
+        const r0 = (try free_it.next()).?;
+        try std.testing.expectEqual(0x1800, r0.start);
+        try std.testing.expectEqual(0x100, r0.len);
+        const r1 = (try free_it.next()).?;
+        try std.testing.expectEqual(0x1A00, r1.start);
+        try std.testing.expectEqual(0x200, r1.len);
+        try std.testing.expectEqual(null, free_it.next());
+
+        const n0 = vma.head.?;
+        try std.testing.expectEqual(0x1C00, n0.range.start);
+        try std.testing.expectEqual(0x200, n0.range.len);
+        try std.testing.expectEqual(null, n0.next);
+    }
+
+    // Remove whatever remains
+    {
+        var free_it = vma.free(0, 0x20000);
+        const r0 = (try free_it.next()).?;
+        try std.testing.expectEqual(0x1C00, r0.start);
+        try std.testing.expectEqual(0x200, r0.len);
+        try std.testing.expectEqual(null, free_it.next());
+
+        try std.testing.expectEqual(null, vma.head);
+    }
+}