nix component: build hook: DRY by adding readStruct() to protocol

Also make strings const and use `std.BufMap` with unmanaged support instead of own map type.
input-output-hk · Jul 15, 2024 · 6a478d4 · 6a478d4
1 parent e5c21d6
commit 6a478d4
Show file tree

Hide file tree

Showing 2 changed files with 67 additions and 55 deletions.
diff --git a/src/components/nix/build-hook/main.zig b/src/components/nix/build-hook/main.zig
@@ -43,13 +43,13 @@ fn innerMain(allocator: std.mem.Allocator) !void {
     }
 
     var settings = try protocol.readStringStringMap(allocator, stdin);
-    defer settings.deinit(allocator);
+    defer settings.deinit();
 
     if (std.log.defaultLogEnabled(.debug)) {
         var settings_msg = std.ArrayList(u8).init(allocator);
         defer settings_msg.deinit();
 
-        var iter = settings.map.iterator();
+        var iter = settings.iterator();
         while (iter.next()) |entry| {
             try settings_msg.appendNTimes(' ', 2);
             try settings_msg.appendSlice(entry.key_ptr.*);
@@ -62,7 +62,7 @@ fn innerMain(allocator: std.mem.Allocator) !void {
     }
 
     var ifds_file = ifds_file: {
-        const builders = settings.map.get("builders").?;
+        const builders = settings.get("builders").?;
         if (builders.len == 0) {
             std.log.err("expected path to write IFDs to in nix config entry `builders` but it is empty", .{});
             return error.NoBuilders;
@@ -123,32 +123,24 @@ fn innerMain(allocator: std.mem.Allocator) !void {
         }
 
         const drv = blk: {
-            const am_willing = try protocol.readBool(stdin);
-            const needed_system = try protocol.readPacket(allocator, stdin);
-            const drv_path = try protocol.readPacket(allocator, stdin);
-            const required_features = try protocol.readPackets(allocator, stdin);
+            const drv = try protocol.readStruct(Drv, allocator, stdin);
 
-            const gop = try drvs.getOrPut(allocator, drv_path);
+            const gop = try drvs.getOrPut(allocator, drv.drv_path);
 
             if (gop.found_existing) {
-                std.debug.assert(am_willing == gop.value_ptr.am_willing);
+                std.debug.assert(drv.am_willing == gop.value_ptr.am_willing);
 
-                std.debug.assert(std.mem.eql(u8, needed_system, gop.value_ptr.needed_system));
+                std.debug.assert(std.mem.eql(u8, drv.needed_system, gop.value_ptr.needed_system));
 
-                std.debug.assert(std.mem.eql(u8, drv_path, gop.value_ptr.drv_path));
+                std.debug.assert(std.mem.eql(u8, drv.drv_path, gop.value_ptr.drv_path));
 
-                std.debug.assert(required_features.len == gop.value_ptr.required_features.len);
-                for (required_features, gop.value_ptr.required_features) |a, b|
+                std.debug.assert(drv.required_features.len == gop.value_ptr.required_features.len);
+                for (drv.required_features, gop.value_ptr.required_features) |a, b|
                     std.debug.assert(std.mem.eql(u8, a, b));
 
-                std.log.debug("known drv: {s}", .{drv_path});
+                std.log.debug("known drv: {s}", .{drv.drv_path});
             } else {
-                gop.value_ptr.* = .{
-                    .am_willing = am_willing,
-                    .needed_system = needed_system,
-                    .drv_path = drv_path,
-                    .required_features = required_features,
-                };
+                gop.value_ptr.* = drv;
 
                 const drv_json = try std.json.stringifyAlloc(allocator, gop.value_ptr.*, .{});
                 defer allocator.free(drv_json);
@@ -218,23 +210,7 @@ const Accepted = struct {
 /// The nix daemon will start another instance of the build hook for the remaining derivations.
 fn accept(allocator: std.mem.Allocator, store_uri: []const u8) !Accepted {
     try stderr.print("# accept\n{s}\n", .{store_uri});
-
-    const inputs = try protocol.readPackets(allocator, stdin);
-    errdefer {
-        for (inputs) |input| allocator.free(input);
-        allocator.free(inputs);
-    }
-
-    const wanted_outputs = try protocol.readPackets(allocator, stdin);
-    errdefer {
-        for (wanted_outputs) |wanted_output| allocator.free(wanted_output);
-        allocator.free(wanted_outputs);
-    }
-
-    return .{
-        .inputs = inputs,
-        .wanted_outputs = wanted_outputs,
-    };
+    return protocol.readStruct(Accepted, allocator, stdin);
 }
 
 test {

diff --git a/src/components/nix/build-hook/protocol.zig b/src/components/nix/build-hook/protocol.zig
@@ -76,15 +76,15 @@ pub fn readBool(reader: anytype) !bool {
     };
 }
 
-pub fn readPacket(allocator: std.mem.Allocator, reader: anytype) ![]u8 {
+pub fn readPacket(allocator: std.mem.Allocator, reader: anytype) ![]const u8 {
     const buf = try allocator.alloc(u8, try readU64(reader));
     errdefer allocator.free(buf);
     try readPadded(reader, buf);
     return buf;
 }
 
-pub fn readPackets(allocator: std.mem.Allocator, reader: anytype) ![][]u8 {
-    const bufs = try allocator.alloc([]u8, try readU64(reader));
+pub fn readPackets(allocator: std.mem.Allocator, reader: anytype) ![]const []const u8 {
+    const bufs = try allocator.alloc([]const u8, try readU64(reader));
     errdefer {
         for (bufs) |buf| allocator.free(buf);
         allocator.free(bufs);
@@ -93,28 +93,64 @@ pub fn readPackets(allocator: std.mem.Allocator, reader: anytype) ![][]u8 {
     return bufs;
 }
 
-const StringStringMapOwned = struct {
-    map: std.StringHashMapUnmanaged([]u8) = .{},
+pub fn readStringStringMap(allocator: std.mem.Allocator, reader: anytype) !std.BufMap {
+    // We build the `std.BufMap`'s underlying `.hash_map` directly
+    // so that we don't have to copy keys and values twice:
+    // Once when reading and once when inserting.
 
-    pub fn deinit(self: *@This(), alloc: std.mem.Allocator) void {
-        var iter = self.map.iterator();
-        while (iter.next()) |entry| {
-            alloc.free(entry.key_ptr.*);
-            alloc.free(entry.value_ptr.*);
-        }
-        self.map.deinit(alloc);
+    var hash_map = std.StringHashMap([]const u8).init(allocator);
+    errdefer {
+        var buf_map = std.BufMap{ .hash_map = hash_map };
+        buf_map.deinit();
     }
-};
 
-pub fn readStringStringMap(allocator: std.mem.Allocator, reader: anytype) !StringStringMapOwned {
-    var map = StringStringMapOwned{};
-    errdefer map.deinit(allocator);
     while (try readU64(reader) != 0) {
         const key = try readPacket(allocator, reader);
         errdefer allocator.free(key);
+
         const value = try readPacket(allocator, reader);
         errdefer allocator.free(value);
-        try map.map.put(allocator, key, value);
+
+        try hash_map.put(key, value);
+    }
+
+    return std.BufMap{ .hash_map = hash_map };
+}
+
+/// Reads fields in declaration order.
+pub fn readStruct(comptime T: type, allocator: std.mem.Allocator, reader: anytype) !T {
+    var strukt: T = undefined;
+
+    const fields = @typeInfo(T).Struct.fields;
+    inline for (fields, 0..) |field, field_idx| {
+        @field(strukt, field.name) = switch (field.type) {
+            []const u8 => readPacket(allocator, reader),
+            []const []const u8 => readPackets(allocator, reader),
+            u64 => readU64(reader),
+            bool => readBool(reader),
+            std.BufMap => readStringStringMap(allocator, reader),
+            std.StringHashMapUnmanaged([]const u8) => if (readStringStringMap(allocator, reader)) |map|
+                map.hash_map.unmanaged
+            else |err|
+                err,
+            else => @compileError("type \"" ++ @typeName(T) ++ "\" does not exist in the nix protocol"),
+        } catch |err| {
+            inline for (fields[0..field_idx]) |field_| {
+                const field_value = @field(strukt, field.name);
+                switch (field_.type) {
+                    []const u8 => allocator.free(field_value),
+                    []const []const u8 => for (field_value) |item| allocator.free(item),
+                    std.BufMap => field_value.deinit(),
+                    std.StringHashMapUnmanaged([]const u8) => {
+                        var map = std.BufMap{ .hash_map = field_value.promote(allocator) };
+                        map.deinit();
+                    },
+                    else => {},
+                }
+            }
+            return err;
+        };
     }
-    return map;
+
+    return strukt;
 }