advent-of-code-2k19/day5/day5.py

#! /usr/bin/env python3
from collections import namedtuple


def decode_opcode(inp):
    opcode = inp % 100  # 2-digit opcode
    a = (inp // 10000) % 10  # mode of 1st param
    b = (inp // 1000) % 10  # mode of 2nd param
    c = (inp // 100) % 10  # mode of 3rd param
    # ABC  DE
    # ^^^  ^^
    # MODE OPCODE
    return a, b, c, opcode


def decode_parameters(input_prog, ip, modes):
    """
    Returns the adress of the parameters according to the input modes.
    if mode[n] == 1: direct access, else indirect access
    """
    param1 = input_prog[ip + 1] if modes[2] == 0 else (ip + 1)
    param2 = input_prog[ip + 2] if modes[1] == 0 else (ip + 2)
    param3 = input_prog[ip + 3] if modes[0] == 0 else (ip + 3)
    return param1, param2, param3



def interpret_intcode(input_prog, stdin=None):
    # Instruction pointer: index of the next element to be executed
    ip = 0
    while ip < len(input_prog):
        instruction = input_prog[ip]
        *modes, opcode = decode_opcode(instruction)
        
        if opcode == 99:
            break

        parameters = decode_parameters(input_prog, ip, modes)        

        if opcode == 1:  # ADD
            print("ADD", " ".join(str(input_prog[p]) for p in parameters))
            input_prog[parameters[2]] = input_prog[parameters[0]] + input_prog[parameters[1]]
            ip += 4

        elif opcode == 2:  # MUL
            print("MUL", " ".join(str(input_prog[p]) for p in parameters))
            input_prog[parameters[2]] = input_prog[parameters[0]] * input_prog[parameters[1]]
            ip += 4

        elif opcode == 3:  # IN
            input_prog[parameters[1]] = stdin
            ip += 2

        elif opcode == 4:  # OUT
            if input_prog[ip + 1] == 99:
                print("End of execution")
                break

            print("OUT ", input_prog[parameters[1]])
            ip += 2

        elif opcode == 5:  # JMP if true
            if input_prog[parameters[0]] != 0:
                print("JT", " ".join(str(input_prog[p]) for p in parameters))
                ip = input_prog[parameters[1]]
            else:
                ip += 3

        elif opcode == 6:  # JMP if false
            if input_prog[parameters[0]] == 0:
                print("JF", " ".join(str(input_prog[p]) for p in parameters))
                ip = input_prog[parameters[1]]
            else:
                ip += 3

        elif opcode == 7:  # LT
            print("LT", " ".join(str(input_prog[p]) for p in parameters))
            res = int(input_prog[parameters[0]] < input_prog[parameters[1]])
            input_prog[parameters[2]] = res
            ip += 4

        elif opcode == 8:  # EQ
            print("EQ", " ".join(str(input_prog[p]) for p in parameters))
            res = int(input_prog[parameters[1]] == input_prog[parameters[2]])
            input_prog[parameters[2]] = res
            ip += 4


def tests():
    inputs = (
        [1, 0, 0, 0, 99],  # ADD 1 + 1 to location 0
        [2, 3, 0, 3, 99],  # MUL 2 * 3 to location 3
        [2, 4, 4, 5, 99, 0],
        [1, 1, 1, 4, 99, 5, 6, 0, 99],
    )
    expected_outputs = (
        [2, 0, 0, 0, 99],  # 1 + 1 = 2
        [2, 3, 0, 6, 99],  # 3 * 2 = 6
        [2, 4, 4, 5, 99, 9801],  # 99 * 99 = 9801
        [30, 1, 1, 4, 2, 5, 6, 0, 99],
    )
    for i, inp in enumerate(inputs):
        result = inp[:]
        interpret_intcode(result)
        assert (
            result == expected_outputs[i]
        ), f"Expected output for {inp} is {expected_outputs[i]}, but found {result} instead."
    print("\nAll tests passed.\n")


if __name__ == "__main__":
    tests()
    # a = [3, 12, 6, 12, 15, 1, 13, 14, 13, 4, 13, 99, -1, 0, 1, 9]
    # b = [3, 3, 1105, -1, 9, 1101, 0, 0, 12, 4, 12, 99, 1]
    
    # interpret_intcode(b, 0)
    # interpret_intcode(a, 0)

    with open("input.txt") as inp:
        print("Start of input program.")
        memory = [int(x) for x in inp.readline().strip().split(",")]
        interpret_intcode(memory, 5)