Advent of Code 2020 solutions in Zig
const std = @import("std");

const PATH = "input/day19.txt";
const Str = []const u8;

const RuleNums = u8;
const Rule = [][]RuleNums;

pub fn first(allocator: ?std.mem.Allocator) anyerror!usize {
    var arena = std.heap.ArenaAllocator.init(allocator.?);
    defer arena.deinit();

    var monster = try parseInput(arena.allocator(), @embedFile(PATH));
    defer monster.rules.deinit();

    var good = std.ArrayList([]Str).init(arena.allocator());
    defer good.deinit();

    try good.append(try monster.genRule(arena.allocator(), 42));
    try good.append(try monster.genRule(arena.allocator(), 31));

    const chunk_size = 8;

    var count: usize = 0;
    messages: for (monster.messages.items) |message| {
        // The only good version is [42][42][31]
        if (message.len != chunk_size * 3) continue;

        var active_chunk: usize = 0;
        var pos: usize = 0;
        chunks: while (pos < message.len - chunk_size + 1) {
            if (pos == chunk_size * 2) active_chunk = 1;
            var chunk = message[pos .. pos + chunk_size];

            for (good.items[active_chunk]) |gm| {
                if (std.mem.eql(u8, chunk, gm)) {
                    pos += chunk_size;
                    continue :chunks;
                }
            }

            if (pos != message.len) {
                continue :messages;
            }
        }
        count += 1;
    }

    return count;
}

// Rule zero is this: 0: 8 11
// We rewrite 8 and 11 to this:
// 8: 42 | 42 8
// 11: 42 31 | 42 11 31
// So the good messages has to start with 1 or more [42] and
// continue with at least one [42][31], but between them 0 or more [42][31] allowed.
//
// Good:
// [42][42][31] (this is the shortest good one)
// [42][42][42][31][31] (embedded one extra [11])
// [42][42][42][42][42][31][31] (embedded one extra [11] and two extra [8] after the first one)
//
// Bad:
// [42][42][31][31] ([42] is missing from [8], or [11] is not symmetric)
pub fn second(allocator: ?std.mem.Allocator) anyerror!usize {
    var arena = std.heap.ArenaAllocator.init(allocator.?);
    defer arena.deinit();

    var monster = try parseInput(arena.allocator(), @embedFile(PATH));
    defer monster.rules.deinit();

    var good = std.ArrayList([]Str).init(arena.allocator());
    defer good.deinit();

    try good.append(try monster.genRule(arena.allocator(), 42));
    try good.append(try monster.genRule(arena.allocator(), 31));

    const chunk_size = 8;

    var count: usize = 0;
    messages: for (monster.messages.items) |message| {
        // shortest message must be [42][42][31] long
        if (message.len < chunk_size * 3) continue;

        var chunk_counter: [2]usize = .{ 0, 0 };

        var active_chunk: usize = 0;
        var pos: usize = 0;
        chunks: while (pos < message.len - chunk_size + 1) {
            // the last chunk has to be a 31, can not be 42
            if (pos == message.len - chunk_size) active_chunk = 1;

            var chunk = message[pos .. pos + chunk_size];

            for (good.items[active_chunk]) |gm| {
                if (std.mem.eql(u8, chunk, gm)) {
                    pos += chunk_size;
                    chunk_counter[active_chunk] += 1;
                    continue :chunks;
                }
            }
            if (active_chunk == 0) {
                active_chunk = 1;
                continue :chunks;
            }

            if (pos != message.len) {
                continue :messages;
            }
        }
        // Number of [42] items has to be more than [31] ones
        // so accidental early [31] matches will not count.
        if (chunk_counter[0] > chunk_counter[1]) count += 1;
    }

    return count;
}

const Monster = struct {
    rules: std.AutoHashMap(RuleNums, Rule),
    cache: std.AutoHashMap(RuleNums, []Str),
    messages: std.ArrayList(Str),

    fn genRule(
        self: *@This(),
        allocator: std.mem.Allocator,
        ruleNum: RuleNums,
    ) anyerror![]Str {
        if (self.cache.get(ruleNum)) |ruleset| return ruleset;

        var ret = std.ArrayList(Str).init(allocator);

        const rule = self.rules.get(ruleNum).?;
        for (rule) |pipe| {
            // rule_items holds each rule Str a separate slice item, so it is easier to
            // generate the cumulated rule for the line.
            var rule_items = std.ArrayList([]Str).init(allocator);
            defer rule_items.deinit();

            for (pipe) |item| {
                const tmp = self.cache.get(item) orelse try self.genRule(allocator, item);
                try rule_items.append(tmp);
            }

            // pipe_rule holds the rule string for a pipe, so the cumulated rule is
            // the multiple of pipe_rule values.
            var pipe_rule = std.ArrayList(Str).init(allocator);
            defer pipe_rule.deinit();

            for (rule_items.items) |item_rule| {
                if (pipe_rule.items.len == 0) {
                    // if pieces are empty, fill it up with first strings
                    try pipe_rule.ensureTotalCapacity(item_rule.len);
                    pipe_rule.appendSliceAssumeCapacity(item_rule);
                } else {
                    // expand pieces, so we can append the new parts after them
                    for (pipe_rule.items) |prev_rules| {
                        try pipe_rule.appendNTimes(prev_rules, item_rule.len - 1);
                    }
                    for (pipe_rule.items) |_, idx| {
                        pipe_rule.items[idx] = try std.fmt.allocPrint(
                            allocator,
                            "{s}{s}",
                            .{ pipe_rule.items[idx], item_rule[(idx + idx / item_rule.len) % item_rule.len] },
                        );
                    }
                }
            }

            try ret.appendSlice(pipe_rule.items);
        }

        // add result to cache
        try self.cache.put(ruleNum, ret.items);

        return ret.items;
    }
};

fn parseInput(allocator: std.mem.Allocator, input: Str) !Monster {
    var m = Monster{
        .rules = std.AutoHashMap(RuleNums, Rule).init(allocator),
        .cache = std.AutoHashMap(RuleNums, []Str).init(allocator),
        .messages = std.ArrayList(Str).init(allocator),
    };

    var lines = std.mem.split(u8, input, "\n");
    lines: while (lines.next()) |line| {
        if (line.len == 0) break; // continue with message lines

        var rules = std.ArrayList([]RuleNums).init(allocator);

        var line_rules = std.mem.tokenize(u8, line, ":");
        const rule_no = try std.fmt.parseUnsigned(RuleNums, line_rules.next().?, 10);

        var pipes = std.mem.tokenize(u8, line_rules.next().?, "|");
        var pipe_num: usize = 0;
        while (pipes.next()) |pipe| : (pipe_num += 1) {
            var rule = std.ArrayList(RuleNums).init(allocator);

            var nums = std.mem.tokenize(u8, pipe, " ");
            var nums_count: usize = 0;
            while (nums.next()) |num| : (nums_count += 1) {
                if (std.fmt.parseUnsigned(RuleNums, num, 10)) |number| {
                    try rule.append(number);
                } else |err| switch (err) {
                    // handle "a" and "b" by presetting cache values for them
                    error.InvalidCharacter => {
                        const char = std.mem.trim(u8, num, " \"")[0];
                        const tmp = try allocator.alloc(Str, 1);
                        std.mem.set(Str, tmp, try std.fmt.allocPrint(allocator, "{c}", .{char}));
                        try m.cache.put(rule_no, tmp);
                        continue :lines;
                    },
                    else => return err,
                }
            }
            try rules.append(try rule.toOwnedSlice());
        }

        try m.rules.put(rule_no, try rules.toOwnedSlice());
    }

    // parse message lines
    while (lines.next()) |line| {
        try m.messages.append(line);
    }

    return m;
}

test "Test input" {
    var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
    defer arena.deinit();

    var monster = try parseInput(arena.allocator(), test_input);
    defer monster.rules.deinit();

    {
        const exp = &[_]Str{"a"};
        const ret = try monster.genRule(arena.allocator(), 4);
        for (ret) |_, idx| {
            try std.testing.expectEqualStrings(exp[idx], ret[idx]);
        }
    }
    {
        const exp = &[_]Str{"b"};
        const ret = try monster.genRule(arena.allocator(), 5);
        for (ret) |_, idx| {
            try std.testing.expectEqualStrings(exp[idx], ret[idx]);
        }
    }
    {
        const exp = &[_]Str{ "aa", "bb" };
        const ret = try monster.genRule(arena.allocator(), 2);
        for (ret) |_, idx| {
            try std.testing.expectEqualStrings(exp[idx], ret[idx]);
        }
    }
    {
        const exp = &[_]Str{ "ab", "ba" };
        const ret = try monster.genRule(arena.allocator(), 3);
        for (ret) |_, idx| {
            try std.testing.expectEqualStrings(exp[idx], ret[idx]);
        }
    }
    {
        const exp = &[_]Str{
            "aaab",
            "bbba",
            "aaba",
            "bbab",
            "abaa",
            "babb",
            "abbb",
            "baaa",
        };
        const ret = try monster.genRule(arena.allocator(), 1);
        for (ret) |_, idx| {
            try std.testing.expectEqualStrings(exp[idx], ret[idx]);
        }
    }

    {
        const exp = &[_]Str{
            "aaaabb",
            "abbbab",
            "aaabab",
            "abbabb",
            "aabaab",
            "ababbb",
            "aabbbb",
            "abaaab",
        };
        const ret = try monster.genRule(arena.allocator(), 0);
        for (ret) |_, idx| {
            try std.testing.expectEqualStrings(exp[idx], ret[idx]);
        }
    }
}

test "Real input" { // 16: "b", 26: "a"
    var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
    defer arena.deinit();

    var monster = try parseInput(arena.allocator(), @embedFile(PATH));
    defer monster.rules.deinit();

    {
        const exp = &[_]Str{ "ba", "aa" };
        const ret = try monster.genRule(arena.allocator(), 39);
        for (ret) |_, idx| {
            try std.testing.expectEqualStrings(exp[idx], ret[idx]);
        }
    }
    {
        const exp = &[_]Str{ "bba", "baa", "aba" };
        // 19: 16 39 | 26 132
        const ret = try monster.genRule(arena.allocator(), 19);
        for (ret) |_, idx| {
            try std.testing.expectEqualStrings(exp[idx], ret[idx]);
        }
    }
    {
        const exp = &[_]Str{ "abb", "baa" };
        // 30: 26 125 | 16 112
        const ret = try monster.genRule(arena.allocator(), 30);
        for (ret) |_, idx| {
            try std.testing.expectEqualStrings(exp[idx], ret[idx]);
        }
    }
    {
        const exp = &[_]Str{ "bbab", "baab", "abab", "abba", "baaa" };
        // 56: 19 16 | 30 26
        const ret = try monster.genRule(arena.allocator(), 56);
        for (ret) |_, idx| {
            try std.testing.expectEqualStrings(exp[idx], ret[idx]);
        }
    }
}

const test_input =
    \\0: 4 1 5
    \\1: 2 3 | 3 2
    \\2: 4 4 | 5 5
    \\3: 4 5 | 5 4
    \\4: "a"
    \\5: "b"
    \\
    \\ababbb
    \\bababa
    \\abbbab
    \\aaabbb
    \\aaaabbb
;

const test_input_part2 =
    \\42: 9 14 | 10 1
    \\9: 14 27 | 1 26
    \\10: 23 14 | 28 1
    \\1: "a"
    \\11: 42 31
    \\5: 1 14 | 15 1
    \\19: 14 1 | 14 14
    \\12: 24 14 | 19 1
    \\16: 15 1 | 14 14
    \\31: 14 17 | 1 13
    \\6: 14 14 | 1 14
    \\2: 1 24 | 14 4
    \\0: 8 11
    \\13: 14 3 | 1 12
    \\15: 1 | 14
    \\17: 14 2 | 1 7
    \\23: 25 1 | 22 14
    \\28: 16 1
    \\4: 1 1
    \\20: 14 14 | 1 15
    \\3: 5 14 | 16 1
    \\27: 1 6 | 14 18
    \\14: "b"
    \\21: 14 1 | 1 14
    \\25: 1 1 | 1 14
    \\22: 14 14
    \\8: 42
    \\26: 14 22 | 1 20
    \\18: 15 15
    \\7: 14 5 | 1 21
    \\24: 14 1
    \\
    \\abbbbbabbbaaaababbaabbbbabababbbabbbbbbabaaaa
    \\bbabbbbaabaabba
    \\babbbbaabbbbbabbbbbbaabaaabaaa
    \\aaabbbbbbaaaabaababaabababbabaaabbababababaaa
    \\bbbbbbbaaaabbbbaaabbabaaa
    \\bbbababbbbaaaaaaaabbababaaababaabab
    \\ababaaaaaabaaab
    \\ababaaaaabbbaba
    \\baabbaaaabbaaaababbaababb
    \\abbbbabbbbaaaababbbbbbaaaababb
    \\aaaaabbaabaaaaababaa
    \\aaaabbaaaabbaaa
    \\aaaabbaabbaaaaaaabbbabbbaaabbaabaaa
    \\babaaabbbaaabaababbaabababaaab
    \\aabbbbbaabbbaaaaaabbbbbababaaaaabbaaabba
;

test "day19a" {
    try std.testing.expectEqual(@as(usize, 208), try first(std.testing.allocator));
}

test "day19b" {
    try std.testing.expectEqual(@as(usize, 316), try second(std.testing.allocator));
}