Combine state definitions during codegen

kernel/modelverse_jit/tree_ir.py

@@ -105,14 +105,18 @@ class Instruction(object):
 class PythonGenerator(object):
     """Generates Python code from instructions."""
 
-    def __init__(self):
+    def __init__(self, combine_state_definitions=True):
         self.code = []
+        self.state_definitions = []
+        self.state_definition_names = set()
         self.indentation_string = ' ' * 4
         self.indentation = 0
-        self.result_value_dict = {}
+        self.result_name_dict = {}
+        self.combine_state_definitions = combine_state_definitions
 
     def append(self, text):
         """Appends the given string to this code generator."""
+        self.flush_state_definitions()
         self.code.append(text)
 
     def append_indentation(self):
@@ -129,25 +133,34 @@ class PythonGenerator(object):
 
     def increase_indentation(self):
         """Increases the code generator's indentation by one indent."""
+        self.flush_state_definitions()
         self.indentation += 1
 
     def decrease_indentation(self):
         """Decreases the code generator's indentation by one indent."""
+        self.flush_state_definitions()
         self.indentation -= 1
 
     def get_result_name(self, instruction, advised_name=None):
         """Gets the name of the given instruction's result variable."""
-        if instruction not in self.result_value_dict:
+        if instruction not in self.result_name_dict:
             override_name = instruction.get_result_name_override(self)
             if override_name is not None:
-                self.result_value_dict[instruction] = override_name
+                self.result_name_dict[instruction] = override_name
             elif advised_name is not None:
-                self.result_value_dict[instruction] = advised_name
+                self.result_name_dict[instruction] = advised_name
             else:
-                self.result_value_dict[instruction] = \
-                    'tmp' + str(len(self.result_value_dict))
+                self.result_name_dict[instruction] = \
+                    'tmp' + str(len(self.result_name_dict))
+
+        result = self.result_name_dict[instruction]
+        if result in self.state_definition_names:
+            # This might mean that we're trying to access a name that is
+            # defined by the current block of state definitions. So we need
+            # to flush the state definitions now.
+            self.flush_state_definitions()
 
-        return self.result_value_dict[instruction]
+        return result
 
     def append_definition(self, lhs, rhs):
         """Defines the first instruction's result variable as the second
@@ -173,13 +186,34 @@ class PythonGenerator(object):
 
     def append_state_definition(self, lhs, opcode, args):
         """Appends a definition that queries the modelverse state."""
-        self.append_line(
-            "%s, = yield [('%s', [%s])]" % (
-                self.get_result_name(lhs),
-                opcode,
-                ', '.join([arg_i.generate_python_use(self) for arg_i in args])))
+        result_name = self.get_result_name(lhs)
+        request_tuple = '(%s, [%s])' % (
+            repr(opcode),
+            ', '.join([arg_i.generate_python_use(self) for arg_i in args]))
+        self.state_definitions.append((result_name, request_tuple))
+        self.state_definition_names.add(result_name)
+        if not self.combine_state_definitions:
+            self.flush_state_definitions()
+
+    def flush_state_definitions(self):
+        """Flushes the list of state-access definitions to the generated code."""
+        state_defs = self.state_definitions
+        if len(state_defs) > 0:
+            # Clear state_definitions _before_ calling append_line, because append_line
+            # will call flush_state_definitions. Clearing state_definitions afterward
+            # will result in infinite recursion.
+            self.state_definitions = []
+            self.state_definition_names = set()
+            if len(state_defs) == 1:
+                self.append_line('%s, = yield [%s]' % state_defs[0])
+            else:
+                self.append_line(
+                    "%s = yield [%s]" % (
+                        ', '.join([name for name, _ in state_defs]),
+                        ', '.join([def_val for _, def_val in state_defs])))
 
     def __str__(self):
+        self.flush_state_definitions()
         return ''.join(self.code)
 
 class VoidInstruction(Instruction):