From 225f44f3406e6df15ca260bc4bf5c980616871b1 Mon Sep 17 00:00:00 2001
From: Chris Heyes <22148308+hazeycode@users.noreply.github.com>
Date: Thu, 8 Feb 2024 23:53:03 +0000
Subject: [PATCH] zemscripten: Allocator wrapping emmalloc

Co-authored-by: Deins <deinsegle@gmail.com>
---
 libs/zemscripten/build.zig    |  8 ++++-
 libs/zemscripten/src/root.zig | 62 +++++++++++++++++++++++++++++++++++
 2 files changed, 69 insertions(+), 1 deletion(-)

diff --git a/libs/zemscripten/build.zig b/libs/zemscripten/build.zig
index 8598db905..f65c6affd 100644
--- a/libs/zemscripten/build.zig
+++ b/libs/zemscripten/build.zig
@@ -23,7 +23,13 @@ pub const Package = struct {
 
         const emcc = b.addSystemCommand(&.{
             b.pathJoin(&.{ pkg.emsdk_path, "upstream/emscripten/emcc" }),
-            emcc_options,
+            std.fmt.allocPrint(
+                b.allocator,
+                "-sMALLOC=emmalloc {s}",
+                .{emcc_options},
+            ) catch |err| switch (err) {
+                error.OutOfMemory => @panic("Out of memory"),
+            },
         });
         emcc.step.dependOn(pkg.emsdk_setup_step);
 
diff --git a/libs/zemscripten/src/root.zig b/libs/zemscripten/src/root.zig
index e69de29bb..9c2f98610 100644
--- a/libs/zemscripten/src/root.zig
+++ b/libs/zemscripten/src/root.zig
@@ -0,0 +1,62 @@
+const std = @import("std");
+
+extern "C" fn emmalloc_memalign(alignment: usize, size: usize) ?*anyopaque;
+extern "C" fn emmalloc_realloc_try(ptr: ?*anyopaque, size: usize) ?*anyopaque;
+extern "C" fn emmalloc_free(ptr: ?*anyopaque) void;
+
+/// Zig Allocator that wraps emmalloc
+/// use with linker flag -sMALLOC=emmalloc
+pub const Allocator = struct {
+    const Self = @This();
+    dummy: u32 = undefined,
+
+    pub fn allocator(self: *Self) std.mem.Allocator {
+        return .{
+            .ptr = self,
+            .vtable = &.{
+                .alloc = &alloc,
+                .resize = &resize,
+                .free = &free,
+            },
+        };
+    }
+
+    fn alloc(
+        ctx: *anyopaque,
+        len: usize,
+        ptr_align_log2: u8,
+        return_address: usize,
+    ) ?[*]u8 {
+        _ = ctx;
+        _ = return_address;
+        const ptr_align = @as(usize, 1) << @as(u5, @intCast(ptr_align_log2));
+        if (!std.math.isPowerOfTwo(ptr_align)) unreachable;
+        const ptr = emmalloc_memalign(ptr_align, len) orelse return null;
+        return @ptrCast(ptr);
+    }
+
+    fn resize(
+        ctx: *anyopaque,
+        buf: []u8,
+        buf_align_log2: u8,
+        new_len: usize,
+        return_address: usize,
+    ) bool {
+        _ = ctx;
+        _ = return_address;
+        _ = buf_align_log2;
+        return emmalloc_realloc_try(buf.ptr, new_len) != null;
+    }
+
+    fn free(
+        ctx: *anyopaque,
+        buf: []u8,
+        buf_align_log2: u8,
+        return_address: usize,
+    ) void {
+        _ = ctx;
+        _ = buf_align_log2;
+        _ = return_address;
+        return emmalloc_free(buf.ptr);
+    }
+};