Fix another thunk-related bug

kernel/modelverse_jit/cfg_optimization.py

@@ -5,6 +5,7 @@ import modelverse_jit.cfg_ir as cfg_ir
 import modelverse_jit.cfg_dominators as cfg_dominators
 import modelverse_jit.cfg_ssa_construction as cfg_ssa_construction
 import modelverse_jit.cfg_data_structures as cfg_data_structures
+import modelverse_jit.tree_ir as tree_ir
 import modelverse_kernel.primitives as primitive_functions
 
 def is_empty_block(block):
@@ -250,7 +251,7 @@ def try_redefine_as_direct_call(definition, jit, called_globals):
                 called_globals.add(loaded_ptr)
     elif isinstance(target, cfg_ir.Literal):
         node_id = target.literal
-        thunk_name = jit.jit_thunk_constant(node_id)
+        thunk_name = jit.jit_thunk_constant_function(node_id)
         definition.redefine(
             cfg_ir.DirectFunctionCall(
                 thunk_name, call.argument_list, cfg_ir.JIT_CALLING_CONVENTION))

kernel/modelverse_jit/jit.py

@@ -106,7 +106,7 @@ class ModelverseJit(object):
             'PrimitiveFinished' : primitive_functions.PrimitiveFinished,
             jit_runtime.CALL_FUNCTION_NAME : jit_runtime.call_function,
             jit_runtime.GET_INPUT_FUNCTION_NAME : jit_runtime.get_input,
-            jit_runtime.JIT_THUNK_CONSTANT_FUNCTION_NAME : self.jit_thunk_constant,
+            jit_runtime.JIT_THUNK_CONSTANT_FUNCTION_NAME : self.jit_thunk_constant_function,
             jit_runtime.JIT_THUNK_GLOBAL_FUNCTION_NAME : self.jit_thunk_global,
             jit_runtime.UNREACHABLE_FUNCTION_NAME : jit_runtime.unreachable
         }
@@ -582,8 +582,8 @@ class ModelverseJit(object):
         self.jit_globals[thunk_name] = __jit_thunk
         return thunk_name
 
-    def jit_thunk_constant(self, body_id):
-        """Creates a thunk from given body id.
+    def jit_thunk_constant_body(self, body_id):
+        """Creates a thunk from the given body id.
            This thunk is a function that will invoke the function whose body id is given.
            The thunk's name in the JIT's global context is returned."""
         self.lookup_compiled_body(body_id)
@@ -597,6 +597,15 @@ class ModelverseJit(object):
             # Looks like we'll just have to build that thunk after all.
             return self.jit_thunk(tree_ir.LiteralInstruction(body_id))
 
+    def jit_thunk_constant_function(self, body_id):
+        """Creates a thunk from the given function id.
+           This thunk is a function that will invoke the function whose function id is given.
+           The thunk's name in the JIT's global context is returned."""
+        return self.jit_thunk(
+            tree_ir.ReadDictionaryValueInstruction(
+                tree_ir.LiteralInstruction(body_id),
+                tree_ir.LiteralInstruction(jit_runtime.FUNCTION_BODY_KEY)))
+
     def jit_thunk_global(self, global_name):
         """Creates a thunk from given global name.
            This thunk is a function that will invoke the function whose body id is given.

kernel/modelverse_jit/runtime.py

@@ -19,8 +19,8 @@ CALL_FUNCTION_NAME = "__call_function"
 GET_INPUT_FUNCTION_NAME = "__get_input"
 """The name of the '__get_input' function, in the jitted function scope."""
 
-JIT_THUNK_CONSTANT_FUNCTION_NAME = "__jit_thunk_constant"
-"""The name of the jit_thunk_constant function in the JIT's global context."""
+JIT_THUNK_CONSTANT_FUNCTION_NAME = "__jit_thunk_constant_function"
+"""The name of the jit_thunk_constant_function function in the JIT's global context."""
 
 JIT_THUNK_GLOBAL_FUNCTION_NAME = "__jit_thunk_global"
 """The name of the jit_thunk_global function in the JIT's global context."""