[Hail] Python common subexpression elimination

hail/python/hail/backend/backend.py

@@ -6,7 +6,7 @@
 from hail.expr.table_type import *
 from hail.expr.matrix_type import *
 from hail.expr.blockmatrix_type import *
-from hail.ir.renderer import Renderer
+from hail.ir.renderer import CSERenderer, Renderer
 from hail.table import Table
 from hail.matrixtable import MatrixTable
 
@@ -99,7 +99,7 @@ def fs(self):
 
     def _to_java_ir(self, ir):
         if not hasattr(ir, '_jir'):
-            r = Renderer(stop_at_jir=True)
+            r = CSERenderer(stop_at_jir=True)
             # FIXME parse should be static
             ir._jir = ir.parse(r(ir), ir_map=r.jirs)
         return ir._jir

hail/python/hail/ir/base_ir.py

@@ -33,7 +33,10 @@ def render_head(self, r):
         head_str = self.head_str()
         if head_str != '':
             head_str = f' {head_str}'
-        return f'({self._ir_name()}{head_str}'
+        trailing_space = ''
+        if len(self.children) > 0:
+            trailing_space = ' '
+        return f'({self._ir_name()}{head_str}{trailing_space}'
 
     def render_tail(self, r):
         return ')'
@@ -89,6 +92,9 @@ def __hash__(self):
     def new_block(self, i: int) -> bool:
         ...
 
+    def renderable_new_block(self, i: int) -> bool:
+        return self.new_block(i)
+
     @staticmethod
     def is_effectful() -> bool:
         return False
@@ -113,9 +119,15 @@ def scan_bindings(self, i: int, default_value=None):
     def uses_agg_context(self, i: int) -> bool:
         return False
 
+    def renderable_uses_agg_context(self, i: int) -> bool:
+        return self.uses_agg_context(i)
+
     def uses_scan_context(self, i: int) -> bool:
         return False
 
+    def renderable_uses_scan_context(self, i: int) -> bool:
+        return self.uses_scan_context(i)
+
     def child_context_without_bindings(self, i: int, parent_context):
         (eval_c, agg_c, scan_c) = parent_context
         if self.uses_agg_context(i):

hail/python/hail/ir/ir.py

@@ -1409,6 +1409,9 @@ def new_block(self, i: int) -> bool:
         n = len(self.constructor_args)
         return self.init_op_args and i < n + len(self.init_op_args)
 
+    def renderable_new_block(self, i: int) -> bool:
+        return i <= 1
+
 
 class ApplyAggOp(BaseApplyAggOp):
     @typecheck_method(agg_op=str,
@@ -1421,6 +1424,9 @@ def __init__(self, agg_op, constructor_args, init_op_args, seq_op_args):
     def uses_agg_context(self, i: int):
         return i >= len(self.constructor_args) + (len(self.init_op_args) if self.init_op_args else 0)
 
+    def renderable_uses_agg_context(self, i: int):
+        return i == 2
+
 
 class ApplyScanOp(BaseApplyAggOp):
     @typecheck_method(agg_op=str,
@@ -1433,6 +1439,9 @@ def __init__(self, agg_op, constructor_args, init_op_args, seq_op_args):
     def uses_scan_context(self, i: int):
         return i >= len(self.constructor_args) + (len(self.init_op_args) if self.init_op_args else 0)
 
+    def renderable_uses_agg_context(self, i: int):
+        return i == 2
+
 
 class Begin(IR):
     @typecheck_method(xs=sequenceof(IR))