Refactor input, output tree construction

kernel/modelverse_jit/bytecode_to_tree.py

@@ -31,6 +31,20 @@ def apply_intrinsic(intrinsic_function, named_args):
             tree_ir.create_block(*store_instructions),
             intrinsic_function(**arg_value_dict))
 
+def retrieve_task_root():
+    """Creates an instruction that stores the task_root variable in a local."""
+    return tree_ir.StoreLocalInstruction(
+        'task_root',
+        tree_ir.LoadIndexInstruction(
+            tree_ir.LoadLocalInstruction(jit_runtime.KWARGS_PARAMETER_NAME),
+            tree_ir.LiteralInstruction('task_root')))
+
+def load_kernel():
+    """Creates an instruction that loads the Modelverse kernel."""
+    return tree_ir.LoadIndexInstruction(
+        tree_ir.LoadLocalInstruction(jit_runtime.KWARGS_PARAMETER_NAME),
+        tree_ir.LiteralInstruction('mvk'))
+
 def create_access(pointer):
     """Creates a tree that loads the given pointer's value."""
     # Accessing a variable is pretty easy. It really just boils
@@ -70,6 +84,138 @@ def create_assign(pointer, value):
         tree_ir.DeleteEdgeInstruction(
             value_link.create_load()))
 
+def create_input(use_input_function=False):
+    """Creates an instruction that pops a value from the input queue."""
+    # Possible alternative to the explicit syntax tree: just call the jit_runtime.__get_input
+    # function.
+    if use_input_function:
+        return tree_ir.create_jit_call(
+            tree_ir.LoadGlobalInstruction(jit_runtime.GET_INPUT_FUNCTION_NAME),
+            [],
+            tree_ir.LoadLocalInstruction(jit_runtime.KWARGS_PARAMETER_NAME))
+
+    # The plan is to generate this tree:
+    #
+    #     value = None
+    #     while True:
+    #         _input = yield [("RD", [task_root, "input"])]
+    #         value = yield [("RD", [_input, "value"])]
+    #
+    #         if value is None:
+    #             kwargs['mvk'].success = False # to avoid blocking
+    #             yield None # nop/interrupt
+    #         else:
+    #             break
+    #
+    #     _next = yield [("RD", [_input, "next"])]
+    #     yield [("CD", [task_root, "input", _next])]
+    #     yield [("CE", [jit_locals, value])]
+    #     yield [("DN", [_input])]
+
+    task_root = retrieve_task_root()
+    _input = tree_ir.StoreLocalInstruction(
+        None,
+        tree_ir.ReadDictionaryValueInstruction(
+            task_root.create_load(),
+            tree_ir.LiteralInstruction('input')))
+
+    value = tree_ir.StoreLocalInstruction(
+        None,
+        tree_ir.ReadDictionaryValueInstruction(
+            _input.create_load(),
+            tree_ir.LiteralInstruction('value')))
+
+    raise primitive_functions.PrimitiveFinished(
+        tree_ir.CompoundInstruction(
+            tree_ir.create_block(
+                task_root,
+                value.create_store(tree_ir.LiteralInstruction(None)),
+                tree_ir.LoopInstruction(
+                    tree_ir.create_block(
+                        _input,
+                        value,
+                        tree_ir.SelectInstruction(
+                            tree_ir.BinaryInstruction(
+                                value.create_load(),
+                                'is',
+                                tree_ir.LiteralInstruction(None)),
+                            tree_ir.create_block(
+                                tree_ir.StoreMemberInstruction(
+                                    load_kernel(),
+                                    'success',
+                                    tree_ir.LiteralInstruction(False)),
+                                tree_ir.NopInstruction()),
+                            tree_ir.BreakInstruction()))),
+                tree_ir.CreateDictionaryEdgeInstruction(
+                    task_root.create_load(),
+                    tree_ir.LiteralInstruction('input'),
+                    tree_ir.ReadDictionaryValueInstruction(
+                        _input.create_load(),
+                        tree_ir.LiteralInstruction('next'))),
+                tree_ir.CreateEdgeInstruction(
+                    tree_ir.LoadLocalInstruction(jit_runtime.LOCALS_NODE_NAME),
+                    value.create_load()),
+                tree_ir.DeleteNodeInstruction(_input.create_load())),
+            value.create_load()))
+
+def create_output(output_value):
+    """Creates an output instruction that outputs the given value."""
+    # The plan is to basically generate this tree:
+    #
+    # value = <some tree>
+    # last_output, last_output_link, new_last_output = \
+    #                 yield [("RD", [task_root, "last_output"]),
+    #                        ("RDE", [task_root, "last_output"]),
+    #                        ("CN", []),
+    #                       ]
+    # _, _, _, _ = \
+    #                 yield [("CD", [last_output, "value", value]),
+    #                        ("CD", [last_output, "next", new_last_output]),
+    #                        ("CD", [task_root, "last_output", new_last_output]),
+    #                        ("DE", [last_output_link])
+    #                       ]
+    # yield None
+
+    value_local = tree_ir.StoreLocalInstruction('value', output_value)
+
+    store_task_root = retrieve_task_root()
+    last_output = tree_ir.StoreLocalInstruction(
+        'last_output',
+        tree_ir.ReadDictionaryValueInstruction(
+            store_task_root.create_load(),
+            tree_ir.LiteralInstruction('last_output')))
+
+    last_output_link = tree_ir.StoreLocalInstruction(
+        'last_output_link',
+        tree_ir.ReadDictionaryEdgeInstruction(
+            store_task_root.create_load(),
+            tree_ir.LiteralInstruction('last_output')))
+
+    new_last_output = tree_ir.StoreLocalInstruction(
+        'new_last_output',
+        tree_ir.CreateNodeInstruction())
+
+    return tree_ir.create_block(
+        value_local,
+        store_task_root,
+        last_output,
+        last_output_link,
+        new_last_output,
+        tree_ir.CreateDictionaryEdgeInstruction(
+            last_output.create_load(),
+            tree_ir.LiteralInstruction('value'),
+            value_local.create_load()),
+        tree_ir.CreateDictionaryEdgeInstruction(
+            last_output.create_load(),
+            tree_ir.LiteralInstruction('next'),
+            new_last_output.create_load()),
+        tree_ir.CreateDictionaryEdgeInstruction(
+            store_task_root.create_load(),
+            tree_ir.LiteralInstruction('last_output'),
+            new_last_output.create_load()),
+        tree_ir.DeleteEdgeInstruction(last_output_link.create_load()),
+        tree_ir.NopInstruction())
+
 class LocalNameMap(object):
     """A map that converts local variable nodes to identifiers."""
     def __init__(self, local_mapping=None):
@@ -111,21 +257,6 @@ class AnalysisState(object):
                 "Local is used as target of function call.")
         self.function_vars.add(local_id)
 
-    def retrieve_task_root(self):
-        """Creates an instruction that stores the task_root variable
-           in a local."""
-        return tree_ir.StoreLocalInstruction(
-            'task_root',
-            tree_ir.LoadIndexInstruction(
-                tree_ir.LoadLocalInstruction(jit_runtime.KWARGS_PARAMETER_NAME),
-                tree_ir.LiteralInstruction('task_root')))
-
-    def load_kernel(self):
-        """Creates an instruction that loads the Modelverse kernel."""
-        return tree_ir.LoadIndexInstruction(
-            tree_ir.LoadLocalInstruction(jit_runtime.KWARGS_PARAMETER_NAME),
-            tree_ir.LiteralInstruction('mvk'))
-
     def analyze(self, instruction):
         """Tries to build an intermediate representation from the instruction with the
         given id."""
@@ -244,139 +375,12 @@ class AnalysisState(object):
 
     def analyze_output(self, instruction):
         """Tries to analyze the given 'output' instruction."""
-        # The plan is to basically generate this tree:
-        #
-        # value = <some tree>
-        # last_output, last_output_link, new_last_output = \
-        #                 yield [("RD", [task_root, "last_output"]),
-        #                        ("RDE", [task_root, "last_output"]),
-        #                        ("CN", []),
-        #                       ]
-        # _, _, _, _ = \
-        #                 yield [("CD", [last_output, "value", value]),
-        #                        ("CD", [last_output, "next", new_last_output]),
-        #                        ("CD", [task_root, "last_output", new_last_output]),
-        #                        ("DE", [last_output_link])
-        #                       ]
-        # yield None
-
         value_val, = yield [("CALL_ARGS", [self.analyze, (instruction.value,)])]
-        value_local = tree_ir.StoreLocalInstruction('value', value_val)
-
-        store_task_root = self.retrieve_task_root()
-        last_output = tree_ir.StoreLocalInstruction(
-            'last_output',
-            tree_ir.ReadDictionaryValueInstruction(
-                store_task_root.create_load(),
-                tree_ir.LiteralInstruction('last_output')))
-
-        last_output_link = tree_ir.StoreLocalInstruction(
-            'last_output_link',
-            tree_ir.ReadDictionaryEdgeInstruction(
-                store_task_root.create_load(),
-                tree_ir.LiteralInstruction('last_output')))
-
-        new_last_output = tree_ir.StoreLocalInstruction(
-            'new_last_output',
-            tree_ir.CreateNodeInstruction())
-
-        result = tree_ir.create_block(
-            value_local,
-            store_task_root,
-            last_output,
-            last_output_link,
-            new_last_output,
-            tree_ir.CreateDictionaryEdgeInstruction(
-                last_output.create_load(),
-                tree_ir.LiteralInstruction('value'),
-                value_local.create_load()),
-            tree_ir.CreateDictionaryEdgeInstruction(
-                last_output.create_load(),
-                tree_ir.LiteralInstruction('next'),
-                new_last_output.create_load()),
-            tree_ir.CreateDictionaryEdgeInstruction(
-                store_task_root.create_load(),
-                tree_ir.LiteralInstruction('last_output'),
-                new_last_output.create_load()),
-            tree_ir.DeleteEdgeInstruction(last_output_link.create_load()),
-            tree_ir.NopInstruction())
-
-        raise primitive_functions.PrimitiveFinished(result)
+        raise primitive_functions.PrimitiveFinished(create_output(value_val))
 
     def analyze_input(self, _):
         """Tries to analyze the given 'input' instruction."""
-
-        # Possible alternative to the explicit syntax tree:
-        if self.jit.input_function_enabled:
-            raise primitive_functions.PrimitiveFinished(
-                tree_ir.create_jit_call(
-                    tree_ir.LoadGlobalInstruction(jit_runtime.GET_INPUT_FUNCTION_NAME),
-                    [],
-                    tree_ir.LoadLocalInstruction(jit_runtime.KWARGS_PARAMETER_NAME)))
-
-        # The plan is to generate this tree:
-        #
-        #     value = None
-        #     while True:
-        #         _input = yield [("RD", [task_root, "input"])]
-        #         value = yield [("RD", [_input, "value"])]
-        #
-        #         if value is None:
-        #             kwargs['mvk'].success = False # to avoid blocking
-        #             yield None # nop/interrupt
-        #         else:
-        #             break
-        #
-        #     _next = yield [("RD", [_input, "next"])]
-        #     yield [("CD", [task_root, "input", _next])]
-        #     yield [("CE", [jit_locals, value])]
-        #     yield [("DN", [_input])]
-
-        task_root = self.retrieve_task_root()
-        _input = tree_ir.StoreLocalInstruction(
-            None,
-            tree_ir.ReadDictionaryValueInstruction(
-                task_root.create_load(),
-                tree_ir.LiteralInstruction('input')))
-
-        value = tree_ir.StoreLocalInstruction(
-            None,
-            tree_ir.ReadDictionaryValueInstruction(
-                _input.create_load(),
-                tree_ir.LiteralInstruction('value')))
-
-        raise primitive_functions.PrimitiveFinished(
-            tree_ir.CompoundInstruction(
-                tree_ir.create_block(
-                    task_root,
-                    value.create_store(tree_ir.LiteralInstruction(None)),
-                    tree_ir.LoopInstruction(
-                        tree_ir.create_block(
-                            _input,
-                            value,
-                            tree_ir.SelectInstruction(
-                                tree_ir.BinaryInstruction(
-                                    value.create_load(),
-                                    'is',
-                                    tree_ir.LiteralInstruction(None)),
-                                tree_ir.create_block(
-                                    tree_ir.StoreMemberInstruction(
-                                        self.load_kernel(),
-                                        'success',
-                                        tree_ir.LiteralInstruction(False)),
-                                    tree_ir.NopInstruction()),
-                                tree_ir.BreakInstruction()))),
-                    tree_ir.CreateDictionaryEdgeInstruction(
-                        task_root.create_load(),
-                        tree_ir.LiteralInstruction('input'),
-                        tree_ir.ReadDictionaryValueInstruction(
-                            _input.create_load(),
-                            tree_ir.LiteralInstruction('next'))),
-                    tree_ir.CreateEdgeInstruction(
-                        tree_ir.LoadLocalInstruction(jit_runtime.LOCALS_NODE_NAME),
-                        value.create_load()),
-                    tree_ir.DeleteNodeInstruction(_input.create_load())),
-                value.create_load()))
+        raise primitive_functions.PrimitiveFinished(create_input(self.jit.input_function_enabled))
 
     def analyze_resolve(self, instruction):
         """Tries to analyze the given 'resolve' instruction."""
@@ -400,7 +404,7 @@ class AnalysisState(object):
             raise primitive_functions.PrimitiveFinished(
                 tree_ir.LoadLocalInstruction(name))
 
-        task_root = self.retrieve_task_root()
+        task_root = retrieve_task_root()
         global_var = tree_ir.StoreLocalInstruction(
             'global_var',
             tree_ir.ReadDictionaryValueInstruction(
@@ -467,7 +471,7 @@ class AnalysisState(object):
         #
         #     tmp = global_var
 
-        task_root = self.retrieve_task_root()
+        task_root = retrieve_task_root()
         _globals = tree_ir.StoreLocalInstruction(
             '_globals',
             tree_ir.ReadDictionaryValueInstruction(