advent-of-code-2019/intcode.py

from collections import defaultdict
from enum import Enum
from inspect import signature
import operator

class IntcodeError(Exception):
    pass

class OpState(Enum):
    """Execution state class."""
    CONTINUE = 0
    HALT     = 1
    WAIT     = 2

class defaultlist(list):
    """Default list class. Allows writing and reading out of bounds."""
    def __init__(self, lst, val_factory):
        super().__init__(lst)
        self.val_factory = val_factory

    def __setitem__(self, i, x):
        for _ in range((i - len(self) + 1)):
            self.append(self.val_factory())
        super().__setitem__(i, x)

    def __getitem__(self, i):
        if i >= len(self):
            return self.val_factory() 
        return super().__getitem__(i)

class Singleton(object):
    def __init__(self, x=None):
        self.x = x

    def append(self, y):
        self.x = y

    def pop(self, _):
        x = self.x
        if x == None:
            raise IndexError("pop from empty singleton")
        self.x = None
        return x 

    def __eq__(self, other):
        return self.x == other

    def __str__(self):
        return str(self.x)

    def format(self):
        if self.x == None:
            return None
        return int(self.x)

    def copy(self):
        return self

def makeIO(in_buff_class, out_buff_class):
    class IO(object):
        def __init__(self, in_buff = in_buff_class()):
            self.in_buff  = in_buff
            self.out_buff = out_buff_class()

        def pop_input(self):
            if self.in_buff == in_buff_class(): 
                return None
            return self.in_buff.pop(0)

        def append_output(self, x):
            self.out_buff.append(x)

        def copy(self):
            new = IO()
            new.in_buff  = self.in_buff.copy()
            new.out_buff = self.out_buff.copy()
            return new

        def flush(self):
            o = self.out_buff
            self.out_buff = out_buff_class()
            try:
                return o.format()
            except AttributeError:
                return o

        def write(self, input):
            if self.in_buff != in_buff_class():
                raise IntcodeError("writing to nonempty input")
            self.in_buff = in_buff_class(input) 

    return IO

StackIO = makeIO(list, list)
SingletonIO = makeIO(Singleton, Singleton)

class Opcode:
    def of_operator(operator):
        """Make opcode from operator."""
        def opcode(state, p1, p2, p3):
            r = operator(state.memory[p1], state.memory[p2])
            state.memory[p3] = r
            state.i += 4
            return OpState.CONTINUE
        return opcode 

    def jump_when(flag):
        def opcode(state, p1, p2):
            if (state.memory[p1] > 0) == flag:
                state.i = state.memory[p2]
            else:
                state.i += 3
            return OpState.CONTINUE
        return opcode 

    def adjust_base(state, p1):
        state.rel_base += state.memory[p1]
        state.i += 2
        return OpState.CONTINUE

    def get_input(state, p1):
        x = state.IO.pop_input()
        if x == None:
            return OpState.WAIT
        state.memory[p1] = x
        state.i += 2
        return OpState.CONTINUE

    def put_output(state, p1):
        state.IO.append_output(state.memory[p1])
        state.i += 2
        return OpState.CONTINUE

    halt = lambda _ : OpState.HALT

    ops = {1  : of_operator(operator.add),
           2  : of_operator(operator.mul),
           3  : get_input,
           4  : put_output,
           5  : jump_when(True), 
           6  : jump_when(False),
           7  : of_operator(operator.lt),
           8  : of_operator(operator.eq),
           9  : adjust_base,
           99 : halt
          }

    def parse(op):
        op = str(op)
        par_modes, op = op[:-2][::-1], int(op[-2:])

        opcode = Opcode.ops[op]
        parnum = len(signature(opcode).parameters) - 1
        par_modes = par_modes + '0'*(parnum - len(par_modes))

        return opcode, map(int, par_modes)

    def run(state):
        opcode, par_modes = Opcode.parse(state.memory[state.i])
        pars = []
        for pn, mode in enumerate(par_modes, start=1):
            p = state.i + pn
            if mode == 0:
                p = state.memory[p]
            elif mode == 2:
                p = state.rel_base + state.memory[p]
            pars.append(p)
        op_state = opcode(state, *pars)
        return op_state



class Interpreter(object):
    def __init__(self, program, IO_class=StackIO, i=0, rel_base=0):
        self.memory = defaultlist(program[:], val_factory = lambda : 0)
        self.IO = IO_class()
        self.i = i
        self.rel_base = rel_base

    def __iter__(self):
        return self
  
    def __next__(self):
        if Opcode.run(self) == OpState.HALT:
            raise StopIteration
        while Opcode.run(self) == OpState.CONTINUE:
            continue
        return self.IO.flush()

    def write(self, in_buff):
        self.IO.write(in_buff) 

    def run(self, in_buff=None):
        if in_buff != None:
            self.write(in_buff)
        out = next(self)
        if Opcode.run(self) != OpState.HALT:
            raise IntcodeError("expecting input")
        return out

    def copy(self):
        memory = self.memory.copy()
        IO = IO.copy()
        return Emulator(memory, IO)


class Display(defaultdict):
    def __init__(self, encoding):
        super(Display, self).__init__(int)
        self.encoding = encoding

    def __str__(self):
        xs, ys = zip(*self.keys())
        X, Y = max(xs)+1, max(ys)+1
        return '\n'.join(''.join(self.encoding[self[(x, y)]]
                         for x in range(X))
                            for y in range(Y))


class DisplayIO(object):
    def __init__(self, display):
        self.display = display
        self.block  = tuple()

    def input(self):
        return NotImplemented

    def wait_for_input(self):
        return NotImplemented

    def output(self, x):
        self.block += (x,)
        if len(self.block) == 3:
            self.display[self.block[:2]] = self.block[2]
            self.block = tuple()