Merge pull request #1238 from lightpanda-io/zigdom-chain-allocation

Prototype Chain Allocations

src/browser/Factory.zig

@@ -17,6 +17,7 @@
 // along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
 const std = @import("std");
+const assert = std.debug.assert;
 const builtin = @import("builtin");
 const reflect = @import("reflect.zig");
 const IS_DEBUG = builtin.mode == .Debug;
@@ -35,21 +36,113 @@ const EventTarget = @import("webapi/EventTarget.zig");
 const XMLHttpRequestEventTarget = @import("webapi/net/XMLHttpRequestEventTarget.zig");
 const Blob = @import("webapi/Blob.zig");
 
-const MemoryPoolAligned = std.heap.MemoryPoolAligned;
-
-// 1. Generally, wrapping an ArenaAllocator within an ArenaAllocator doesn't make
-// much sense. But wrapping a MemoryPool within an Arena does. Specifically, by
-// doing so, we solve a major issue with Arena: freed memory can be re-used [for
-// more of the same size].
-// 2. Normally, you have a MemoryPool(T) where T is a `User` or something. Then
-// the MemoryPool can be used for creating users. But in reality, that memory
-// created by that pool could be re-used for anything with the same size (or less)
-// than a User (and a compatible alignment). So that's what we do - we have size
-// (and alignment) based pools.
 const Factory = @This();
 _page: *Page,
 _slab: SlabAllocator,
 
+fn PrototypeChain(comptime types: []const type) type {
+    return struct {
+        const Self = @This();
+        memory: []u8,
+
+        fn totalSize() usize {
+            var size: usize = 0;
+            for (types) |T| {
+                size = std.mem.alignForward(usize, size, @alignOf(T));
+                size += @sizeOf(T);
+            }
+            return size;
+        }
+
+        fn maxAlign() std.mem.Alignment {
+            var alignment: std.mem.Alignment = .@"1";
+
+            for (types) |T| {
+                alignment = std.mem.Alignment.max(alignment, std.mem.Alignment.of(T));
+            }
+
+            return alignment;
+        }
+
+        fn getType(comptime index: usize) type {
+            return types[index];
+        }
+
+        fn allocate(allocator: std.mem.Allocator) !Self {
+            const size = comptime Self.totalSize();
+            const alignment = comptime Self.maxAlign();
+
+            const memory = try allocator.alignedAlloc(u8, alignment, size);
+            return .{ .memory = memory };
+        }
+
+        fn get(self: *const Self, comptime index: usize) *getType(index) {
+            var offset: usize = 0;
+            inline for (types, 0..) |T, i| {
+                offset = std.mem.alignForward(usize, offset, @alignOf(T));
+
+                if (i == index) {
+                    return @as(*T, @ptrCast(@alignCast(self.memory.ptr + offset)));
+                }
+                offset += @sizeOf(T);
+            }
+            unreachable;
+        }
+
+        fn set(self: *const Self, comptime index: usize, value: getType(index)) void {
+            const ptr = self.get(index);
+            ptr.* = value;
+        }
+
+        fn setRoot(self: *const Self, comptime T: type) void {
+            const ptr = self.get(0);
+            ptr.* = .{ ._type = unionInit(T, self.get(1)) };
+        }
+
+        fn setMiddle(self: *const Self, comptime index: usize, comptime T: type) void {
+            assert(index >= 1);
+            assert(index < types.len);
+
+            const ptr = self.get(index);
+            ptr.* = .{ ._proto = self.get(index - 1), ._type = unionInit(T, self.get(index + 1)) };
+        }
+
+        fn setMiddleWithValue(self: *const Self, comptime index: usize, comptime T: type, value: anytype) void {
+            assert(index >= 1);
+
+            const ptr = self.get(index);
+            ptr.* = .{ ._proto = self.get(index - 1), ._type = unionInit(T, value) };
+        }
+
+        fn setLeaf(self: *const Self, comptime index: usize, value: anytype) void {
+            assert(index >= 1);
+
+            const ptr = self.get(index);
+            ptr.* = value;
+            ptr._proto = self.get(index - 1);
+        }
+    };
+}
+
+fn AutoPrototypeChain(comptime types: []const type) type {
+    return struct {
+        fn create(allocator: std.mem.Allocator, leaf_value: anytype) !*@TypeOf(leaf_value) {
+            const chain = try PrototypeChain(types).allocate(allocator);
+
+            const RootType = types[0];
+            chain.setRoot(RootType.Type);
+
+            inline for (1..types.len - 1) |i| {
+                const MiddleType = types[i];
+                chain.setMiddle(i, MiddleType.Type);
+            }
+
+            chain.setLeaf(types.len - 1, leaf_value);
+            return chain.get(types.len - 1);
+        }
+    };
+}
+
 pub fn init(page: *Page) Factory {
     return .{
         ._page = page,
@@ -59,168 +152,163 @@ pub fn init(page: *Page) Factory {
 
 // this is a root object
 pub fn eventTarget(self: *Factory, child: anytype) !*@TypeOf(child) {
-    const child_ptr = try self.createT(@TypeOf(child));
-    child_ptr.* = child;
+    const allocator = self._slab.allocator();
+    const chain = try PrototypeChain(
+        &.{ EventTarget, @TypeOf(child) },
+    ).allocate(allocator);
 
-    const et = try self.createT(EventTarget);
-    child_ptr._proto = et;
-    et.* = .{ ._type = unionInit(EventTarget.Type, child_ptr) };
-    return child_ptr;
-}
+    const event_ptr = chain.get(0);
+    event_ptr.* = .{
+        ._type = unionInit(EventTarget.Type, chain.get(1)),
+    };
+    chain.setLeaf(1, child);
 
-pub fn node(self: *Factory, child: anytype) !*@TypeOf(child) {
-    const child_ptr = try self.createT(@TypeOf(child));
-    child_ptr.* = child;
-    child_ptr._proto = try self.eventTarget(Node{
-        ._proto = undefined,
-        ._type = unionInit(Node.Type, child_ptr),
-    });
-    return child_ptr;
-}
-
-pub fn document(self: *Factory, child: anytype) !*@TypeOf(child) {
-    const child_ptr = try self.createT(@TypeOf(child));
-    child_ptr.* = child;
-    child_ptr._proto = try self.node(Document{
-        ._proto = undefined,
-        ._type = unionInit(Document.Type, child_ptr),
-    });
-    return child_ptr;
-}
-
-pub fn documentFragment(self: *Factory, child: anytype) !*@TypeOf(child) {
-    const child_ptr = try self.createT(@TypeOf(child));
-    child_ptr.* = child;
-    child_ptr._proto = try self.node(Node.DocumentFragment{
-        ._proto = undefined,
-        ._type = unionInit(Node.DocumentFragment.Type, child_ptr),
-    });
-    return child_ptr;
-}
-
-pub fn element(self: *Factory, child: anytype) !*@TypeOf(child) {
-    const child_ptr = try self.createT(@TypeOf(child));
-    child_ptr.* = child;
-    child_ptr._proto = try self.node(Element{
-        ._proto = undefined,
-        ._type = unionInit(Element.Type, child_ptr),
-    });
-    return child_ptr;
-}
-
-pub fn htmlElement(self: *Factory, child: anytype) !*@TypeOf(child) {
-    if (comptime fieldIsPointer(Element.Html.Type, @TypeOf(child))) {
-        const child_ptr = try self.createT(@TypeOf(child));
-        child_ptr.* = child;
-        child_ptr._proto = try self.element(Element.Html{
-            ._proto = undefined,
-            ._type = unionInit(Element.Html.Type, child_ptr),
-        });
-        return child_ptr;
-    }
-
-    // Our union type fields are usually pointers. But, at the leaf, they
-    // can be struct (if all they contain is the `_proto` field, then we might
-    // as well store it directly in the struct).
-
-    const html = try self.element(Element.Html{
-        ._proto = undefined,
-        ._type = unionInit(Element.Html.Type, child),
-    });
-    const field_name = comptime unionFieldName(Element.Html.Type, @TypeOf(child));
-    var child_ptr = &@field(html._type, field_name);
-    child_ptr._proto = html;
-    return child_ptr;
-}
-
-pub fn svgElement(self: *Factory, tag_name: []const u8, child: anytype) !*@TypeOf(child) {
-    if (@TypeOf(child) == Element.Svg) {
-        return self.element(child);
-    }
-
-    // will never allocate, can't fail
-    const tag_name_str = String.init(self._page.arena, tag_name, .{}) catch unreachable;
-
-    if (comptime fieldIsPointer(Element.Svg.Type, @TypeOf(child))) {
-        const child_ptr = try self.createT(@TypeOf(child));
-        child_ptr.* = child;
-        child_ptr._proto = try self.element(Element.Svg{
-            ._proto = undefined,
-            ._tag_name = tag_name_str,
-            ._type = unionInit(Element.Svg.Type, child_ptr),
-        });
-        return child_ptr;
-    }
-
-    // Our union type fields are usually pointers. But, at the leaf, they
-    // can be struct (if all they contain is the `_proto` field, then we might
-    // as well store it directly in the struct).
-    const svg = try self.element(Element.Svg{
-        ._proto = undefined,
-        ._tag_name = tag_name_str,
-        ._type = unionInit(Element.Svg.Type, child),
-    });
-    const field_name = comptime unionFieldName(Element.Svg.Type, @TypeOf(child));
-    var child_ptr = &@field(svg._type, field_name);
-    child_ptr._proto = svg;
-    return child_ptr;
+    return chain.get(1);
 }
 
 // this is a root object
 pub fn event(self: *Factory, typ: []const u8, child: anytype) !*@TypeOf(child) {
-    const child_ptr = try self.createT(@TypeOf(child));
-    child_ptr.* = child;
+    const allocator = self._slab.allocator();
 
-    const e = try self.createT(Event);
-    child_ptr._proto = e;
-    e.* = .{
-        ._type = unionInit(Event.Type, child_ptr),
+    // Special case: Event has a _type_string field, so we need manual setup
+    const chain = try PrototypeChain(
+        &.{ Event, @TypeOf(child) },
+    ).allocate(allocator);
+
+    const event_ptr = chain.get(0);
+    event_ptr.* = .{
+        ._type = unionInit(Event.Type, chain.get(1)),
         ._type_string = try String.init(self._page.arena, typ, .{}),
     };
-    return child_ptr;
-}
+    chain.setLeaf(1, child);
 
-pub fn xhrEventTarget(self: *Factory, child: anytype) !*@TypeOf(child) {
-    const et = try self.eventTarget(XMLHttpRequestEventTarget{
-        ._proto = undefined,
-        ._type = unionInit(XMLHttpRequestEventTarget.Type, child),
-    });
-    const field_name = comptime unionFieldName(XMLHttpRequestEventTarget.Type, @TypeOf(child));
-    var child_ptr = &@field(et._type, field_name);
-    child_ptr._proto = et;
-    return child_ptr;
+    return chain.get(1);
 }
 
 pub fn blob(self: *Factory, child: anytype) !*@TypeOf(child) {
-    const child_ptr = try self.createT(@TypeOf(child));
-    child_ptr.* = child;
+    const allocator = self._slab.allocator();
 
-    const b = try self.createT(Blob);
-    child_ptr._proto = b;
-    b.* = .{
-        ._type = unionInit(Blob.Type, child_ptr),
+    // Special case: Blob has slice and mime fields, so we need manual setup
+    const chain = try PrototypeChain(
+        &.{ Blob, @TypeOf(child) },
+    ).allocate(allocator);
+
+    const blob_ptr = chain.get(0);
+    blob_ptr.* = .{
+        ._type = unionInit(Blob.Type, chain.get(1)),
         .slice = "",
         .mime = "",
     };
-    return child_ptr;
+    chain.setLeaf(1, child);
+
+    return chain.get(1);
 }
 
-pub fn create(self: *Factory, value: anytype) !*@TypeOf(value) {
-    const ptr = try self.createT(@TypeOf(value));
-    ptr.* = value;
-    return ptr;
-}
-
-pub fn createT(self: *Factory, comptime T: type) !*T {
+pub fn node(self: *Factory, child: anytype) !*@TypeOf(child) {
     const allocator = self._slab.allocator();
-    return try allocator.create(T);
+    return try AutoPrototypeChain(
+        &.{ EventTarget, Node, @TypeOf(child) },
+    ).create(allocator, child);
+}
+
+pub fn document(self: *Factory, child: anytype) !*@TypeOf(child) {
+    const allocator = self._slab.allocator();
+    return try AutoPrototypeChain(
+        &.{ EventTarget, Node, Document, @TypeOf(child) },
+    ).create(allocator, child);
+}
+
+pub fn documentFragment(self: *Factory, child: anytype) !*@TypeOf(child) {
+    const allocator = self._slab.allocator();
+    return try AutoPrototypeChain(
+        &.{ EventTarget, Node, Node.DocumentFragment, @TypeOf(child) },
+    ).create(allocator, child);
+}
+
+pub fn element(self: *Factory, child: anytype) !*@TypeOf(child) {
+    const allocator = self._slab.allocator();
+    return try AutoPrototypeChain(
+        &.{ EventTarget, Node, Element, @TypeOf(child) },
+    ).create(allocator, child);
+}
+
+pub fn htmlElement(self: *Factory, child: anytype) !*@TypeOf(child) {
+    const allocator = self._slab.allocator();
+    return try AutoPrototypeChain(
+        &.{ EventTarget, Node, Element, Element.Html, @TypeOf(child) },
+    ).create(allocator, child);
+}
+
+pub fn svgElement(self: *Factory, tag_name: []const u8, child: anytype) !*@TypeOf(child) {
+    const allocator = self._slab.allocator();
+    const ChildT = @TypeOf(child);
+
+    if (ChildT == Element.Svg) {
+        return self.element(child);
+    }
+
+    const chain = try PrototypeChain(
+        &.{ EventTarget, Node, Element, Element.Svg, ChildT },
+    ).allocate(allocator);
+
+    chain.setRoot(EventTarget.Type);
+    chain.setMiddle(1, Node.Type);
+    chain.setMiddle(2, Element.Type);
+
+    // will never allocate, can't fail
+    const tag_name_str = String.init(self._page.arena, tag_name, .{}) catch unreachable;
+
+    // Manually set Element.Svg with the tag_name
+    chain.set(3, .{
+        ._proto = chain.get(2),
+        ._tag_name = tag_name_str,
+        ._type = unionInit(Element.Svg.Type, chain.get(4)),
+    });
+
+    chain.setLeaf(4, child);
+    return chain.get(4);
+}
+
+pub fn xhrEventTarget(self: *Factory, child: anytype) !*@TypeOf(child) {
+    const allocator = self._slab.allocator();
+
+    return try AutoPrototypeChain(
+        &.{ EventTarget, XMLHttpRequestEventTarget, @TypeOf(child) },
+    ).create(allocator, child);
+}
+
+fn hasChainRoot(comptime T: type) bool {
+    // Check if this is a root
+    if (@hasDecl(T, "_prototype_root")) {
+        return true;
+    }
+
+    // If no _proto field, we're at the top but not a recognized root
+    if (!@hasField(T, "_proto")) return false;
+
+    // Get the _proto field's type and recurse
+    const fields = @typeInfo(T).@"struct".fields;
+    inline for (fields) |field| {
+        if (std.mem.eql(u8, field.name, "_proto")) {
+            const ProtoType = reflect.Struct(field.type);
+            return hasChainRoot(ProtoType);
+        }
+    }
+
+    return false;
+}
+
+fn isChainType(comptime T: type) bool {
+    if (@hasField(T, "_proto")) return false;
+    return comptime hasChainRoot(T);
 }
 
 pub fn destroy(self: *Factory, value: anytype) void {
     const S = reflect.Struct(@TypeOf(value));
+
     if (comptime IS_DEBUG) {
         // We should always destroy from the leaf down.
-        if (@hasField(S, "_type") and @typeInfo(@TypeOf(value._type)) == .@"union") {
+        if (@hasDecl(S, "_prototype_root")) {
             // A Event{._type == .generic} (or any other similar types)
             // _should_ be destoyed directly. The _type = .generic is a pseudo
             // child
@@ -231,14 +319,48 @@ pub fn destroy(self: *Factory, value: anytype) void {
         }
     }
 
-    self.destroyChain(value, true);
+    if (comptime isChainType(S)) {
+        self.destroyChain(value, true, 0, std.mem.Alignment.@"1");
+    } else {
+        self.destroyStandalone(value);
+    }
 }
 
-fn destroyChain(self: *Factory, value: anytype, comptime first: bool) void {
+pub fn destroyStandalone(self: *Factory, value: anytype) void {
     const S = reflect.Struct(@TypeOf(value));
+    assert(!@hasDecl(S, "_prototype_root"));
 
     const allocator = self._slab.allocator();
 
+    if (@hasDecl(S, "deinit")) {
+        // And it has a deinit, we'll call it
+        switch (@typeInfo(@TypeOf(S.deinit)).@"fn".params.len) {
+            1 => value.deinit(),
+            2 => value.deinit(self._page),
+            else => @compileLog(@typeName(S) ++ " has an invalid deinit function"),
+        }
+    }
+
+    allocator.destroy(value);
+}
+
+fn destroyChain(
+    self: *Factory,
+    value: anytype,
+    comptime first: bool,
+    old_size: usize,
+    old_align: std.mem.Alignment,
+) void {
+    const S = reflect.Struct(@TypeOf(value));
+    const allocator = self._slab.allocator();
+
+    // aligns the old size to the alignment of this element
+    const current_size = std.mem.alignForward(usize, old_size, @alignOf(S));
+    const alignment = std.mem.Alignment.fromByteUnits(@alignOf(S));
+
+    const new_align = std.mem.Alignment.max(old_align, alignment);
+    const new_size = current_size + @sizeOf(S);
+
     // This is initially called from a deinit. We don't want to call that
     // same deinit. So when this is the first time destroyChain is called
     // we don't call deinit (because we're in that deinit)
@@ -255,44 +377,33 @@ fn destroyChain(self: *Factory, value: anytype, comptime first: bool) void {
     }
 
     if (@hasField(S, "_proto")) {
-        self.destroyChain(value._proto, false);
+        self.destroyChain(value._proto, false, new_size, new_align);
     } else if (@hasDecl(S, "JsApi")) {
         // Doesn't have a _proto, but has a JsApi.
         if (self._page.js.removeTaggedMapping(@intFromPtr(value))) |tagged| {
             allocator.destroy(tagged);
         }
-    }
+    } else {
+        // no proto so this is the head of the chain.
+        // we use this as the ptr to the start of the chain.
+        // and we have summed up the length.
+        assert(@hasDecl(S, "_prototype_root"));
 
-    // Leaf types are allowed by be placed directly within their _proto
-    // (which makes sense when the @sizeOf(Leaf) == 8). These don't need to
-    // be (cannot be) freed. But we'll still free the chain.
-    if (comptime wasAllocated(S)) {
-        allocator.destroy(value);
+        const memory_ptr: [*]const u8 = @ptrCast(value);
+        const len = std.mem.alignForward(usize, new_size, new_align.toByteUnits());
+        allocator.free(memory_ptr[0..len]);
     }
 }
 
-fn wasAllocated(comptime S: type) bool {
-    // Whether it's heap allocate or not, we should have a pointer.
-    // (If it isn't heap allocated, it'll be a pointer from the proto's type
-    // e.g. &html._type.title)
-    if (!@hasField(S, "_proto")) {
-        // a root is always on the heap.
-        return true;
-    }
+pub fn createT(self: *Factory, comptime T: type) !*T {
+    const allocator = self._slab.allocator();
+    return try allocator.create(T);
+}
 
-    // the _proto type
-    const P = reflect.Struct(std.meta.fieldInfo(S, ._proto).type);
-
-    // the _proto._type type (the parent's _type union)
-    const U = std.meta.fieldInfo(P, ._type).type;
-    inline for (@typeInfo(U).@"union".fields) |field| {
-        if (field.type == S) {
-            // One of the types in the proto's _type union is this non-pointer
-            // structure, so it isn't heap allocted.
-            return false;
-        }
-    }
-    return true;
+pub fn create(self: *Factory, value: anytype) !*@TypeOf(value) {
+    const ptr = try self.createT(@TypeOf(value));
+    ptr.* = value;
+    return ptr;
 }
 
 fn unionInit(comptime T: type, value: anytype) T {
@@ -316,15 +427,3 @@ fn unionFieldName(comptime T: type, comptime V: type) []const u8 {
     }
     @compileError(@typeName(V) ++ " is not a valid type for " ++ @typeName(T) ++ ".type");
 }
-
-fn fieldIsPointer(comptime T: type, comptime V: type) bool {
-    inline for (@typeInfo(T).@"union".fields) |field| {
-        if (field.type == V) {
-            return false;
-        }
-        if (field.type == *V) {
-            return true;
-        }
-    }
-    @compileError(@typeName(V) ++ " is not a valid type for " ++ @typeName(T) ++ ".type");
-}

src/browser/Page.zig

@@ -177,7 +177,9 @@ pub fn deinit(self: *Page) void {
 
         // Uncomment if you want slab statistics to print.
         // const stats = self._factory._slab.getStats(self.arena) catch unreachable;
-        // stats.print() catch unreachable;
+        // var buffer: [256]u8 = undefined;
+        // var stream = std.fs.File.stderr().writer(&buffer).interface;
+        // stats.print(&stream) catch unreachable;
     }
 
     self.js.deinit();

src/browser/webapi/Blob.zig

@@ -26,7 +26,9 @@ const Page = @import("../Page.zig");
 /// https://developer.mozilla.org/en-US/docs/Web/API/Blob
 const Blob = @This();
 
+const _prototype_root = true;
 _type: Type,
+
 /// Immutable slice of blob.
 /// Note that another blob may hold a pointer/slice to this,
 /// so its better to leave the deallocation of it to arena allocator.

src/browser/webapi/Event.zig

@@ -26,7 +26,9 @@ const String = @import("../../string.zig").String;
 
 pub const Event = @This();
 
+const _prototype_root = true;
 _type: Type,
+
 _bubbles: bool = false,
 _cancelable: bool = false,
 _composed: bool = false,

src/browser/webapi/EventTarget.zig

@@ -26,6 +26,7 @@ const Event = @import("Event.zig");
 
 const EventTarget = @This();
 
+const _prototype_root = true;
 _type: Type,
 
 pub const Type = union(enum) {

src/browser/webapi/element/Html.zig

@@ -67,36 +67,36 @@ pub fn construct(page: *Page) !*Element {
 }
 
 pub const Type = union(enum) {
-    anchor: Anchor,
-    body: Body,
-    br: BR,
-    button: Button,
+    anchor: *Anchor,
+    body: *Body,
+    br: *BR,
+    button: *Button,
     custom: *Custom,
-    dialog: Dialog,
-    div: Div,
-    form: Form,
+    dialog: *Dialog,
+    div: *Div,
+    form: *Form,
     generic: *Generic,
     heading: *Heading,
-    head: Head,
-    html: Html,
-    hr: HR,
-    img: Image,
-    iframe: IFrame,
+    head: *Head,
+    html: *Html,
+    hr: *HR,
+    img: *Image,
+    iframe: *IFrame,
     input: *Input,
-    li: LI,
-    link: Link,
-    meta: Meta,
-    ol: OL,
+    li: *LI,
+    link: *Link,
+    meta: *Meta,
+    ol: *OL,
     option: *Option,
-    p: Paragraph,
+    p: *Paragraph,
     script: *Script,
-    select: Select,
-    slot: Slot,
-    style: Style,
+    select: *Select,
+    slot: *Slot,
+    style: *Style,
     template: *Template,
     text_area: *TextArea,
-    title: Title,
-    ul: UL,
+    title: *Title,
+    ul: *UL,
     unknown: *Unknown,
 };
 

src/browser/webapi/net/XMLHttpRequestEventTarget.zig

@@ -35,7 +35,7 @@ _on_progress: ?js.Function = null,
 _on_timeout: ?js.Function = null,
 
 pub const Type = union(enum) {
-    request: @import("XMLHttpRequest.zig"),
+    request: *@import("XMLHttpRequest.zig"),
     // TODO: xml_http_request_upload
 };
 

src/slab.zig

@@ -258,47 +258,47 @@ pub const SlabAllocator = struct {
         utilization_ratio: f64,
         slabs: []const Slab.Stats,
 
-        pub fn print(self: *const Stats) !void {
-            std.debug.print("\n", .{});
-            std.debug.print("\n=== Slab Allocator Statistics ===\n", .{});
-            std.debug.print("Overall Memory:\n", .{});
-            std.debug.print("  Total allocated: {} bytes ({d:.2} MB)\n", .{
+        pub fn print(self: *const Stats, stream: *std.io.Writer) !void {
+            try stream.print("\n", .{});
+            try stream.print("\n=== Slab Allocator Statistics ===\n", .{});
+            try stream.print("Overall Memory:\n", .{});
+            try stream.print("  Total allocated: {} bytes ({d:.2} MB)\n", .{
                 self.total_allocated_bytes,
                 @as(f64, @floatFromInt(self.total_allocated_bytes)) / 1_048_576.0,
             });
-            std.debug.print("  In use:          {} bytes ({d:.2} MB)\n", .{
+            try stream.print("  In use:          {} bytes ({d:.2} MB)\n", .{
                 self.bytes_in_use,
                 @as(f64, @floatFromInt(self.bytes_in_use)) / 1_048_576.0,
             });
-            std.debug.print("  Free:            {} bytes ({d:.2} MB)\n", .{
+            try stream.print("  Free:            {} bytes ({d:.2} MB)\n", .{
                 self.bytes_free,
                 @as(f64, @floatFromInt(self.bytes_free)) / 1_048_576.0,
             });
 
-            std.debug.print("\nOverall Structure:\n", .{});
-            std.debug.print("  Slab Count:    {}\n", .{self.slab_count});
-            std.debug.print("  Total chunks:    {}\n", .{self.total_chunks});
-            std.debug.print("  Total slots:     {}\n", .{self.total_slots});
-            std.debug.print("  Slots in use:    {}\n", .{self.slots_in_use});
-            std.debug.print("  Slots free:      {}\n", .{self.slots_free});
+            try stream.print("\nOverall Structure:\n", .{});
+            try stream.print("  Slab Count:    {}\n", .{self.slab_count});
+            try stream.print("  Total chunks:    {}\n", .{self.total_chunks});
+            try stream.print("  Total slots:     {}\n", .{self.total_slots});
+            try stream.print("  Slots in use:    {}\n", .{self.slots_in_use});
+            try stream.print("  Slots free:      {}\n", .{self.slots_free});
 
-            std.debug.print("\nOverall Efficiency:\n", .{});
-            std.debug.print("  Utilization:     {d:.1}%\n", .{self.utilization_ratio * 100.0});
-            std.debug.print("  Fragmentation:   {d:.1}%\n", .{self.fragmentation_ratio * 100.0});
+            try stream.print("\nOverall Efficiency:\n", .{});
+            try stream.print("  Utilization:     {d:.1}%\n", .{self.utilization_ratio * 100.0});
+            try stream.print("  Fragmentation:   {d:.1}%\n", .{self.fragmentation_ratio * 100.0});
 
             if (self.slabs.len > 0) {
-                std.debug.print("\nPer-Slab Breakdown:\n", .{});
-                std.debug.print(
+                try stream.print("\nPer-Slab Breakdown:\n", .{});
+                try stream.print(
                     "  {s:>5} | {s:>4} | {s:>6} | {s:>6} | {s:>6} | {s:>10} | {s:>6}\n",
                     .{ "Size", "Algn", "Chunks", "Slots", "InUse", "Bytes", "Util%" },
                 );
-                std.debug.print(
+                try stream.print(
                     "  {s:-<5}-+-{s:-<4}-+-{s:-<6}-+-{s:-<6}-+-{s:-<6}-+-{s:-<10}-+-{s:-<6}\n",
                     .{ "", "", "", "", "", "", "" },
                 );
 
                 for (self.slabs) |slab| {
-                    std.debug.print("  {d:5} | {d:4} | {d:6} | {d:6} | {d:6} | {d:10} | {d:5.1}%\n", .{
+                    try stream.print("  {d:5} | {d:4} | {d:6} | {d:6} | {d:6} | {d:10} | {d:5.1}%\n", .{
                         slab.key.size,
                         @intFromEnum(slab.key.alignment),
                         slab.chunk_count,
@@ -376,23 +376,25 @@ pub const SlabAllocator = struct {
         const self: *Self = @ptrCast(@alignCast(ctx));
         _ = ret_addr;
 
+        const aligned_len = std.mem.alignForward(usize, len, alignment.toByteUnits());
+
         const list_gop = self.slabs.getOrPut(
             self.child_allocator,
-            SlabKey{ .size = len, .alignment = alignment },
+            SlabKey{ .size = aligned_len, .alignment = alignment },
         ) catch return null;
 
         if (!list_gop.found_existing) {
             list_gop.value_ptr.* = Slab.init(
                 self.child_allocator,
                 alignment,
-                len,
+                aligned_len,
                 self.max_slot_count,
             ) catch return null;
         }
 
         const list = list_gop.value_ptr;
         const buf = list.alloc(self.child_allocator) catch return null;
-        return buf.ptr;
+        return buf[0..len].ptr;
     }
 
     fn free(ctx: *anyopaque, memory: []u8, alignment: Alignment, ret_addr: usize) void {
@@ -401,8 +403,9 @@ pub const SlabAllocator = struct {
 
         const ptr = memory.ptr;
         const len = memory.len;
+        const aligned_len = std.mem.alignForward(usize, len, alignment.toByteUnits());
 
-        const list = self.slabs.getPtr(.{ .size = len, .alignment = alignment }).?;
+        const list = self.slabs.getPtr(.{ .size = aligned_len, .alignment = alignment }).?;
         list.free(ptr);
     }
 };
@@ -822,3 +825,39 @@ test "slab allocator - different size classes don't interfere" {
     allocator.free(ptr_128);
     allocator.free(ptr_64_again);
 }
+
+test "slab allocator - 16-byte alignment" {
+    var slab_alloc = TestSlabAllocator.init(testing.allocator, 16);
+    defer slab_alloc.deinit();
+
+    const allocator = slab_alloc.allocator();
+
+    // Request 16-byte aligned memory
+    const ptr = try allocator.alignedAlloc(u8, .@"16", 152);
+    defer allocator.free(ptr);
+
+    // Verify alignment
+    const addr = @intFromPtr(ptr.ptr);
+    try testing.expect(addr % 16 == 0);
+
+    // Make sure we can use it
+    @memset(ptr, 0xFF);
+}
+
+test "slab allocator - various alignments" {
+    var slab_alloc = TestSlabAllocator.init(testing.allocator, 16);
+    defer slab_alloc.deinit();
+
+    const allocator = slab_alloc.allocator();
+
+    const alignments = [_]std.mem.Alignment{ .@"1", .@"2", .@"4", .@"8", .@"16" };
+
+    inline for (alignments) |alignment| {
+        const ptr = try allocator.alignedAlloc(u8, alignment, 100);
+        defer allocator.free(ptr);
+
+        const addr = @intFromPtr(ptr.ptr);
+        const align_value = alignment.toByteUnits();
+        try testing.expect(addr % align_value == 0);
+    }
+}