Add deep rewrite for mutual recursion support

opshin/compiler.py

@@ -39,6 +39,12 @@
 
 from .compiler_config import DEFAULT_CONFIG
 from .optimize.optimize_const_folding import OptimizeConstantFolding
+from .rewrite.rewrite_expanded_union_calls import (
+    RewriteExpandedUnionCalls,
+)
+from .rewrite.rewrite_function_closures import (
+    RewriteFunctionClosures,
+)
 from .optimize.optimize_remove_deadconstants import OptimizeRemoveDeadConstants
 from .optimize.optimize_remove_deadconds import OptimizeRemoveDeadConditions
 from .optimize.optimize_fold_if_fallthrough import OptimizeFoldIfFallthrough
@@ -1234,6 +1240,8 @@ def compile(
         RewriteAnnotateFallthrough(),
         # The type inference needs to be run after complex python operations were rewritten
         AggressiveTypeInferencer(config.allow_isinstance_anything),
+        (RewriteExpandedUnionCalls() if config.expand_union_types else NoOp()),
+        RewriteFunctionClosures(),
         # Rewrites that circumvent the type inference or use its results
         OptimizeFoldBoolCast(),
         RewriteAssertNone(),

opshin/optimize/optimize_fold_if_fallthrough.py

@@ -1,28 +1,27 @@
 from ast import *
 from copy import copy
 
-from ..util import CompilingNodeTransformer
+from ..typed_util import (
+    ScopedSequenceNodeTransformer,
+    annotate_compound_statement_fallthrough,
+)
 
 """
 If exactly one branch of an if-statement can fall through, fold the following
 statements in the enclosing sequence into that branch.
 """
 
 
-def sequence_can_fall_through(statements):
-    for stmt in statements:
-        if not getattr(stmt, "can_fall_through", True):
-            return False
-    return True
-
-
-class OptimizeFoldIfFallthrough(CompilingNodeTransformer):
+class OptimizeFoldIfFallthrough(ScopedSequenceNodeTransformer):
     step = "Folding trailing statements into sole fallthrough if-branches"
 
     def fold_sequence(self, statements):
         folded = []
         i = 0
         while i < len(statements):
+            if statements[i] is None:
+                i += 1
+                continue
             stmt = self.visit(statements[i])
             if stmt is None:
                 i += 1
@@ -38,68 +37,45 @@ def fold_sequence(self, statements):
                         stmt.body = self.fold_sequence(stmt.body + trailing)
                     else:
                         stmt.orelse = self.fold_sequence(stmt.orelse + trailing)
-                    stmt.body_can_fall_through = sequence_can_fall_through(stmt.body)
-                    stmt.orelse_can_fall_through = sequence_can_fall_through(
-                        stmt.orelse
-                    )
-                    stmt.can_fall_through = (
-                        stmt.body_can_fall_through or stmt.orelse_can_fall_through
-                    )
-                    folded.append(stmt)
+                    folded.append(annotate_compound_statement_fallthrough(stmt))
                     break
             folded.append(stmt)
             i += 1
         return folded
 
     def visit_Module(self, node: Module) -> Module:
-        node_cp = copy(node)
-        node_cp.body = self.fold_sequence(node.body)
-        node_cp.can_fall_through = sequence_can_fall_through(node_cp.body)
-        return node_cp
+        node_cp = super().visit_Module(node)
+        node_cp.body = self.fold_sequence(node_cp.body)
+        return annotate_compound_statement_fallthrough(node_cp)
 
     def visit_FunctionDef(self, node: FunctionDef) -> FunctionDef:
-        node_cp = copy(node)
-        node_cp.body = self.fold_sequence(node.body)
-        node_cp.body_can_fall_through = sequence_can_fall_through(node_cp.body)
-        node_cp.can_fall_through = True
-        return node_cp
+        node_cp = super().visit_FunctionDef(node)
+        node_cp.body = self.fold_sequence(node_cp.body)
+        return annotate_compound_statement_fallthrough(node_cp)
 
     def visit_ClassDef(self, node: ClassDef) -> ClassDef:
-        node_cp = copy(node)
-        node_cp.body = self.fold_sequence(node.body)
-        node_cp.body_can_fall_through = sequence_can_fall_through(node_cp.body)
-        node_cp.can_fall_through = True
-        return node_cp
+        node_cp = super().visit_ClassDef(node)
+        node_cp.body = self.fold_sequence(node_cp.body)
+        return annotate_compound_statement_fallthrough(node_cp)
 
     def visit_If(self, node: If) -> If:
         node_cp = copy(node)
         node_cp.test = self.visit(node.test)
         node_cp.body = self.fold_sequence(node.body)
         node_cp.orelse = self.fold_sequence(node.orelse)
-        node_cp.body_can_fall_through = sequence_can_fall_through(node_cp.body)
-        node_cp.orelse_can_fall_through = sequence_can_fall_through(node_cp.orelse)
-        node_cp.can_fall_through = (
-            node_cp.body_can_fall_through or node_cp.orelse_can_fall_through
-        )
-        return node_cp
+        return annotate_compound_statement_fallthrough(node_cp)
 
     def visit_While(self, node: While) -> While:
         node_cp = copy(node)
         node_cp.test = self.visit(node.test)
         node_cp.body = self.fold_sequence(node.body)
         node_cp.orelse = self.fold_sequence(node.orelse)
-        node_cp.body_can_fall_through = sequence_can_fall_through(node_cp.body)
-        node_cp.orelse_can_fall_through = sequence_can_fall_through(node_cp.orelse)
-        node_cp.can_fall_through = node_cp.orelse_can_fall_through
-        return node_cp
+        return annotate_compound_statement_fallthrough(node_cp)
 
     def visit_For(self, node: For) -> For:
         node_cp = copy(node)
         node_cp.target = self.visit(node.target)
         node_cp.iter = self.visit(node.iter)
         node_cp.body = self.fold_sequence(node.body)
         node_cp.orelse = self.fold_sequence(node.orelse)
-        node_cp.body_can_fall_through = sequence_can_fall_through(node_cp.body)
-        node_cp.orelse_can_fall_through = sequence_can_fall_through(node_cp.orelse)
-        node_cp.can_fall_through = node_cp.orelse_can_fall_through
-        return node_cp
+        return annotate_compound_statement_fallthrough(node_cp)

opshin/optimize/optimize_remove_deadconds.py

@@ -1,87 +1,53 @@
 from ast import *
-from copy import deepcopy, copy
+from copy import copy
 from typing import Any, Union
 
-from ..util import CompilingNodeTransformer
+from ..typed_util import FlatteningScopedSequenceNodeTransformer
 
 """
 Removes if/while branches that are never executed
 """
 
 
-class OptimizeRemoveDeadConditions(CompilingNodeTransformer):
-    def visit_FunctionDef(self, node: FunctionDef) -> Any:
-        node = copy(node)
-        node.body = self.visit_sequence(node.body)
-        return node
+class OptimizeRemoveDeadConditions(FlatteningScopedSequenceNodeTransformer):
+    def expression_guaranteed_tf(self, expr: expr) -> Union[bool, None]:
+        """
+        Returns True if the expression is guaranteed to be truthy.
+        Returns False if the expression is guaranteed to be falsy.
+        Returns None if it cannot be determined.
 
-    def visit_sequence(self, stmts):
-        new_stmts = []
-        for stmt in stmts:
-            s = self.visit(stmt)
-            if s is None:
-                continue
-            if isinstance(s, list):
-                new_stmts.extend(s)
-            else:
-                new_stmts.append(s)
-        return new_stmts
+        Needs to be run after self.visit has been called on expr.
+        """
+        if isinstance(expr, Constant):
+            return bool(expr.value)
+        return None
 
     def visit_If(self, node: If) -> Any:
         node = copy(node)
         node.test = self.visit(node.test)
         node.body = self.visit_sequence(node.body)
         node.orelse = self.visit_sequence(node.orelse)
-        if isinstance(node.test, Constant):
-            if node.test.value:
-                return node.body
-            else:
-                return node.orelse
+        test_value = self.expression_guaranteed_tf(node.test)
+        if test_value is True:
+            return node.body
+        if test_value is False:
+            return node.orelse
         return node
 
     def visit_While(self, node: While) -> Any:
         node = copy(node)
         node.test = self.visit(node.test)
         node.body = self.visit_sequence(node.body)
         node.orelse = self.visit_sequence(node.orelse)
-        if isinstance(node.test, Constant):
-            if node.test.value:
-                raise ValueError(
-                    "While loop with constant True condition is not allowed (infinite loop)"
-                )
-            else:
-                return node.orelse
-        return node
-
-    def visit_IfExp(self, node: IfExp) -> Any:
-        node = copy(node)
-        node.test = self.visit(node.test)
-        node.body = self.visit(node.body)
-        node.orelse = self.visit(node.orelse)
-
-        # Simplify if the test condition is a constant
-        if isinstance(node.test, Constant):
-            if node.test.value:
-                return node.body
-            else:
-                return node.orelse
-
+        test_value = self.expression_guaranteed_tf(node.test)
+        if test_value is True:
+            raise ValueError(
+                "While loop with constant True condition is not allowed (infinite loop)"
+            )
+        if test_value is False:
+            return node.orelse
         return node
 
-    # Expression simplification logic
-
-    def expression_guaranteed_tf(self, expr: expr) -> Union[bool, None]:
-        """
-        Returns True if the expression is guaranteed to be truthy
-        Returns False if the expression is guaranteed to be falsy
-        Returns None if it cannot be determined
-
-        Needs to be run after self.visit has been called on expr
-        """
-        if isinstance(expr, Constant):
-            return expr.value
-        return None
-
     def visit_IfExp(self, node: IfExp) -> expr:
         ex = copy(node)
         ex.test = self.visit(ex.test)

opshin/optimize/optimize_remove_unreachable.py

@@ -1,62 +1,48 @@
 from ast import *
 from copy import copy
 
-from ..util import CompilingNodeTransformer
+from ..typed_util import ScopedSequenceNodeTransformer
 
 """
 Removes statements that are unreachable because a previous statement in the same
 sequence is known not to fall through.
 """
 
 
-class OptimizeRemoveUnreachable(CompilingNodeTransformer):
+class OptimizeRemoveUnreachable(ScopedSequenceNodeTransformer):
     step = "Removing unreachable statements"
 
-    @staticmethod
-    def visit_sequence(statements, visitor):
+    def visit_sequence(self, statements):
         visited = []
         for stmt in statements:
-            stmt_cp = visitor.visit(stmt)
+            if stmt is None:
+                continue
+            stmt_cp = self.visit(stmt)
             if stmt_cp is None:
                 continue
             visited.append(stmt_cp)
             if not getattr(stmt_cp, "can_fall_through", True):
                 break
         return visited
 
-    def visit_Module(self, node: Module) -> Module:
-        node_cp = copy(node)
-        node_cp.body = self.visit_sequence(node.body, self)
-        return node_cp
-
-    def visit_FunctionDef(self, node: FunctionDef) -> FunctionDef:
-        node_cp = copy(node)
-        node_cp.body = self.visit_sequence(node.body, self)
-        return node_cp
-
-    def visit_ClassDef(self, node: ClassDef) -> ClassDef:
-        node_cp = copy(node)
-        node_cp.body = self.visit_sequence(node.body, self)
-        return node_cp
-
     def visit_If(self, node: If) -> If:
         node_cp = copy(node)
         node_cp.test = self.visit(node.test)
-        node_cp.body = self.visit_sequence(node.body, self)
-        node_cp.orelse = self.visit_sequence(node.orelse, self)
+        node_cp.body = self.visit_sequence(node.body)
+        node_cp.orelse = self.visit_sequence(node.orelse)
         return node_cp
 
     def visit_While(self, node: While) -> While:
         node_cp = copy(node)
         node_cp.test = self.visit(node.test)
-        node_cp.body = self.visit_sequence(node.body, self)
-        node_cp.orelse = self.visit_sequence(node.orelse, self)
+        node_cp.body = self.visit_sequence(node.body)
+        node_cp.orelse = self.visit_sequence(node.orelse)
         return node_cp
 
     def visit_For(self, node: For) -> For:
         node_cp = copy(node)
         node_cp.target = self.visit(node.target)
         node_cp.iter = self.visit(node.iter)
-        node_cp.body = self.visit_sequence(node.body, self)
-        node_cp.orelse = self.visit_sequence(node.orelse, self)
+        node_cp.body = self.visit_sequence(node.body)
+        node_cp.orelse = self.visit_sequence(node.orelse)
         return node_cp