Make scheduling code iterative instead of recursive

kernel/modelverse_jit/cfg_to_tree.py

@@ -501,7 +501,7 @@ class DefinitionScheduler(object):
         else:
             return tree_ir.EmptyInstruction()
 
-    def __schedule_definition(self, definition, store_and_forget=False):
+    def __schedule_definition(self, definition, store_and_forget):
         """Schedules the given definition, excluding definitions that are dependent on it.
            A list of trees is returned."""
         assert not self.has_scheduled(definition)
@@ -535,27 +535,42 @@ class DefinitionScheduler(object):
 
         return statements
 
+    def __schedule_top_of_stack(self, stack, statements, store_and_forget):
+        """Tries to schedule the given stack's top-of-stack element."""
+        definition = stack[-1]
+        if self.has_scheduled(definition):
+            # Definition has already been scheduled. Load it if it is the only
+            # element on the stack. Otherwise, do nothing.
+            stack.pop()
+            if (not store_and_forget
+                    and definition.has_value()
+                    and not cfg_ir.is_value_def(definition, LoweringState.inline_value_types)):
+                statements.append(self.lowering_state.create_definition_load(definition))
+            return statements
+
+        dependents = self.get_dependents(definition)
+        if len(dependents) > 0:
+            dependent = next(iter(dependents))
+            stack.append(dependent)
+            return statements
+        else:
+            # Definition has not been scheduled, and all of its dependent definitions
+            # have been scheduled.
+            stack.pop()
+            return self.__schedule_definition(definition, store_and_forget) + statements
+
     def schedule(self, definition, store_and_forget=False):
         """Schedules the given definition. The resulting tree is returned."""
         assert not self.has_scheduled(definition)
 
         statements = []
-        dependents = self.get_dependents(definition)
-        while not self.has_scheduled(definition) and len(dependents) > 0:
-            dependent = next(iter(dependents))
-            dependent_tree = self.schedule(dependent, True)
-            statements.append(dependent_tree)
+        schedule_stack = [definition]
 
-        statements.reverse()
-        if not self.has_scheduled(definition):
-            statements = self.__schedule_definition(definition, store_and_forget) + statements
-        elif (not store_and_forget
-              and definition.has_value()
-              and not cfg_ir.is_value_def(definition, LoweringState.inline_value_types)):
-            statements.append(self.lowering_state.create_definition_load(definition))
-
-        result = tree_ir.create_block(*statements)
-        return result
+        while len(schedule_stack) > 0:
+            statements = self.__schedule_top_of_stack(
+                schedule_stack, statements, store_and_forget or len(schedule_stack) > 1)
+
+        return tree_ir.create_block(*statements)
 
 class LoweringState(object):
     """Stores information related to the relooper->tree conversion."""