Implement 2019 day 5 part 1

2019/aoc2019/day05.py

@@ -0,0 +1,29 @@
+from typing import TextIO
+
+from aoc2019.intcode import read_program, Computer
+
+
+def part1(data: TextIO) -> int:
+    program = read_program(data)
+
+    computer = Computer(program)
+
+    # Enter the required starting code
+    computer.input.append(1)
+
+    computer.run()
+
+    return computer.output.pop()
+
+
+def part2(data: TextIO) -> int:
+    program = read_program(data)
+
+    computer = Computer(program)
+
+    # Enter the required starting code
+    computer.input.append(5)
+
+    computer.run()
+
+    return computer.output.pop()

2019/aoc2019/intcode.py

@@ -1,4 +1,5 @@
-from typing import List, TextIO
+import collections
+from typing import List, TextIO, Tuple, Union
 
 
 def read_program(data: TextIO) -> List[int]:
@@ -10,18 +11,46 @@ def read_program(data: TextIO) -> List[int]:
 class Computer:
     program: List[int]
     pointer: int
+    input: collections.deque[int]
+    output: collections.deque[int]
 
     def __init__(self, program: List[int], pointer: int = 0) -> None:
         self.program = program
         self.pointer = pointer
+        self.input = collections.deque()
+        self.output = collections.deque()
 
-    def __getitem__(self, item: int) -> int:
-        self._ensure_length(item + 1)
-        return self.program[item]
+    def _mode_and_key(self, item: Union[int, Tuple[int, int]]) -> Tuple[int, int]:
+        if type(item) == int:
+            mode = 0
+            key = item
+        else:
+            mode, key = item
+            key = self.program[self.pointer + key]
 
-    def __setitem__(self, key: int, value: int) -> None:
+        return mode, key
+
+    def __getitem__(self, item: Union[int, Tuple[int, int]]) -> int:
+        mode, key = self._mode_and_key(item)
+
+        if mode == 0:
+            self._ensure_length(key + 1)
+            return self.program[key]
+        elif mode == 1:
+            return key
+        else:
+            raise ValueError(f'Unsupported mode "{mode}"')
+
+    def __setitem__(self, item: Union[int, Tuple[int, int]], value: int) -> None:
+        mode, key = self._mode_and_key(item)
+
+        if mode == 0:
             self._ensure_length(key + 1)
             self.program[key] = value
+        elif mode == 1:
+            raise ValueError('Cannot assign to an immediate')
+        else:
+            raise ValueError(f'Unsupported mode "{mode}"')
 
     def _ensure_length(self, length: int) -> None:
         if len(self.program) < length:
@@ -29,7 +58,7 @@ class Computer:
             self.program.extend(0 for _ in range(len(self.program)))
 
     def run(self) -> None:
-        """ Run until failure"""
+        """ Run until failure """
         while self._execute_current():
             pass
 
@@ -39,16 +68,30 @@ class Computer:
         :return: True if the program should continue
         """
         pointer = self.pointer
-        opcode = self[pointer]
+        instruction = self[pointer]
+        opcode = instruction % 100
+
+        mode = [
+            (instruction // 100) % 10,
+            (instruction // 1000) % 10,
+            (instruction // 10000) % 10,
+        ]
 
         if opcode == 1:
             # Add
-            self[self[pointer + 3]] = self[self[pointer + 1]] + self[self[pointer + 2]]
+            self[mode[2], 3] = self[mode[0], 1] + self[mode[1], 2]
             self.pointer += 4
         elif opcode == 2:
             # Multiply
-            self[self[pointer + 3]] = self[self[pointer + 1]] * self[self[pointer + 2]]
+            self[mode[2], 3] = self[mode[0], 1] * self[mode[1], 2]
             self.pointer += 4
+        elif opcode == 3:
+            # Input
+            self[mode[0], 1] = self.input.popleft()
+            self.pointer += 2
+        elif opcode == 4:
+            self.output.append(self[mode[0], 1])
+            self.pointer += 2
         elif opcode == 99:
             # Halt
             return False

2019/tests/test_intcode.py

@@ -10,7 +10,9 @@ from aoc2019.intcode import Computer
     ([1, 0, 0, 0, 99], [2, 0, 0, 0, 99]),
     ([2, 3, 0, 3, 99], [2, 3, 0, 6, 99]),
     ([2, 4, 4, 5, 99, 0], [2, 4, 4, 5, 99, 9801]),
-    ([1, 1, 1, 4, 99, 5, 6, 0, 99], [30, 1, 1, 4, 2, 5, 6, 0, 99])
+    ([1, 1, 1, 4, 99, 5, 6, 0, 99], [30, 1, 1, 4, 2, 5, 6, 0, 99]),
+    # This is technically part of day 5 but it fits the pattern
+    ([1002, 4, 3, 4, 33], [1002, 4, 3, 4, 99])
 ])
 def test_instructions_day2(program: List[int], expected: List[int]) -> None:
     computer = Computer(program)