From 4f0762af62a42e8e8713dd17f265b19ea72ee9ea Mon Sep 17 00:00:00 2001
From: Auguste Rame <19855629+SuperAuguste@users.noreply.github.com>
Date: Sun, 4 Jun 2023 02:42:56 -0400
Subject: [PATCH] Use @silversquirl's binned allocator (#1134)

* Use @silversquirl's binned allocator

* Finish this
---
 build.zig     |  6 ++++++
 build.zig.zon |  4 ++++
 src/main.zig  | 18 ++++++++++++++----
 3 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/build.zig b/build.zig
index ff04fe4..3ceacfa 100644
--- a/build.zig
+++ b/build.zig
@@ -39,6 +39,7 @@ pub fn build(b: *std.build.Builder) !void {
     exe_options.addOption(bool, "enable_tracy_callstack", b.option(bool, "enable_tracy_callstack", "Enable callstack graphs.") orelse enable_tracy);
     exe_options.addOption(bool, "enable_failing_allocator", b.option(bool, "enable_failing_allocator", "Whether to use a randomly failing allocator.") orelse false);
     exe_options.addOption(u32, "enable_failing_allocator_likelihood", b.option(u32, "enable_failing_allocator_likelihood", "The chance that an allocation will fail is `1/likelihood`") orelse 256);
+    exe_options.addOption(bool, "use_gpa", b.option(bool, "use_gpa", "Good for debugging") orelse (optimize == .Debug));
 
     const version = v: {
         const version_string = b.fmt("{d}.{d}.{d}", .{ zls_version.major, zls_version.minor, zls_version.patch });
@@ -88,6 +89,9 @@ pub fn build(b: *std.build.Builder) !void {
     const diffz_module = b.dependency("diffz", .{}).module("diffz");
     exe.addModule("diffz", diffz_module);
 
+    const binned_allocator_module = b.dependency("binned_allocator", .{}).module("binned_allocator");
+    exe.addModule("binned_allocator", binned_allocator_module);
+
     if (enable_tracy) {
         const client_cpp = "src/tracy/public/TracyClient.cpp";
 
@@ -119,6 +123,7 @@ pub fn build(b: *std.build.Builder) !void {
             .{ .name = "known-folders", .module = known_folders_module },
             .{ .name = "tres", .module = tres_module },
             .{ .name = "diffz", .module = diffz_module },
+            .{ .name = "binned_allocator", .module = binned_allocator_module },
             .{ .name = "build_options", .module = build_options_module },
         },
     });
@@ -154,6 +159,7 @@ pub fn build(b: *std.build.Builder) !void {
     tests.addModule("zls", zls_module);
     tests.addModule("tres", tres_module);
     tests.addModule("diffz", diffz_module);
+    tests.addModule("binned_allocator", binned_allocator_module);
     test_step.dependOn(&b.addRunArtifact(tests).step);
 
     var src_tests = b.addTest(.{
diff --git a/build.zig.zon b/build.zig.zon
index b5132f5..6abcca1 100644
--- a/build.zig.zon
+++ b/build.zig.zon
@@ -15,5 +15,9 @@
             .url = "https://github.com/ziglibs/diffz/archive/b966296b4489eb082b0831ec9a37d6f5e1906040.tar.gz",
             .hash = "1220ed4aed884221108ad39f2658b69a91653e0bbc8ce429bc7f1bc4e58f6a751553",
         },
+        .binned_allocator = .{
+            .url = "https://gist.github.com/silversquirl/c1e4840048fdf48e669b6eac76d80634/archive/8bbe137e65f26854ff936046d884a45d4fa156de.tar.gz",
+            .hash = "1220044bd25cc02da476d0ddf988667a29751dae6bf988128b32ac258b21c23d0f47",
+        },
     },
 }
diff --git a/src/main.zig b/src/main.zig
index c2d7fa4..20454c9 100644
--- a/src/main.zig
+++ b/src/main.zig
@@ -9,6 +9,7 @@ const configuration = @import("configuration.zig");
 const Server = @import("Server.zig");
 const Header = @import("Header.zig");
 const debug = @import("debug.zig");
+const binned_allocator = @import("binned_allocator");
 
 const logger = std.log.scoped(.zls_main);
 const message_logger = std.log.scoped(.message);
@@ -361,11 +362,20 @@ const stack_frames = switch (zig_builtin.mode) {
 };
 
 pub fn main() !void {
-    var gpa_state = std.heap.GeneralPurposeAllocator(.{ .stack_trace_frames = stack_frames }){};
-    defer std.debug.assert(gpa_state.deinit() == .ok);
+    var allocator_state = if (build_options.use_gpa)
+        std.heap.GeneralPurposeAllocator(.{ .stack_trace_frames = stack_frames }){}
+    else
+        binned_allocator.BinnedAllocator(.{ .thread_safe = false }){};
 
-    var tracy_state = if (tracy.enable_allocation) tracy.tracyAllocator(gpa_state.allocator()) else void{};
-    const inner_allocator: std.mem.Allocator = if (tracy.enable_allocation) tracy_state.allocator() else gpa_state.allocator();
+    defer {
+        if (build_options.use_gpa)
+            std.debug.assert(allocator_state.deinit() == .ok)
+        else
+            allocator_state.deinit();
+    }
+
+    var tracy_state = if (tracy.enable_allocation) tracy.tracyAllocator(allocator_state.allocator()) else void{};
+    const inner_allocator: std.mem.Allocator = if (tracy.enable_allocation) tracy_state.allocator() else allocator_state.allocator();
 
     var failing_allocator_state = if (build_options.enable_failing_allocator) debug.FailingAllocator.init(inner_allocator, build_options.enable_failing_allocator_likelihood) else void{};
     const allocator: std.mem.Allocator = if (build_options.enable_failing_allocator) failing_allocator_state.allocator() else inner_allocator;