Optimize CFG flow before lowering

kernel/modelverse_jit/cfg_ir.py

@@ -515,3 +515,27 @@ class Binary(Value):
 def create_jump(block, arguments=None):
     """Creates a jump to the given block with the given argument list."""
     return JumpFlow(Branch(block, arguments))
+
+def apply_to_value(function, def_or_value):
+    """Applies the given function to the specified value, or the underlying value of the
+       given definition."""
+    if isinstance(def_or_value, Definition):
+        return apply_to_value(function, def_or_value.value)
+    else:
+        return function(def_or_value)
+
+def is_literal(value):
+    """Tests if the given value is a literal."""
+    return isinstance(value, Literal)
+
+def is_literal_def(def_or_value):
+    """Tests if the given value is a literal or a definition with an underlying literal."""
+    return apply_to_value(is_literal, def_or_value)
+
+def get_literal_value(value):
+    """Gets the value of the given literal value."""
+    return value.literal
+
+def get_literal_def_value(def_or_value):
+    """Gets the value of the given literal value or definition with an underlying literal."""
+    return apply_to_value(get_literal_value, def_or_value)

kernel/modelverse_jit/cfg_optimization.py

@@ -31,3 +31,40 @@ def get_all_reachable_blocks(entry_point):
         if block not in results:
             results[block] = get_reachable_blocks(block)
     return results
+
+def is_empty_block(block):
+    """Tests if the given block contains no parameters or definitions."""
+    return len(block.parameters) == 0 and len(block.definitions) == 0
+
+def optimize_flow(block):
+    """Optimizes the given block's flow instruction."""
+    changed = True
+    while changed:
+        changed = False
+
+        # Select flow with a literal condition can be optimized to a direct jump.
+        if (isinstance(block.flow, cfg_ir.SelectFlow)
+                and cfg_ir.is_literal_def(block.flow.condition)):
+            literal = cfg_ir.get_literal_def_value(block.flow.condition)
+            block.flow = cfg_ir.JumpFlow(
+                block.flow.if_branch if literal else block.flow.else_branch)
+            changed = True
+
+        # Jumps to blocks which contain no parameters or definitions can be replaced
+        # by the target block's flow.
+        if (isinstance(block.flow, cfg_ir.JumpFlow)
+                and is_empty_block(block.flow.branch.block)
+                and block.flow.branch.block is not block):
+            block.flow = block.flow.branch.block.flow
+            changed = True
+
+def optimize_graph_flow(entry_point):
+    """Optimizes all flow instructions in the graph defined by the given entry point."""
+    all_blocks = set(get_reachable_blocks(entry_point))
+    all_blocks.add(entry_point)
+    for block in all_blocks:
+        optimize_flow(block)
+
+def optimize(entry_point):
+    """Optimizes the control-flow graph defined by the given entry point."""
+    optimize_graph_flow(entry_point)

kernel/modelverse_jit/jit.py

@@ -440,6 +440,7 @@ class ModelverseJit(object):
         if self.jit_code_log_function is not None:
             bytecode_analyzer = bytecode_to_cfg.AnalysisState()
             bytecode_analyzer.analyze(body_bytecode)
+            cfg_optimization.optimize(bytecode_analyzer.entry_point)
             self.jit_code_log_function(
                 "CFG for function '%s' at '%d':\n%s" % (
                     function_name, body_id,