Map global names to body ids in the JIT instead of the kernel

kernel/modelverse_jit/jit.py

@@ -53,13 +53,19 @@ class ModelverseJit(object):
     def __init__(self, max_instructions=None, compiled_function_lookup=None):
         self.todo_entry_points = set()
         self.no_jit_entry_points = set()
-        self.jitted_entry_points = {}
         self.jitted_parameters = {}
         self.jit_globals = {
             'PrimitiveFinished' : primitive_functions.PrimitiveFinished,
             jit_runtime.CALL_FUNCTION_NAME : jit_runtime.call_function,
             jit_runtime.GET_INPUT_FUNCTION_NAME : jit_runtime.get_input
         }
+        # jitted_entry_points maps body ids to values in jit_globals.
+        self.jitted_entry_points = {}
+        # global_functions maps global value names to body ids.
+        self.global_functions = {}
+        # global_functions_inv maps body ids to global value names.
+        self.global_functions_inv = {}
+        # bytecode_graphs maps body ids to their parsed bytecode graphs.
         self.bytecode_graphs = {}
         self.jit_count = 0
         self.max_instructions = max_instructions
@@ -156,15 +162,39 @@ class ModelverseJit(object):
             self.jit_count += 1
             return function_name
 
-    def generate_function_name(self, suggested_name=None):
+    def generate_function_name(self, body_id, suggested_name=None):
         """Generates a new function name or picks the suggested name if it is still
            available."""
+        if suggested_name is None:
+            suggested_name = self.get_global_name(body_id)
+
         return self.generate_name('func', suggested_name)
 
+    def register_global(self, body_id, global_name):
+        """Associates the given body id with the given global name."""
+        self.global_functions[global_name] = body_id
+        self.global_functions_inv[body_id] = global_name
+
+    def get_global_name(self, body_id):
+        """Gets the name of the global function with the given body id.
+           Returns None if no known global exists with the given id."""
+        if body_id in self.global_functions_inv:
+            return self.global_functions_inv[body_id]
+        else:
+            return None
+
+    def get_global_body_id(self, global_name):
+        """Gets the body id of the global function with the given name.
+           Returns None if no known global exists with the given name."""
+        if global_name in self.global_functions:
+            return self.global_functions[global_name]
+        else:
+            return None
+
     def register_compiled(self, body_id, compiled_function, function_name=None):
         """Registers a compiled entry point with the JIT."""
         # Get the function's name.
-        function_name = self.generate_function_name(function_name)
+        function_name = self.generate_function_name(body_id, function_name)
         # Map the body id to the given parameter list.
         self.jitted_entry_points[body_id] = function_name
         self.jit_globals[function_name] = compiled_function
@@ -280,7 +310,7 @@ class ModelverseJit(object):
 
         # Generate a name for the function we're about to analyze, and pretend that
         # it already exists. (we need to do this for recursive functions)
-        function_name = self.generate_function_name(suggested_name)
+        function_name = self.generate_function_name(body_id, suggested_name)
         self.jitted_entry_points[body_id] = function_name
         self.jit_globals[function_name] = None
 

kernel/modelverse_kernel/main.py

@@ -31,10 +31,9 @@ class ModelverseKernel(object):
         self.allow_compiled = True
         #self.allow_compiled = False
 
-        # suggested_function_names maps body ids to suggested function names.
-        # You can tell the kernel to stop caring about getting the function names
-        # right by setting this to `None`.
-        self.suggested_function_names = {}
+        # Set `self.suggest_function_names` to True to associate global function names
+        # with their function bodies.
+        self.suggest_function_names = True
 
         # `self.jit` handles most JIT-related functionality.
         self.jit = jit.ModelverseJit()
@@ -182,11 +181,7 @@ class ModelverseKernel(object):
 
     def jit_compile(self, user_root, inst):
         # Try to retrieve the suggested name.
-        if self.suggested_function_names is not None and inst in self.suggested_function_names:
-            suggested_name = self.suggested_function_names[inst]
-        else:
-            suggested_name = None
-
+        suggested_name = self.jit.get_global_name(inst)
         # Have the JIT compile the function.
         return self.jit.jit_compile(user_root, inst, suggested_name)
 
@@ -528,18 +523,18 @@ class ModelverseKernel(object):
                     # We have a compiled function ready!
                     # Now we have to bind the ID to the compiled functions
                     # For this, we read out the body of the resolved data
-                    compiler_val, =  yield [("RD", [variable, "value"])]
+                    compiler_val, = yield [("RD", [variable, "value"])]
                     compiler_body, = yield [("RD", [compiler_val, "body"])]
                     self.jit.register_compiled(compiler_body, compiled_function, var_name)
 
             # If we're dealing with a function, then we might want to figure out what its body id
             # is now so we can suggest a name to the JIT later.
-            if self.suggested_function_names is not None:
-                compiler_val, =  yield [("RD", [variable, "value"])]
+            if self.suggest_function_names and self.jit.get_global_body_id(var_name) is None:
+                compiler_val, = yield [("RD", [variable, "value"])]
                 if compiler_val is not None:
                     compiler_body, = yield [("RD", [compiler_val, "body"])]
                     if compiler_body is not None:
-                        self.suggested_function_names[compiler_body] = var_name
+                        self.jit.register_global(compiler_body, var_name)
 
         else:
             phase_link, returnvalue_link, new_phase = \