solve day 20 part 2

src/day20.zig

@@ -59,8 +59,10 @@ const ModuleId = u8;
 const InternalModule = struct {
     const max_inputs = 12;
 
+    const StateBitSet = std.bit_set.IntegerBitSet(max_inputs);
+
     module_type: ModuleType,
-    state: std.bit_set.IntegerBitSet(max_inputs),
+    state: StateBitSet,
     inputs: []ModuleId,
     outputs: []ModuleId,
 
@@ -177,6 +179,16 @@ const MachineState = struct {
         }
         return null;
     }
+
+    fn get_name(self: MachineState, id: ModuleId) []const u8 {
+        return self.names[id].constSlice();
+    }
+
+    fn reset(self: *MachineState) void {
+        for (self.state) |*module| {
+            module.state = InternalModule.StateBitSet.initEmpty();
+        }
+    }
 };
 
 fn parse_module(allocator: Allocator, line: []const u8) !Module {
@@ -226,6 +238,10 @@ const Pulse = struct {
     from: ModuleId,
     to: ModuleId,
 };
+const PulseCount = struct {
+    low: u64,
+    high: u64
+};
 
 fn simulate_pulse(pulse_queue: *std.ArrayList(Pulse), machine_state: *MachineState, pulse: Pulse) !void {
     var to_module = &machine_state.state[pulse.to];
@@ -251,7 +267,7 @@ fn simulate_pulse(pulse_queue: *std.ArrayList(Pulse), machine_state: *MachineSta
     }
 }
 
-fn simulate_button_press(allocator: Allocator, machine_state: *MachineState, broadcaster: ModuleId) !struct { low: u64, high: u64 } {
+fn simulate_button_press(allocator: Allocator, machine_state: *MachineState, broadcaster: ModuleId) !PulseCount {
     var pulses = std.ArrayList(Pulse).init(allocator);
     defer pulses.deinit();
 
@@ -300,6 +316,165 @@ pub fn part1(input: *aoc.Input) !aoc.Result {
     return .{ .uint = low_count * high_count };
 }
 
+const MachineCounter = struct {
+    flip_flops: []ModuleId,
+    input: ModuleId,
+    output: ModuleId
+};
+
+const MachineAnalysis = struct {
+    allocator: Allocator,
+    counters: []MachineCounter,
+    counter_merge: ModuleId,
+
+    fn deinit(self: MachineAnalysis) void {
+        for (self.counters) |counter| {
+            self.allocator.free(counter.flip_flops);
+        }
+        self.allocator.free(self.counters);
+    }
+};
+
+const ModuleIdHashSet = std.AutoHashMap(ModuleId, void);
+
+fn find_coverage(allocator: Allocator, machine_state: MachineState, from: ModuleId) !ModuleIdHashSet {
+    var seen = ModuleIdHashSet.init(allocator);
+    errdefer seen.deinit();
+
+    var stack = std.ArrayList(ModuleId).init(allocator);
+    defer stack.deinit();
+
+    try stack.append(from);
+    try seen.put(from, {});
+
+    while (stack.popOrNull()) |node_idx| {
+        for (machine_state.state[node_idx].outputs) |output| {
+            if (seen.contains(output)) continue;
+
+            try stack.append(output);
+            try seen.put(output, {});
+        }
+    }
+
+    return seen;
+}
+
+fn analyze_machine(allocator: Allocator, machine_state: MachineState, input_module: ModuleId, output_module: ModuleId) !MachineAnalysis {
+    const counters = try allocator.alloc(MachineCounter, machine_state.state[input_module].outputs.len);
+    errdefer allocator.free(counters);
+
+    var coverage_counts = try allocator.alloc(u32, machine_state.state.len);
+    defer allocator.free(coverage_counts);
+    @memset(coverage_counts, 0);
+
+    for (0.., machine_state.state[input_module].outputs) |i, counter_input| {
+        var coverage = try find_coverage(allocator, machine_state, counter_input);
+        defer coverage.deinit();
+
+        {
+            var coverage_iter = coverage.keyIterator();
+            while (coverage_iter.next()) |covered_module| {
+                coverage_counts[covered_module.*] += 1;
+            }
+        }
+
+        counters[i].input = counter_input;
+
+        assert(coverage.contains(output_module));
+        counters[i].output = output_module: {
+            var prev = output_module;
+            var current = output_module;
+            outer: while (true) {
+                var covered_input: ?ModuleId = null;
+                for (machine_state.state[current].inputs) |input| {
+                    if (!coverage.contains(input)) continue;
+
+                    if (covered_input != null) break :outer; // If a second covered input was found, stop search
+                    covered_input = input;
+                }
+                if (covered_input == null) break;
+
+                prev = current;
+                current = covered_input.?;
+            }
+
+            break :output_module prev;
+        };
+
+        var flip_flop_count: u32 = 0;
+        {
+            var iter = coverage.keyIterator();
+            while (iter.next()) |id| {
+                if (machine_state.state[id.*].module_type == .FlipFlop) {
+                    flip_flop_count += 1;
+                }
+            }
+        }
+
+        counters[i].flip_flops = try allocator.alloc(ModuleId, flip_flop_count);
+        {
+            var flip_flop_i: usize = 0;
+            var iter = coverage.keyIterator();
+            while (iter.next()) |id| {
+                if (machine_state.state[id.*].module_type == .FlipFlop) {
+                    counters[i].flip_flops[flip_flop_i] = id.*;
+                    flip_flop_i += 1;
+                }
+            }
+        }
+    }
+
+    var counter_merge: ?ModuleId = null;
+    for (0.., coverage_counts) |id, count| {
+        if (count != machine_state.state[input_module].outputs.len) continue;
+        if (id == output_module) continue;
+
+        counter_merge = @intCast(id);
+    }
+
+    if (counter_merge == null) return error.NoMergeModule;
+
+    return MachineAnalysis{
+        .allocator = allocator,
+        .counters = counters,
+        .counter_merge = counter_merge.?,
+    };
+}
+
+fn find_counter_cycle_size(allocator: Allocator, machine_state: *MachineState, counter: MachineCounter, input_module: ModuleId) !u64 {
+    machine_state.reset();
+
+    var pulses = std.ArrayList(Pulse).init(allocator);
+    defer pulses.deinit();
+
+    var pulse_count: u64 = 0;
+    pulse_loop: while (true) {
+        pulse_count += 1;
+
+        try pulses.append(Pulse{
+            .is_high = false,
+            .to = counter.input,
+            .from = input_module,
+        });
+
+        while (pulses.items.len > 0) {
+            const pulse: Pulse = pulses.orderedRemove(0);
+
+            if (pulse.to == counter.output) {
+                if (!pulse.is_high) break :pulse_loop;
+            } else {
+                try simulate_pulse(&pulses, machine_state, pulse);
+            }
+        }
+    }
+
+    return pulse_count;
+}
+
+fn lcm(a: u64, b: u64) u64 {
+    return a * b / std.math.gcd(a, b);
+}
+
 pub fn part2(input: *aoc.Input) !aoc.Result {
     const allocator = input.allocator;
     const parsed = try parse_input(allocator, input.lines);
@@ -308,7 +483,19 @@ pub fn part2(input: *aoc.Input) !aoc.Result {
     var machine_state = try MachineState.init(allocator, parsed.modules.items);
     defer machine_state.deinit();
 
-    return .{ .uint = 0 };
+    const broadcaster_idx = machine_state.get_name_idx(try NameString.fromSlice("broadcaster")).?;
+    const rx_idx = machine_state.get_name_idx(try NameString.fromSlice("rx")).?;
+
+    const analysis = try analyze_machine(allocator, machine_state, broadcaster_idx, rx_idx);
+    defer analysis.deinit();
+
+    var answer: u64 = 1;
+    for (analysis.counters) |counter| {
+        const cycle_size = try find_counter_cycle_size(allocator, &machine_state, counter, broadcaster_idx);
+        answer = lcm(answer, cycle_size);
+    }
+
+    return .{ .uint = answer };
 }
 
 const example_input1 = [_][]const u8{