🍺

y2019/intcode.py

@@ -3,11 +3,13 @@ import itertools
 import re
 
 
-def parse(string):
-    ns = [int(n) for n in re.findall(r'-?\d+', string)]
-    memory = collections.defaultdict(int)
-    memory.update(enumerate(ns))
-    return memory
+MODES = {
+    '0': lambda ns, x, relbase: ns[x],
+    '1': lambda ns, x, relbase: x,
+    '2': lambda ns, x, relbase: ns[x + relbase],
+    '0-special': lambda ns, x, relbase: x,
+    '2-special': lambda ns, x, relbase: x + relbase,
+}
 
 
 def get_parameters(ns, pos, modes, N, writes, relbase):
@@ -15,16 +17,10 @@ def get_parameters(ns, pos, modes, N, writes, relbase):
         # paradox: return immediate mode to use positionally outside
         modes[ord(c) - ord('a')] += '-special'
 
-    for mode, x in zip(modes, [ns[y] for y in range(pos, pos + N)]):
-        yield {
-            '0': lambda: ns[x],
-            '1': lambda: x,
-            '2': lambda: ns[x + relbase],
-            '0-special': lambda: x,
-            '2-special': lambda: x + relbase,
-        }[mode]()
-
     yield pos + N
+    for i in range(N):
+        mode, x = modes[i], ns[pos + i]
+        yield MODES[mode](ns, x, relbase)
 
 
 def compute(ns, in_iter):
@@ -43,45 +39,45 @@ def compute(ns, in_iter):
         op = ns[pos] % 100
         # instructions stupidly say ABC referring to parameters 3, 2, 1
         # we do a, b, c
-        modes = list(str(ns[pos] // 100).zfill(3)[::-1])
+        modes = list(str(ns[pos] // 100)[::-1] + '000')
         pos += 1
 
         if op == 1:
-            a, b, c, pos = consume(3, 'c')
+            pos, a, b, c = consume(3, 'c')
             ns[c] = a + b
 
         elif op == 2:
-            a, b, c, pos = consume(3, 'c')
+            pos, a, b, c = consume(3, 'c')
             ns[c] = a * b
 
         elif op == 3:
-            a, pos = consume(1, 'a')
+            pos, a = consume(1, 'a')
             ns[a] = next(in_iter)
 
         elif op == 4:
-            a, pos = consume(1)
+            pos, a = consume(1)
             yield a
 
         elif op == 5:
-            a, b, pos = consume(2)
+            pos, a, b = consume(2)
             if a != 0:
                 pos = b
 
         elif op == 6:
-            a, b, pos = consume(2)
+            pos, a, b = consume(2)
             if a == 0:
                 pos = b
 
         elif op == 7:
-            a, b, c, pos = consume(3, 'c')
+            pos, a, b, c = consume(3, 'c')
             ns[c] = int(a < b)
 
         elif op == 8:
-            a, b, c, pos = consume(3, 'c')
+            pos, a, b, c = consume(3, 'c')
             ns[c] = int(a == b)
 
         elif op == 9:
-            a, pos = consume(1)
+            pos, a = consume(1)
             relbase += a
 
         elif op == 99:
@@ -89,3 +85,10 @@ def compute(ns, in_iter):
 
         else:
             raise RuntimeError(op)
+
+
+def parse(string):
+    ns = [int(n) for n in re.findall(r'-?\d+', string)]
+    memory = collections.defaultdict(int)
+    memory.update(enumerate(ns))
+    return memory