From b65c805891560ae7f84c2411499b40d62c4eab08 Mon Sep 17 00:00:00 2001
From: Bert Peters <bert@bertptrs.nl>
Date: Sun, 24 Jan 2021 10:16:14 +0100
Subject: [PATCH] Implement 2019 day 9

---
 2019/aoc2019/day09.py      | 25 +++++++++++++++++++++++++
 2019/aoc2019/intcode.py    | 38 ++++++++++++++++++++++++++++----------
 2019/tests/test_intcode.py | 15 +++++++++++++++
 3 files changed, 68 insertions(+), 10 deletions(-)
 create mode 100644 2019/aoc2019/day09.py

diff --git a/2019/aoc2019/day09.py b/2019/aoc2019/day09.py
new file mode 100644
index 0000000..8c483f5
--- /dev/null
+++ b/2019/aoc2019/day09.py
@@ -0,0 +1,25 @@
+import sys
+from typing import TextIO
+
+from aoc2019.intcode import read_program, Computer
+
+
+def run_machine(data: TextIO, initializer: int) -> int:
+    program = read_program(data)
+    computer = Computer(program)
+    computer.input.append(initializer)
+
+    computer.run()
+
+    if len(computer.output) > 1:
+        sys.exit(computer.output)
+    else:
+        return computer.output.pop()
+
+
+def part1(data: TextIO) -> int:
+    return run_machine(data, 1)
+
+
+def part2(data: TextIO) -> int:
+    return run_machine(data, 2)
diff --git a/2019/aoc2019/intcode.py b/2019/aoc2019/intcode.py
index 406b28d..d00801d 100644
--- a/2019/aoc2019/intcode.py
+++ b/2019/aoc2019/intcode.py
@@ -11,12 +11,14 @@ def read_program(data: TextIO) -> List[int]:
 class Computer:
     program: List[int]
     pointer: int
+    relative_base: 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.relative_base = 0
         self.input = collections.deque()
         self.output = collections.deque()
 
@@ -30,29 +32,41 @@ class Computer:
     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:
+        if mode == 1:
             return key
+        elif mode == 0:
+            pass  # Nothing to do here, handled below
+        elif mode == 2:
+            key += self.relative_base
         else:
             raise ValueError(f'Unsupported mode "{mode}"')
 
+        self._ensure_length(key + 1)
+        return self.program[key]
+
     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:
+        if mode == 1:
             raise ValueError('Cannot assign to an immediate')
+        elif mode == 0:
+            pass # Nothing to do here, handled below
+        elif mode == 2:
+            key += self.relative_base
         else:
             raise ValueError(f'Unsupported mode "{mode}"')
 
+        self._ensure_length(key + 1)
+        self.program[key] = value
+
     def _ensure_length(self, length: int) -> None:
         if len(self.program) < length:
-            # Double current program size with 0s
-            self.program.extend(0 for _ in range(len(self.program)))
+            if 2 * len(self.program) >= length:
+                # Double current program size with 0s
+                self.program.extend(0 for _ in range(len(self.program)))
+            else:
+                # Resize until the desired length
+                self.program.extend(0 for _ in range(length - len(self.program)))
 
     def run(self) -> None:
         """ Run until failure """
@@ -116,6 +130,10 @@ class Computer:
             else:
                 self[mode[2], 3] = 0
             self.pointer += 4
+        elif opcode == 9:
+            # Adjust relative base
+            self.relative_base += self[mode[0], 1]
+            self.pointer += 2
         elif opcode == 99:
             # Halt
             return False
diff --git a/2019/tests/test_intcode.py b/2019/tests/test_intcode.py
index ea4cbf5..b733951 100644
--- a/2019/tests/test_intcode.py
+++ b/2019/tests/test_intcode.py
@@ -64,3 +64,18 @@ def test_day5_example(inputs: int, expected: int):
     computer.run()
 
     assert computer.output.pop() == expected
+
+
+@pytest.mark.parametrize('program,output', [
+    ([109, 1, 204, -1, 1001, 100, 1, 100, 1008, 100, 16, 101, 1006, 101, 0, 99],
+     [109, 1, 204, -1, 1001, 100, 1, 100, 1008, 100, 16, 101, 1006, 101, 0, 99]),
+    ([1102, 34915192, 34915192, 7, 4, 7, 99, 0], [1219070632396864]),
+    ([104, 1125899906842624, 99], [1125899906842624]),
+])
+def test_instructions_day9(program: List[int], output: List[int]) -> None:
+    computer = Computer(program)
+    computer.run()
+
+    result = list(computer.output)
+
+    assert result == output