aoc2024-03-1.zig

// Written by claude-3-5-sonnet-20241022, based on https://github.com/paulsmith/aoc2024/blob/main/03/solution.zig

const std = @import("std");
const startsWith = std.mem.startsWith;
const input = @embedFile("input.txt");

const OpType = enum {
    mul,
    do_op,
    dont_op,
};

const Operation = struct {
    op_type: OpType,
    operand_a: ?i64 = null,
    operand_b: ?i64 = null,
};

const Parser = struct {
    pos: usize = 0,
    input: []const u8,

    fn init(s: []const u8) Parser {
        return .{ .input = s };
    }

    fn parseOperation(self: *Parser) !?Operation {
        // Skip until we find an operation
        while (self.pos < self.input.len) {
            if (try self.matchOperation()) |op| {
                return op;
            }
            self.pos += 1;
        }
        return null;
    }

    fn matchOperation(self: *Parser) !?Operation {
        const remaining = self.input[self.pos..];

        // Match operation type
        if (startsWith(u8, remaining, "mul")) {
            self.pos += "mul".len;
            if (try self.parseOperands()) |operands| {
                return Operation{
                    .op_type = .mul,
                    .operand_a = operands[0],
                    .operand_b = operands[1],
                };
            }
        } else if (startsWith(u8, remaining, "do()")) {
            self.pos += "do()".len;
            return Operation{ .op_type = .do_op };
        } else if (startsWith(u8, remaining, "don't()")) {
            self.pos += "don't()".len;
            return Operation{ .op_type = .dont_op };
        }

        return null;
    }

    fn parseOperands(self: *Parser) !?[2]i64 {
        if (self.pos >= self.input.len or self.input[self.pos] != '(') return null;
        self.pos += 1;

        const opa = try self.parseNumber() orelse return null;

        if (self.pos >= self.input.len or self.input[self.pos] != ',') return null;
        self.pos += 1;

        const opb = try self.parseNumber() orelse return null;

        if (self.pos >= self.input.len or self.input[self.pos] != ')') return null;
        self.pos += 1;

        return [2]i64{ opa, opb };
    }

    fn parseNumber(self: *Parser) !?i64 {
        if (matchDigitRun(self.input[self.pos..])) |len| {
            const num = try std.fmt.parseInt(i64, self.input[self.pos .. self.pos + len], 10);
            self.pos += len;
            return num;
        }
        return null;
    }
};

fn matchDigitRun(s: []const u8) ?u64 {
    var len: u64 = 0;
    while (len < s.len and std.ascii.isDigit(s[len])) len += 1;
    return if (len == 0) null else len;
}

pub fn main() !void {
    try part1();
    try part2();
}

fn part1() !void {
    var parser = Parser.init(input);
    var sum: i64 = 0;

    while (try parser.parseOperation()) |op| {
        if (op.op_type == .mul) {
            sum += op.operand_a.? * op.operand_b.?;
        }
    }

    std.debug.print("{}\n", .{sum});
}

fn part2() !void {
    var parser = Parser.init(input);
    var sum: i64 = 0;
    var enabled = true;

    while (try parser.parseOperation()) |op| {
        switch (op.op_type) {
            .mul => if (enabled) {
                sum += op.operand_a.? * op.operand_b.?;
            },
            .do_op => enabled = true,
            .dont_op => enabled = false,
        }
    }

    std.debug.print("{}\n", .{sum});
}