Refactored with Interpreter.pipe_from

day07.py

@@ -1,49 +1,52 @@
 #!/usr/bin/env python3
 
-from intcode import Interpreter, makeIO, Singleton, Halted
+from intcode import Interpreter, SingletonIO 
 from itertools import permutations
+from lib import memoize
 
 stack_size = 5
 fst_amp_input = 0
 
-IO = makeIO(list, Singleton)
 
 def preproc(puzzle_input):
     program = list(map(int, puzzle_input.split(',')))
-    amp = lambda : Interpreter(program, IO)
-    return amp
+    def make_amp(phase):
+        amp = Interpreter(program, SingletonIO)
+        amp.eval(phase)
+        return amp
+    return make_amp
 
 
 
-def partI(amp):
+def partI(make_amp):
+    @memoize
+    def run_amp(phase, amp_in):
+        amp = make_amp(phase)
+        return amp.run(amp_in)
+
     phase_range = range(stack_size)
     best = 0 
 
     for perm in permutations(phase_range):
         amp_in = fst_amp_input
         for phase in perm:
-            amp_in = amp().run([phase, amp_in])
+            amp_in = run_amp(phase, amp_in)
         best = max(amp_in, best)
     return best
    
-def partII(amp):
+def partII(make_amp):
     phase_range = range(5, 10)
     best = 0 
 
     for perm in permutations(phase_range):
-        amps = [amp() for _ in range(stack_size)]
-        for p, a in zip(perm, amps):
-            a.eval([p])
-
-        amp_in = fst_amp_input
-        while True:
-            try:
-                for a in amps:
-                    amp_in = a.eval([amp_in])
-            except Halted:
-                break
+        amps = list(map(make_amp, perm))
+        for prev, amp in zip(amps, amps[1:]):
+            amp.pipe_from(prev)
+        
+        amp.write(fst_amp_input)
+        for amp_in in amp:
+            amp.write(amp_in)
         best = max(amp_in, best)
-    
     return best
 
 from main import Tests