Start of AST based optimization pass - support unary ops

DinoV · DinoV · commit 986522d2c18b · 2021-01-28T11:14:38.000-08:00
Summary:
This adds the first new AST based optimization - folding unary ops into a new constant.

This also introduces the usage of the ast.Constant class instead of using dedicated classes (e.g. ast.Num, ast.Str).  All constants end up being of this new type.  So we also get support for emitting a LOAD_CONST for nodes of these types.

Test Plan: ./python -m test.test_compiler
diff --git a/compiler/optimizer.py b/compiler/optimizer.py
@@ -0,0 +1,59 @@
+import ast
+import operator
+from ast import Constant, Num, Str, Bytes, Ellipsis, NameConstant, copy_location
+
+from compiler.visitor import ASTRewriter
+
+
+def is_const(node):
+    return isinstance(node, (Constant, Num, Str, Bytes, Ellipsis, NameConstant))
+
+
+def get_const_value(node):
+    if isinstance(node, (Constant, NameConstant)):
+        return node.value
+    elif isinstance(node, Num):
+        return node.n
+    elif isinstance(node, (Str, Bytes)):
+        return node.s
+    elif isinstance(node, Ellipsis):
+        return ...
+
+    raise TypeError("Bad constant value")
+
+
+UNARY_OPS = {
+    ast.Invert: operator.invert,
+    ast.Not: operator.not_,
+    ast.UAdd: operator.pos,
+    ast.USub: operator.neg,
+}
+INVERSE_OPS = {
+    ast.Is: ast.IsNot,
+    ast.IsNot: ast.Is,
+    ast.In: ast.NotIn,
+    ast.NotIn: ast.In,
+}
+
+
+class AstOptimizer(ASTRewriter):
+    def visitUnaryOp(self, node: ast.UnaryOp):
+        op = self.visit(node.operand)
+        if is_const(op):
+            conv = UNARY_OPS[type(node.op)]
+            val = get_const_value(op)
+            try:
+                return copy_location(Constant(conv(val)), node)
+            except:
+                pass
+        elif (
+            isinstance(node.op, ast.Not)
+            and isinstance(node.operand, ast.Compare)
+            and len(node.operand.ops) == 1
+        ):
+            cmp_op = node.operand.ops[0]
+            new_op = INVERSE_OPS.get(type(cmp_op))
+            if new_op is not None:
+                return self.update_node(node.operand, ops=[new_op()])
+
+        return self.update_node(node, operand=op)
diff --git a/compiler/pycodegen.py b/compiler/pycodegen.py
@@ -90,6 +90,8 @@ def get_bool_const(node):
     if isinstance(node, ast.Name):
         if node.id == "__debug__":
             return not OPTIMIZE
+    if isinstance(node, ast.Constant):
+        return bool(node.value)
 
 
 def is_constant_false(node):
@@ -2059,6 +2061,10 @@ def compileJumpIf(self, test, next, is_if_true):
         self.emit('POP_JUMP_IF_TRUE' if is_if_true else 'POP_JUMP_IF_FALSE', next)
         return True
 
+    def visitConstant(self, node: ast.Constant):
+        self.update_lineno(node)
+        self.emit('LOAD_CONST', node.value)
+
 
 def get_default_generator():
     if sys.version_info >= (3, 7):
diff --git a/test_compiler/test_py37.py b/test_compiler/test_py37.py
@@ -1,6 +1,9 @@
+import ast
 import dis
 from .common import CompilerTest
 from compiler.pycodegen import CodeGenerator, Python37CodeGenerator
+from compiler.optimizer import AstOptimizer
+from compiler.unparse import to_expr
 from compiler.consts import (
     CO_OPTIMIZED,
     CO_NOFREE,
@@ -143,3 +146,21 @@ def test_compile_opt_chained_cmp_op(self):
 
         graph = self.to_graph('assert not a > b > c', CodeGenerator)
         self.assertInGraph(graph, 'UNARY_NOT')
+
+    def test_ast_optimizer(self):
+        cases = [
+            ("+1", "1"),
+            ("--1", "1"),
+            ("~1", "-2"),
+            ("not 1", "False"),
+            ("not x is y", "x is not y"),
+            ("not x is not y", "x is y"),
+            ("not x in y", "x not in y"),
+            ("~1.1", "~1.1"),
+            ("+'str'", "+'str'"),
+        ]
+        for inp, expected in cases:
+            optimizer = AstOptimizer()
+            tree = ast.parse(inp)
+            optimized = to_expr(optimizer.visit(tree).body[0].value)
+            self.assertEqual(expected, optimized, "Input was: " + inp)
