Merge pull request #40 from nix-community/feat/better-alloc

Fix memory leaks

Author: RossComputerGuy

src/Dependency.zig

@@ -1,4 +1,13 @@
+const std = @import("std");
+const Dependency = @This();
+
 url: []const u8,
 rev: ?[]const u8,
 nix_hash: ?[]const u8,
 done: bool,
+
+pub fn deinit(self: Dependency, alloc: std.mem.Allocator) void {
+    alloc.free(self.url);
+    if (self.rev) |rev| alloc.free(rev);
+    if (self.nix_hash) |nix_hash| alloc.free(nix_hash);
+}

src/codegen.zig

@@ -19,7 +19,10 @@ pub fn write(alloc: Allocator, out: anytype, deps: StringHashMap(Dependency)) !v
     );
 
     const len = deps.count();
+
     var entries = try alloc.alloc(Entry, len);
+    defer alloc.free(entries);
+
     var iter = deps.iterator();
     for (0..len) |i| {
         entries[i] = iter.next().?;

src/fetch.zig

@@ -27,6 +27,7 @@ const Worker = struct {
 
 pub fn fetch(alloc: Allocator, deps: *StringHashMap(Dependency)) !void {
     var workers = try ArrayList(Worker).initCapacity(alloc, deps.count());
+    defer workers.deinit();
     var done = false;
 
     while (!done) {
@@ -45,12 +46,17 @@ pub fn fetch(alloc: Allocator, deps: *StringHashMap(Dependency)) !void {
                         break :base try fmt.allocPrint(alloc, "tarball+{s}", .{dep.url});
                     }
                 };
+
                 const revi = mem.lastIndexOf(u8, base, "rev=") orelse break :ref base;
                 const refi = mem.lastIndexOf(u8, base, "ref=") orelse break :ref base;
 
+                defer alloc.free(base);
+
                 const i = @min(revi, refi);
-                break :ref base[0..(i - 1)];
+                break :ref try alloc.dupe(u8, base[0..(i - 1)]);
             };
+            defer alloc.free(ref);
+
             log.debug("running \"nix flake prefetch --json --extra-experimental-features 'flakes nix-command' {s}\"", .{ref});
             const argv = &[_][]const u8{ nix, "flake", "prefetch", "--json", "--extra-experimental-features", "flakes nix-command", ref };
             child.* = ChildProcess.init(argv, alloc);
@@ -67,15 +73,18 @@ pub fn fetch(alloc: Allocator, deps: *StringHashMap(Dependency)) !void {
             const child = worker.child;
             const dep = worker.dep;
 
+            defer alloc.destroy(child);
+
             const buf = try child.stdout.?.readToEndAlloc(alloc, std.math.maxInt(usize));
             defer alloc.free(buf);
 
             log.debug("nix prefetch for \"{s}\" returned: {s}", .{ dep.url, buf });
 
-            var fbs = std.io.fixedBufferStream(buf);
-
-            var reader = json.reader(alloc, fbs.reader());
-            const res = try json.parseFromTokenSourceLeaky(Prefetch, alloc, &reader, .{ .ignore_unknown_fields = true });
+            const res = try json.parseFromSlice(Prefetch, alloc, buf, .{
+                .ignore_unknown_fields = true,
+                .allocate = .alloc_always,
+            });
+            defer res.deinit();
 
             switch (try child.wait()) {
                 .Exited => |code| if (code != 0) {
@@ -92,13 +101,15 @@ pub fn fetch(alloc: Allocator, deps: *StringHashMap(Dependency)) !void {
                 },
             }
 
-            assert(res.hash.len != 0);
-            log.debug("hash for \"{s}\" is {s}", .{ dep.url, res.hash });
+            assert(res.value.hash.len != 0);
+            log.debug("hash for \"{s}\" is {s}", .{ dep.url, res.value.hash });
 
-            dep.nix_hash = res.hash;
+            dep.nix_hash = try alloc.dupe(u8, res.value.hash);
             dep.done = true;
 
-            const path = try fmt.allocPrint(alloc, "{s}" ++ fs.path.sep_str ++ "build.zig.zon", .{res.storePath});
+            const path = try fmt.allocPrint(alloc, "{s}" ++ fs.path.sep_str ++ "build.zig.zon", .{res.value.storePath});
+            defer alloc.free(path);
+
             const file = fs.openFileAbsolute(path, .{}) catch |err| switch (err) {
                 error.FileNotFound => continue,
                 else => return err,

src/main.zig

@@ -1,3 +1,4 @@
+const builtin = @import("builtin");
 const std = @import("std");
 const StringHashMap = std.StringHashMap;
 const fs = std.fs;
@@ -10,6 +11,20 @@ const fetch = @import("fetch.zig").fetch;
 const parse = @import("parse.zig").parse;
 const write = @import("codegen.zig").write;
 
+const zig_legacy_version = (std.SemanticVersion{
+    .major = builtin.zig_version.major,
+    .minor = builtin.zig_version.minor,
+    .patch = builtin.zig_version.patch,
+}).order(.{
+    .major = 0,
+    .minor = 14,
+    .patch = 0,
+}) == .lt;
+
+const DebugAllocator = @field(std.heap, if (zig_legacy_version) "GeneralPurposeAllocator" else "DebugAllocator");
+
+var debug_allocator: DebugAllocator(.{}) = if (zig_legacy_version) .{} else .init;
+
 pub fn main() !void {
     var args = process.args();
     _ = args.skip();
@@ -24,16 +39,31 @@ pub fn main() !void {
         dir.openFile("build.zig.zon", .{});
     defer file.close();
 
-    var arena = heap.ArenaAllocator.init(heap.page_allocator);
-    defer arena.deinit();
-    const alloc = arena.allocator();
+    const gpa, const is_debug = gpa: {
+        break :gpa switch (builtin.mode) {
+            .Debug, .ReleaseSafe => .{ debug_allocator.allocator(), true },
+            .ReleaseFast, .ReleaseSmall => .{ std.heap.smp_allocator, false },
+        };
+    };
+    defer if (is_debug) {
+        _ = debug_allocator.deinit();
+    };
+
+    var deps = StringHashMap(Dependency).init(gpa);
+    defer {
+        var iter = deps.iterator();
+        while (iter.next()) |entry| {
+            gpa.free(entry.key_ptr.*);
+            entry.value_ptr.deinit(gpa);
+        }
+        deps.deinit();
+    }
 
-    var deps = StringHashMap(Dependency).init(alloc);
-    try parse(alloc, &deps, file);
-    try fetch(alloc, &deps);
+    try parse(gpa, &deps, file);
+    try fetch(gpa, &deps);
 
     var out = io.bufferedWriter(io.getStdOut().writer());
-    try write(alloc, out.writer(), deps);
+    try write(gpa, out.writer(), deps);
     try out.flush();
 }
 

src/parse.zig

@@ -23,17 +23,23 @@ const zig_legacy_version = (std.SemanticVersion{
 
 pub fn parse(alloc: Allocator, deps: *StringHashMap(Dependency), file: File) !void {
     const content = try alloc.allocSentinel(u8, try file.getEndPos(), 0);
+    defer alloc.free(content);
+
     _ = try file.reader().readAll(content);
 
-    const ast = try Ast.parse(alloc, content, .zon);
+    var ast = try Ast.parse(alloc, content, .zon);
+    defer ast.deinit(alloc);
 
     var root_buf: [2]Index = undefined;
     const root_init = ast.fullStructInit(&root_buf, @field(ast.nodes.items(.data)[0], if (zig_legacy_version) "lhs" else "node")) orelse {
         return error.ParseError;
     };
 
     for (root_init.ast.fields) |field_idx| {
-        if (!mem.eql(u8, try parseFieldName(alloc, ast, field_idx), "dependencies")) {
+        const field_name = try parseFieldName(alloc, ast, field_idx);
+        defer alloc.free(field_name);
+
+        if (!mem.eql(u8, field_name, "dependencies")) {
             continue;
         }
 
@@ -61,18 +67,21 @@ pub fn parse(alloc: Allocator, deps: *StringHashMap(Dependency), file: File) !vo
 
             for (dep_init.ast.fields) |dep_field_idx| {
                 const name = try parseFieldName(alloc, ast, dep_field_idx);
+                defer alloc.free(name);
 
                 if (mem.eql(u8, name, "url")) {
                     const parsed_url = try parseString(alloc, ast, dep_field_idx);
                     if (std.mem.startsWith(u8, parsed_url, "https://")) {
                         url = parsed_url;
                     } else if (std.mem.startsWith(u8, parsed_url, "git+https://")) {
+                        defer alloc.free(parsed_url);
+
                         const url_end = std.mem.indexOf(u8, parsed_url[0..], "#").?;
                         const raw_url = parsed_url[4..url_end];
                         const hash_start = url_end + 1; // +1 to skip the '#'
                         const git_hash = parsed_url[hash_start..];
-                        url = raw_url;
-                        dep.rev = git_hash;
+                        url = try alloc.dupe(u8, raw_url);
+                        dep.rev = try alloc.dupe(u8, git_hash);
                     }
                 } else if (mem.eql(u8, name, "hash")) {
                     hash = try parseString(alloc, ast, dep_field_idx);
@@ -91,7 +100,7 @@ pub fn parse(alloc: Allocator, deps: *StringHashMap(Dependency), file: File) !vo
 
 fn parseFieldName(alloc: Allocator, ast: Ast, idx: Index) ![]const u8 {
     const name = ast.tokenSlice(ast.firstToken(idx) - 2);
-    return if (name[0] == '@') string_literal.parseAlloc(alloc, name[1..]) else name;
+    return if (name[0] == '@') string_literal.parseAlloc(alloc, name[1..]) else alloc.dupe(u8, name);
 }
 
 fn parseString(alloc: Allocator, ast: Ast, idx: Index) ![]const u8 {