Better debug allocator

src/debug_allocator.zig

@@ -0,0 +1,124 @@
+//! This allocator collects information about allocation sizes
+
+const std = @import("std");
+
+const DebugAllocator = @This();
+
+fn toMB(value: var) f64 {
+    return switch (@TypeOf(value)) {
+        f64 => value / (1024 * 1024),
+        else => @intToFloat(f64, value) / (1024 * 1024),
+    };
+}
+
+const Stats = struct {
+    mean: f64 = 0,
+    mean_of_squares: f64 = 0,
+    total: usize = 0,
+    count: usize = 0,
+
+    fn addSample(self: *Stats, value: usize) void {
+        const count_f64 = @intToFloat(f64, self.count);
+        self.mean = (self.mean * count_f64 + @intToFloat(f64, value)) / (count_f64 + 1);
+        self.mean_of_squares = (self.mean_of_squares * count_f64 + @intToFloat(f64, value * value)) / (count_f64 + 1);
+        self.total += value;
+        self.count += 1;
+    }
+
+    fn stdDev(self: Stats) f64 {
+        return std.math.sqrt(self.mean_of_squares - self.mean * self.mean);
+    }
+};
+
+pub const AllocationInfo = struct {
+    allocation_stats: Stats = Stats{},
+    deallocation_count: usize = 0,
+    deallocation_total: usize = 0,
+
+    reallocation_stats: Stats = Stats{},
+    shrink_stats: Stats = Stats{},
+
+    pub fn format(
+        self: AllocationInfo,
+        comptime fmt: []const u8,
+        options: std.fmt.FormatOptions,
+        out_stream: var,
+    ) !void {
+        @setEvalBranchQuota(2000);
+
+        // TODO: Make these behave like {Bi}, which doesnt work on floating point numbers.
+        return std.fmt.format(
+            out_stream,
+            \\------------------------------------------ Allocation info ------------------------------------------
+            \\{} total allocations (total: {d:.2} MB, mean: {d:.2} MB, std. dev: {d:.2} MB), {} deallocations
+            \\{} current allocations ({d:.2} MB)
+            \\{} reallocations (total: {d:.2} MB, mean: {d:.2} MB, std. dev: {d:.2} MB)
+            \\{} shrinks (total: {d:.2} MB, mean: {d:.2} MB, std. dev: {d:.2} MB)
+            \\-----------------------------------------------------------------------------------------------------
+        ,
+            .{
+                self.allocation_stats.count,
+                toMB(self.allocation_stats.total),
+                toMB(self.allocation_stats.mean),
+                toMB(self.allocation_stats.stdDev()),
+                self.deallocation_count,
+                self.allocation_stats.count - self.deallocation_count,
+                toMB(self.allocation_stats.total + self.reallocation_stats.total - self.deallocation_total - self.shrink_stats.total),
+                self.reallocation_stats.count,
+                toMB(self.reallocation_stats.total),
+                toMB(self.reallocation_stats.mean),
+                toMB(self.reallocation_stats.stdDev()),
+                self.shrink_stats.count,
+                toMB(self.shrink_stats.total),
+                toMB(self.shrink_stats.mean),
+                toMB(self.shrink_stats.stdDev()),
+            },
+        );
+    }
+};
+
+base_allocator: *std.mem.Allocator,
+info: AllocationInfo,
+
+// Interface implementation
+allocator: std.mem.Allocator,
+
+pub fn init(base_allocator: *std.mem.Allocator) DebugAllocator {
+    return .{
+        .base_allocator = base_allocator,
+        .info = .{},
+        .allocator = .{
+            .reallocFn = realloc,
+            .shrinkFn = shrink,
+        },
+    };
+}
+
+fn realloc(allocator: *std.mem.Allocator, old_mem: []u8, old_align: u29, new_size: usize, new_align: u29) ![]u8 {
+    const self = @fieldParentPtr(DebugAllocator, "allocator", allocator);
+    var data = try self.base_allocator.reallocFn(self.base_allocator, old_mem, old_align, new_size, new_align);
+    if (old_mem.len == 0) {
+        self.info.allocation_stats.addSample(new_size);
+    } else if (new_size > old_mem.len) {
+        self.info.reallocation_stats.addSample(new_size - old_mem.len);
+    } else if (new_size < old_mem.len) {
+        self.info.shrink_stats.addSample(old_mem.len - new_size);
+    }
+    return data;
+}
+
+fn shrink(allocator: *std.mem.Allocator, old_mem: []u8, old_align: u29, new_size: usize, new_align: u29) []u8 {
+    const self = @fieldParentPtr(DebugAllocator, "allocator", allocator);
+    if (new_size == 0) {
+        if (self.info.allocation_stats.count == self.info.deallocation_count) {
+            @panic("error - too many calls to free, most likely double free");
+        }
+        self.info.deallocation_total += old_mem.len;
+        self.info.deallocation_count += 1;
+    } else if (new_size < old_mem.len) {
+        self.info.shrink_stats.addSample(old_mem.len - new_size);
+    } else if (new_size > old_mem.len) {
+        @panic("error - trying to shrink to a bigger size");
+    }
+    return self.base_allocator.shrinkFn(self.base_allocator, old_mem, old_align, new_size, new_align);
+}

src/main.zig

@@ -3,6 +3,7 @@ const build_options = @import("build_options");
 
 const Config = @import("config.zig");
 const DocumentStore = @import("document_store.zig");
+const DebugAllocator = @import("debug_allocator.zig");
 const data = @import("data/" ++ build_options.data_version ++ ".zig");
 const types = @import("types.zig");
 const analysis = @import("analysis.zig");
@@ -565,9 +566,9 @@ fn processJsonRpc(parser: *std.json.Parser, json: []const u8, config: Config) !v
     }
 }
 
-var debug_alloc_state: std.testing.LeakCountAllocator = undefined;
+var debug_alloc_state: DebugAllocator = undefined;
 // We can now use if(leak_count_alloc) |alloc| { ... } as a comptime check.
-const debug_alloc: ?*std.testing.LeakCountAllocator = if (build_options.allocation_info) &debug_alloc_state else null;
+const debug_alloc: ?*DebugAllocator = if (build_options.allocation_info) &debug_alloc_state else null;
 
 pub fn main() anyerror!void {
     // TODO: Use a better purpose general allocator once std has one.
@@ -578,7 +579,7 @@ pub fn main() anyerror!void {
     if (build_options.allocation_info) {
         // TODO: Use a better debugging allocator, track size in bytes, memory reserved etc..
         // Initialize the leak counting allocator.
-        debug_alloc_state = std.testing.LeakCountAllocator.init(allocator);
+        debug_alloc_state = DebugAllocator.init(allocator);
         allocator = &debug_alloc_state.allocator;
     }
 
@@ -708,7 +709,7 @@ pub fn main() anyerror!void {
         offset += bytes_read;
 
         if (debug_alloc) |dbg| {
-            try log("Allocations alive: {}", .{dbg.count});
+            try log("{}", .{dbg.info});
         }
     }
 }