From 49d8723408233ac1c8942eca0a52ca12188040ce Mon Sep 17 00:00:00 2001
From: Lee Cannon <leecannon@leecannon.xyz>
Date: Sun, 31 Oct 2021 14:16:59 +0000
Subject: [PATCH] add functionality to trace allocations

---
 build.zig     | 2 ++
 src/main.zig  | 7 +++++++
 src/tracy.zig | 3 ++-
 3 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/build.zig b/build.zig
index 3b875c1d5e..56abec9e26 100644
--- a/build.zig
+++ b/build.zig
@@ -111,6 +111,7 @@ pub fn build(b: *Builder) !void {
 
     const tracy = b.option([]const u8, "tracy", "Enable Tracy integration. Supply path to Tracy source");
     const tracy_callstack = b.option(bool, "tracy-callstack", "Include callstack information with Tracy data. Does nothing if -Dtracy is not provided") orelse false;
+    const tracy_allocation = b.option(bool, "tracy-allocation", "Include allocation information with Tracy data. Does nothing if -Dtracy is not provided") orelse false;
     const link_libc = b.option(bool, "force-link-libc", "Force self-hosted compiler to link libc") orelse enable_llvm;
     const strip = b.option(bool, "strip", "Omit debug information") orelse false;
 
@@ -266,6 +267,7 @@ pub fn build(b: *Builder) !void {
     exe_options.addOption(bool, "enable_link_snapshots", enable_link_snapshots);
     exe_options.addOption(bool, "enable_tracy", tracy != null);
     exe_options.addOption(bool, "enable_tracy_callstack", tracy_callstack);
+    exe_options.addOption(bool, "enable_tracy_allocation", tracy_allocation);
     exe_options.addOption(bool, "is_stage1", is_stage1);
     exe_options.addOption(bool, "omit_stage2", omit_stage2);
     if (tracy) |tracy_path| {
diff --git a/src/main.zig b/src/main.zig
index 280a69a425..ab9ec7c2d6 100644
--- a/src/main.zig
+++ b/src/main.zig
@@ -10,6 +10,7 @@ const ArrayList = std.ArrayList;
 const Ast = std.zig.Ast;
 const warn = std.log.warn;
 
+const tracy = @import("tracy.zig");
 const Compilation = @import("Compilation.zig");
 const link = @import("link.zig");
 const Package = @import("Package.zig");
@@ -155,6 +156,12 @@ pub fn main() anyerror!void {
     const arena = &arena_instance.allocator;
 
     const args = try process.argsAlloc(arena);
+
+    if (tracy.enable_allocation) {
+        var gpa_tracy = tracy.tracyAllocator(gpa);
+        return mainArgs(&gpa_tracy.allocator, arena, args);
+    }
+
     return mainArgs(gpa, arena, args);
 }
 
diff --git a/src/tracy.zig b/src/tracy.zig
index 9af1368d7e..3ef2df7ca0 100644
--- a/src/tracy.zig
+++ b/src/tracy.zig
@@ -2,7 +2,8 @@ const std = @import("std");
 const builtin = @import("builtin");
 
 pub const enable = if (builtin.is_test) false else @import("build_options").enable_tracy;
-const enable_callstack = @import("build_options").enable_tracy_callstack;
+pub const enable_allocation = enable and @import("build_options").enable_tracy_allocation;
+pub const enable_callstack = enable and @import("build_options").enable_tracy_callstack;
 
 // TODO: make this configurable
 const callstack_depth = 10;