Optimize calls in tail position to tail call instructions

kernel/modelverse_jit/tree_ir.py

@@ -306,6 +306,30 @@ class ReturnInstruction(VoidInstruction):
         value, = new_children
         return ReturnInstruction(value)
 
+    def simplify_node(self):
+        """Applies basic simplification to this instruction only."""
+        # If we have a return whose value is a call, then we can rewrite
+        # that as a tail call and save us a stack frame.
+        def rewrite_call(instruction):
+            """Rewrites the given instruction in tail-call form, or returns
+               None."""
+            if isinstance(instruction, RunGeneratorFunctionInstruction):
+                return RunTailGeneratorFunctionInstruction(*instruction.get_children())
+            elif isinstance(instruction, CompoundInstruction):
+                snd_rewritten = rewrite_call(instruction.second)
+                if snd_rewritten is not None:
+                    return CompoundInstruction(instruction.first, snd_rewritten)
+
+            return None
+
+        rewritten_value = rewrite_call(self.value)
+        if rewritten_value is None:
+            return self
+        else:
+            # We don't even need to create a return here, because tail calls terminate
+            # the current stack frame anyway.
+            return rewritten_value
+
     def generate_python_def(self, code_generator):
         """Generates Python code for this instruction."""
         if self.value.has_definition():
@@ -691,6 +715,21 @@ class RunGeneratorFunctionInstruction(StateInstruction):
         """Gets this state instruction's argument list."""
         return [self.function, self.argument_dict]
 
+class RunTailGeneratorFunctionInstruction(StateInstruction):
+    """An instruction that runs a generator function."""
+    def __init__(self, function, argument_dict):
+        StateInstruction.__init__(self)
+        self.function = function
+        self.argument_dict = argument_dict
+
+    def get_opcode(self):
+        """Gets the opcode for this state instruction."""
+        return "TAIL_CALL_KWARGS"
+
+    def get_arguments(self):
+        """Gets this state instruction's argument list."""
+        return [self.function, self.argument_dict]
+
 class VariableName(object):
     """A data structure that unifies names across instructions that access the
        same variable."""