Refacto day 5 intcode computer

2019/day05-intcode.py

@@ -1,59 +1,79 @@
-ZERO_PARAM_CODES = [99]
-ONE_PARAM_CODES = [3, 4]
-THREE_PARAMS_CODES = [1, 2]
+NUMBER_OF_PARAMS_MAP = {
+    1: 3,
+    2: 3,
+    3: 1,
+    4: 1,
+    99: 0,
+}
+
+LAST_IS_RESULT_MAP = {
+    1: True,
+    2: True,
+    3: True,
+    4: False,
+}
+
+
+def get_value(program, args, param_modes, index):
+    if param_modes[index] == 0:
+        return program[args[index]]
+    return args[index]
+
+
+def parse_args(program, raw_args, param_modes, last_is_result=False):
+    args = []
+    limit = -1 if last_is_result else None
+    for i, arg in enumerate(raw_args[:limit]):
+        args.append(get_value(program, raw_args, param_modes, i))
+    if last_is_result:
+        args.append(raw_args[-1])
+    return args
 
 
 def compute(lst):
-    result = lst.copy()
+    program = lst.copy()
 
-    i = 0
-    while i < len(result):
-        first_instruction = str(result[i])
-        code = int(first_instruction[-2:])
-        param_modes = first_instruction[:-2]
+    pointer = 0
+    while pointer < len(program):
+        pointer_moved = False
+        instruction = str(program[pointer])
+        code = int(instruction[-2:])
         if code == 99:
             print("Halting due to code 99")
-            return result
-        if code in THREE_PARAMS_CODES:
-            offset = 4
-            param_modes = param_modes.zfill(offset - 1)
-            param_modes_parsed = list(map(int, reversed(param_modes)))
+            return program
 
-            position1, position2 = result[i + 1], result[i + 2]
-            if param_modes_parsed[0] == 0:
-                item1 = result[position1]
-            else:
-                item1 = position1
-            if param_modes_parsed[1] == 0:
-                item2 = result[position2]
-            else:
-                item2 = position2
-            result_index = result[i + 3]
+        number_of_params = NUMBER_OF_PARAMS_MAP[code]
+        offset = number_of_params + 1
+        param_modes = instruction[:-2]
+        param_modes = param_modes.zfill(number_of_params)
+        param_modes = list(map(int, reversed(param_modes)))
+        raw_params = []
+        for i in range(1, offset):
+            raw_params.append(program[pointer + i])
 
-            if code == 1:
-                result[result_index] = item1 + item2
-            elif code == 2:
-                result[result_index] = item1 * item2
+        last_is_result = LAST_IS_RESULT_MAP[code]
+        params = parse_args(program, raw_params, param_modes, last_is_result)
 
-        elif code in ONE_PARAM_CODES:
-            offset = 2
-            param_mode = int(param_modes.zfill(offset - 1))
-            position = result[i + 1]
-            if code == 3:
-                result[position] = int(input(f"Input for instruction {i}"))
-            elif code == 4:
-                if param_mode == 0:
-                    res = result[position]
-                else:
-                    res = position
-                print(res)
+        if code == 1:
+            # Addition
+            program[params[2]] = params[0] + params[1]
+        elif code == 2:
+            # Multiplication
+            program[params[2]] = params[0] * params[1]
+        elif code == 3:
+            # Input
+            program[params[0]] = int(input(f"Input for instruction {pointer}"))
+        elif code == 4:
+            # Output
+            print(params[0])
 
         else:
             raise ValueError(f"Something bad happened, code={code}")
 
-        i += offset
+        if not pointer_moved:
+            pointer += offset
 
-    return result
+    return program
 
 
 if __name__ == "__main__":