Factor out create_function from jit_recompile

kernel/modelverse_jit/jit.py

@@ -44,6 +44,51 @@ def optimize_tree_ir(instruction):
     """Optimizes an IR tree."""
     return map_and_simplify_generator(expand_constant_read, instruction)
 
+def create_function(
+        function_name, parameter_list, param_dict,
+        body_param_dict, function_body):
+    """Creates a function from the given function name, parameter list,
+       variable-to-parameter name map, variable-to-local name map and
+       function body."""
+    # Write a prologue and prepend it to the generated function body.
+    prologue_statements = []
+    # Create a LOCALS_NODE_NAME node, and connect it to the user root.
+    prologue_statements.append(
+        tree_ir.create_new_local_node(
+            jit_runtime.LOCALS_NODE_NAME,
+            tree_ir.LoadIndexInstruction(
+                tree_ir.LoadLocalInstruction(jit_runtime.KWARGS_PARAMETER_NAME),
+                tree_ir.LiteralInstruction('user_root')),
+            jit_runtime.LOCALS_EDGE_NAME))
+    for (key, val) in param_dict.items():
+        arg_ptr = tree_ir.create_new_local_node(
+            body_param_dict[key],
+            tree_ir.LoadLocalInstruction(jit_runtime.LOCALS_NODE_NAME))
+        prologue_statements.append(arg_ptr)
+        prologue_statements.append(
+            tree_ir.CreateDictionaryEdgeInstruction(
+                tree_ir.LoadLocalInstruction(body_param_dict[key]),
+                tree_ir.LiteralInstruction('value'),
+                tree_ir.LoadLocalInstruction(val)))
+
+    constructed_body = tree_ir.create_block(
+        *(prologue_statements + [function_body]))
+
+    # Optimize the function's body.
+    constructed_body, = yield [("CALL_ARGS", [optimize_tree_ir, (constructed_body,)])]
+
+    # Shield temporaries from the GC.
+    constructed_body = tree_ir.protect_temporaries_from_gc(
+        constructed_body, tree_ir.LoadLocalInstruction(jit_runtime.LOCALS_NODE_NAME))
+
+    # Wrap the IR in a function definition, give it a unique name.
+    constructed_function = tree_ir.DefineFunctionInstruction(
+        function_name,
+        parameter_list + ['**' + jit_runtime.KWARGS_PARAMETER_NAME],
+        constructed_body)
+
+    raise primitive_functions.PrimitiveFinished(constructed_function)
+
 def print_value(val):
     """A thin wrapper around 'print'."""
     print(val)
@@ -356,42 +401,12 @@ class ModelverseJit(object):
         yield [("END_TRY", [])]
         del self.compilation_dependencies[body_id]
 
-        # Write a prologue and prepend it to the generated function body.
-        prologue_statements = []
-        # Create a LOCALS_NODE_NAME node, and connect it to the user root.
-        prologue_statements.append(
-            tree_ir.create_new_local_node(
-                jit_runtime.LOCALS_NODE_NAME,
-                tree_ir.LoadIndexInstruction(
-                    tree_ir.LoadLocalInstruction(jit_runtime.KWARGS_PARAMETER_NAME),
-                    tree_ir.LiteralInstruction('user_root')),
-                jit_runtime.LOCALS_EDGE_NAME))
-        for (key, val) in param_dict.items():
-            arg_ptr = tree_ir.create_new_local_node(
-                body_param_dict[key],
-                tree_ir.LoadLocalInstruction(jit_runtime.LOCALS_NODE_NAME))
-            prologue_statements.append(arg_ptr)
-            prologue_statements.append(
-                tree_ir.CreateDictionaryEdgeInstruction(
-                    tree_ir.LoadLocalInstruction(body_param_dict[key]),
-                    tree_ir.LiteralInstruction('value'),
-                    tree_ir.LoadLocalInstruction(val)))
-
-        constructed_body = tree_ir.create_block(
-            *(prologue_statements + [constructed_body]))
-
-        # Optimize the function's body.
-        constructed_body, = yield [("CALL_ARGS", [optimize_tree_ir, (constructed_body,)])]
-
-        # Shield temporaries from the GC.
-        constructed_body = tree_ir.protect_temporaries_from_gc(
-            constructed_body, tree_ir.LoadLocalInstruction(jit_runtime.LOCALS_NODE_NAME))
-
-        # Wrap the IR in a function definition, give it a unique name.
-        constructed_function = tree_ir.DefineFunctionInstruction(
-            function_name,
-            parameter_list + ['**' + jit_runtime.KWARGS_PARAMETER_NAME],
-            constructed_body)
+        # Wrap the tree IR in a function definition.
+        constructed_function, = yield [
+            ("CALL_ARGS",
+             [create_function,
+              (function_name, parameter_list, param_dict, body_param_dict, constructed_body)])]
+
         # Convert the function definition to Python code, and compile it.
         exec(str(constructed_function), self.jit_globals)
         # Extract the compiled function from the JIT global state.