aoc2020 day14

def decode(line):
    lhs, rhs = line.split(' = ')
    if lhs == 'mask':
        mask1 = int(rhs.replace('X', '1'), base=2)
        mask2 = int(rhs.replace('X', '0'), base=2)
        return lhs, (mask1, mask2, rhs)
    elif lhs[:3] == 'mem':
        address = int(lhs[4:-1])
        value = int(rhs)
        return address, value


def addresses(base, mask, offset=0):
    for i, bit in enumerate(reversed(mask)):
        j = i + offset
        if bit == 'X':
            new_address = base ^ (1 << j)
            for address in addresses(new_address, mask[:-(i + 1)], j + 1):
                yield address

    yield base


def solve(lines, callback):
    memory = {}
    current_mask = 0
    for lhs, rhs in map(decode, lines):
        if lhs == 'mask':
            current_mask = rhs
        else:
            callback(memory, lhs, rhs, current_mask)

    return sum(memory.values())


def solve_part1(memory, address, value, mask):
    value &= mask[0]
    value |= mask[1]
    memory[address] = value


def solve_part2(memory, address, value, mask):
    for address in addresses(address | mask[1], mask[2]):
        memory[address] = value


with open('input_day14') as f:
    lines = f.read().strip().split('\n')
    print(solve(lines, solve_part1))
    print(solve(lines, solve_part2))