Make JIT-compiled function names more natural

The JIT now tries to use the source code names of the functions it compiles.

kernel/modelverse_jit/jit.py

@@ -63,17 +63,20 @@ class ModelverseJit(object):
         if body_id in self.todo_entry_points:
             self.todo_entry_points.remove(body_id)
 
-    def generate_function_name(self):
-        """Generates a new function name,"""
-        function_name = 'jit_func%d' % self.jit_count
-        self.jit_count += 1
-        return function_name
+    def generate_function_name(self, suggested_name=None):
+        """Generates a new function name or picks the suggested name if it is still
+           available."""
+        if suggested_name is not None and suggested_name not in self.jit_globals:
+            self.jit_count += 1
+            return suggested_name
+        else:
+            function_name = 'jit_func%d' % self.jit_count
+            self.jit_count += 1
+            return function_name
 
     def register_compiled(self, body_id, compiled_function, function_name=None):
         """Registers a compiled entry point with the JIT."""
-        if function_name is None:
-            function_name = self.generate_function_name()
-
+        function_name = self.generate_function_name(function_name)
         self.jitted_entry_points[body_id] = function_name
         self.jit_globals[function_name] = compiled_function
         if body_id in self.todo_entry_points:
@@ -105,7 +108,7 @@ class ModelverseJit(object):
 
         raise primitive_functions.PrimitiveFinished(self.jitted_parameters[body_id])
 
-    def jit_compile(self, user_root, body_id):
+    def jit_compile(self, user_root, body_id, suggested_name=None):
         """Tries to jit the function defined by the given entry point id and parameter list."""
         # The comment below makes pylint shut up about our (hopefully benign) use of exec here.
         # pylint: disable=I0011,W0122
@@ -121,7 +124,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()
+        function_name = self.generate_function_name(suggested_name)
         self.jitted_entry_points[body_id] = function_name
         self.jit_globals[function_name] = None
 
@@ -175,7 +178,7 @@ class ModelverseJit(object):
         # Extract the compiled function from the JIT global state.
         compiled_function = self.jit_globals[function_name]
 
-        # print(constructed_function)
+        print(constructed_function)
         raise primitive_functions.PrimitiveFinished(compiled_function)
 
 class AnalysisState(object):
@@ -681,7 +684,7 @@ class AnalysisState(object):
         if compiled_func is None:
             # Compile the callee.
             try:
-                gen = self.jit.jit_compile(self.user_root, body_id)
+                gen = self.jit.jit_compile(self.user_root, body_id, callee_name)
                 inp = None
                 while True:
                     inp = yield gen.send(inp)