Move jit_define_function calls into the jit compilation functions

kernel/modelverse_jit/jit.py

@@ -514,15 +514,12 @@ class ModelverseJit(object):
             raise JitCompilationFailedException(
                 "Function was marked '%s'." % jit_runtime.MUTABLE_FUNCTION_KEY)
 
-        constructed_function, = yield [
+        compiled_function, = yield [
             ("CALL_ARGS", [compile_function_body, (self, function_name, body_id, task_root)])]
 
         yield [("END_TRY", [])]
         del self.compilation_dependencies[body_id]
 
-        # Convert the function definition to Python code, and compile it.
-        compiled_function = self.jit_define_function(function_name, constructed_function)
-
         if self.jit_success_log_function is not None:
             self.jit_success_log_function(
                 "JIT compilation successful: (function '%s' at %d)" % (function_name, body_id))
@@ -546,7 +543,6 @@ class ModelverseJit(object):
            and extracts the resulting function."""
         # The comment below makes pylint shut up about our (hopefully benign) use of exec here.
         # pylint: disable=I0011,W0122
-
         if self.jit_code_log_function is not None:
             self.jit_code_log_function(function_def)
 
@@ -737,10 +733,13 @@ def compile_function_body_baseline(jit, function_name, body_id, task_root, heade
     constructed_body, = yield [("CALL_ARGS", [optimize_tree_ir, (constructed_body,)])]
 
     # Wrap the tree IR in a function definition.
+    constructed_function = create_function(
+        function_name, parameter_list, param_dict,
+        body_param_dict, constructed_body, jit.get_source_map_name(function_name))
+
+    # Convert the function definition to Python code, and compile it.
     raise primitive_functions.PrimitiveFinished(
-        create_function(
-            function_name, parameter_list, param_dict,
-            body_param_dict, constructed_body, jit.get_source_map_name(function_name)))
+        jit.jit_define_function(function_name, constructed_function))
 
 def compile_function_body_fast(jit, function_name, body_id, _):
     """Have the fast JIT compile the function with the given name and body id."""
@@ -763,10 +762,11 @@ def compile_function_body_fast(jit, function_name, body_id, _):
 
     # Optimize the tree that was generated.
     constructed_body, = yield [("CALL_ARGS", [optimize_tree_ir, (constructed_body,)])]
+    constructed_function = create_bare_function(function_name, parameter_list, constructed_body)
+
+    # Convert the function definition to Python code, and compile it.
     raise primitive_functions.PrimitiveFinished(
-        create_bare_function(
-            function_name, parameter_list,
-            constructed_body))
+        jit.jit_define_function(function_name, constructed_function))
 
 def favor_large_functions(body_bytecode):
     """Computes the initial temperature of a function based on the size of