From 71797b20fd3a761dccce32a0326b7ef0a40a64ca Mon Sep 17 00:00:00 2001
From: shilangyu <xmarcinmarcin@gmail.com>
Date: Wed, 24 Jan 2024 22:43:46 +0100
Subject: [PATCH] Abstract the allocator

---
 src/main.zig   |  7 +++----
 src/search.zig | 12 ++++++------
 src/utils.zig  | 41 +++++++++++++++++++++++++++++++++++++++++
 3 files changed, 50 insertions(+), 10 deletions(-)

diff --git a/src/main.zig b/src/main.zig
index 1d7d6ed..3a7a940 100644
--- a/src/main.zig
+++ b/src/main.zig
@@ -27,10 +27,9 @@ const SearchResult = struct {
 };
 
 pub fn main() !void {
-    // TODO: replace allocators, maybe https://github.com/kprotty/zap/blob/54cd494257915e6c126a0b70f95789b669b49b96/benchmarks/zig/async.zig#L60
-    var gpa = std.heap.GeneralPurposeAllocator(.{}){};
-    const allocator = gpa.allocator();
-    defer _ = gpa.deinit();
+    var heap = utils.HeapAllocator.init();
+    defer heap.deinit();
+    const allocator = heap.allocator();
 
     var args = try ParsedArgs.parse(allocator);
     defer args.deinit();
diff --git a/src/search.zig b/src/search.zig
index 586e094..73fe12c 100644
--- a/src/search.zig
+++ b/src/search.zig
@@ -6,10 +6,10 @@ const Box = utils.Box;
 const ThreadPoolState = struct {
     matches: std.ArrayList(SearchMatch),
     read_buffer: []u8,
-    allocator: Box(std.heap.GeneralPurposeAllocator(.{})),
+    allocator: Box(utils.HeapAllocator),
 
     fn create(allocator: std.mem.Allocator) !@This() {
-        var allocat = try Box(std.heap.GeneralPurposeAllocator(.{})).init(allocator, std.heap.GeneralPurposeAllocator(.{}){});
+        var allocat = try Box(utils.HeapAllocator).init(allocator, utils.HeapAllocator.init());
 
         return .{
             .matches = std.ArrayList(SearchMatch).init(allocat.ptr.allocator()),
@@ -24,7 +24,7 @@ const ThreadPoolState = struct {
         for (states.items) |*e| totalMatches += e.matches.items.len;
 
         var matches = try std.ArrayList(SearchMatch).initCapacity(allocator, totalMatches);
-        var allocators = try std.ArrayList(Box(std.heap.GeneralPurposeAllocator(.{}))).initCapacity(allocator, states.items.len);
+        var allocators = try std.ArrayList(Box(utils.HeapAllocator)).initCapacity(allocator, states.items.len);
         for (states.items) |*e| {
             e.allocator.ptr.allocator().free(e.read_buffer);
 
@@ -48,7 +48,7 @@ const ThreadPoolState = struct {
 
         self.allocator.ptr.allocator().free(self.read_buffer);
 
-        std.debug.assert(self.allocator.ptr.deinit() == .ok);
+        self.allocator.ptr.deinit();
         self.allocator.deinit();
     }
 };
@@ -56,7 +56,7 @@ const ThreadPoolState = struct {
 /// A merged result of a single bucket search.
 pub const SearchResult = struct {
     matches: std.ArrayList(SearchMatch),
-    allocators: std.ArrayList(Box(std.heap.GeneralPurposeAllocator(.{}))),
+    allocators: std.ArrayList(Box(utils.HeapAllocator)),
 
     fn sortMatches(self: *@This()) void {
         const Sort = struct {
@@ -74,7 +74,7 @@ pub const SearchResult = struct {
         self.matches.deinit();
 
         for (self.allocators.items) |*e| {
-            std.debug.assert(e.ptr.deinit() == .ok);
+            e.ptr.deinit();
             e.deinit();
         }
         self.allocators.deinit();
diff --git a/src/utils.zig b/src/utils.zig
index b42a3a3..19998f7 100644
--- a/src/utils.zig
+++ b/src/utils.zig
@@ -1,6 +1,47 @@
 const std = @import("std");
+const builtin = @import("builtin");
 const assert = std.debug.assert;
 
+/// A heap allocator optimized for each platform/compilation mode.
+pub const HeapAllocator = struct {
+    const is_windows = builtin.target.os.tag == .windows;
+
+    backing_allocator: if (builtin.mode == .Debug)
+        std.heap.GeneralPurposeAllocator(.{})
+    else if (is_windows)
+        std.heap.HeapAllocator
+    else if (builtin.link_libc)
+        void
+    else
+        @compileError("When not running in debug mode, you must use windows or link to libc as zig does not provide a fast, libc-less, general purpose allocator (yet)"),
+
+    pub fn init() @This() {
+        return .{ .backing_allocator = if (builtin.mode == .Debug)
+            std.heap.GeneralPurposeAllocator(.{}){}
+        else if (is_windows)
+            std.heap.HeapAllocator.init()
+        else if (builtin.link_libc) {} };
+    }
+
+    pub fn allocator(self: *@This()) std.mem.Allocator {
+        if (builtin.mode == .Debug) {
+            return self.backing_allocator.allocator();
+        } else if (is_windows) {
+            return self.backing_allocator.allocator();
+        } else if (builtin.link_libc) {
+            return std.heap.c_allocator;
+        }
+    }
+
+    pub fn deinit(self: *@This()) void {
+        if (builtin.mode == .Debug) {
+            assert(self.backing_allocator.deinit() == .ok);
+        } else if (is_windows) {
+            self.backing_allocator.deinit();
+        }
+    }
+};
+
 /// Concatenates two slices into a newly allocated buffer. Returns an owned reference slice into the contents.
 pub fn concatOwned(allocator: std.mem.Allocator, a: []const u8, b: []const u8) ![]const u8 {
     const result = try allocator.alloc(u8, a.len + b.len);