Improve codegen slightly

kernel/modelverse_jit/jit.py

@@ -76,7 +76,7 @@ class ModelverseJit(object):
         constructed_function = tree_ir.DefineFunctionInstruction(
             'jit_func%d' % self.jit_count,
             parameter_list + ['**' + KWARGS_PARAMETER_NAME],
-            constructed_body)
+            constructed_body.simplify())
         self.jit_count += 1
         # Convert the function definition to Python code, and compile it.
         exec(str(constructed_function), self.jit_globals)

kernel/modelverse_jit/tree_ir.py

@@ -182,14 +182,17 @@ class SelectInstruction(Instruction):
         code_generator.append_line(
             'if ' + self.condition.generate_python_use(code_generator) + ':')
         code_generator.increase_indentation()
-        self.if_clause.generate_python_def(code_generator)
+        if self.if_clause.has_definition() or (not if_has_result):
+            self.if_clause.generate_python_def(code_generator)
         if if_has_result:
             code_generator.append_definition(self, self.if_clause)
         code_generator.decrease_indentation()
-        if self.else_clause.has_definition() or if_has_result:
+        else_has_def = self.else_clause.has_definition()
+        if else_has_def or if_has_result:
             code_generator.append_line('else:')
             code_generator.increase_indentation()
-            self.else_clause.generate_python_def(code_generator)
+            if else_has_def:
+                self.else_clause.generate_python_def(code_generator)
             if if_has_result:
                 code_generator.append_definition(self, self.else_clause)
             code_generator.decrease_indentation()
@@ -639,8 +642,8 @@ def create_block(*statements):
         return statements[0]
     else:
         return CompoundInstruction(
-            statements[0],
-            create_block(*statements[1:]))
+            create_block(*statements[:-1]),
+            statements[-1])
 
 if __name__ == "__main__":
     example_tree = SelectInstruction(