[pypy-svn] r29834 - in pypy/dist/pypy/lib: . test2

[mwh at codespeak.net]

Author: mwh
Date: Sat Jul  8 16:09:01 2006
New Revision: 29834

Added:
   pypy/dist/pypy/lib/optimizers.py   (contents, props changed)
   pypy/dist/pypy/lib/test2/test_optimizers.py   (contents, props changed)
Log:
(misto, mwh)
A first cut at a constant folding optimization step.
Not enabled by default (it's rather slow)
import parser, optimizers
parser.install_compiler_hook(optimizers.hook)
will enable.


Added: pypy/dist/pypy/lib/optimizers.py
==============================================================================
--- (empty file)
+++ pypy/dist/pypy/lib/optimizers.py	Sat Jul  8 16:09:01 2006
@@ -0,0 +1,69 @@
+import parser, operator
+
+def binaryVisit(operation):
+    def visit(self, node):
+        left = node.left
+        right = node.right
+        if isinstance(left, parser.ASTConst) and \
+               isinstance(right, parser.ASTConst):
+            if type(left.value) == type(right.value):
+                return parser.ASTConst(operation(left.value, right.value))
+        return node
+    return visit
+
+def bitopVisit(astclass, operation):
+    def compress(values):
+        while len(values) > 1:
+            values[0] = operation(values[0], values[1])
+            del values[1]
+        return values[0]
+    def visit(self, node):
+        values = []
+        for i, n in enumerate(node.nodes):
+            if not isinstance(n, parser.ASTConst):
+                if values:
+                    return astclass([compress(values)] + node.nodes[i:])
+                else:
+                    return node
+            values.append(n.value)
+        return parser.ASTConst(compress(values))
+    return visit
+        
+
+class Folder:
+    def __init__(self):
+        pass
+        
+    def defaultvisit(self, node):
+        return node
+
+    def __getattr__(self, attrname):
+        if attrname.startswith('visit'):
+            return self.defaultvisit
+        raise AttributeError(attrname)
+
+    visitAdd = binaryVisit(operator.add)
+    visitSub = binaryVisit(operator.sub)
+    visitMul = binaryVisit(operator.mul)
+    visitDiv = binaryVisit(operator.div)
+
+    visitBitand = bitopVisit(parser.ASTBitand, operator.and_)
+    visitBitor  = bitopVisit(parser.ASTBitor, operator.or_)
+    visitBitxor = bitopVisit(parser.ASTBitxor, operator.xor)
+
+    def visitTuple(self, node):
+        contents = []
+        for n in node.nodes:
+            if not isinstance(n, parser.ASTConst):
+                return node
+            contents.append(n.value)
+        return parser.ASTConst(tuple(contents))
+
+    def visitDiscard(self, node):
+        if isinstance(node, parser.ASTConst):
+            return None
+        else:
+            return node
+
+def hook(ast, enc):
+    return ast.mutate(Folder())

Added: pypy/dist/pypy/lib/test2/test_optimizers.py
==============================================================================
--- (empty file)
+++ pypy/dist/pypy/lib/test2/test_optimizers.py	Sat Jul  8 16:09:01 2006
@@ -0,0 +1,57 @@
+
+def expr(src):
+    from parser import source2ast
+    module = source2ast(src)
+    return module.node.nodes[0].expr
+
+class AppTestFolder:
+    def setup_class(cls):
+        cls.w_expr, cls.w_check_const_fold = cls.space.unpackiterable(cls.space.appexec([], '''():
+        from parser import source2ast, ASTConst, ASTNode
+        from optimizers import Folder
+
+        def expr(cls, src):
+            module = source2ast(src)
+            expr = module.node.nodes[0].expr
+            assert isinstance(expr, cls)
+            return expr
+
+        def check_const_fold(src, value, cls=ASTNode):
+            exprnode = expr(cls, src)
+            exprnode = exprnode.mutate(Folder())
+            assert isinstance(exprnode, ASTConst) and exprnode.value == value
+
+        return expr, check_const_fold'''))
+
+    def test_expr(self):
+        from parser import ASTConst
+        raises(AssertionError, self.expr, ASTConst, "a")
+
+    def test_bitfold(self):
+        from parser import ASTBitand, ASTBitor, ASTBitxor
+
+        self.check_const_fold('1&3', 1, ASTBitand)
+        self.check_const_fold('1|3', 3, ASTBitor)
+        self.check_const_fold('1^3', 2, ASTBitxor)
+
+        self.check_const_fold('1&3&5', 1, ASTBitand)
+        self.check_const_fold('1|3|5', 7, ASTBitor)
+        self.check_const_fold('1^3^5', 7, ASTBitxor)
+
+    def test_binaryfold(self):
+        from parser import ASTAdd, ASTSub, ASTMul, ASTDiv
+
+        self.check_const_fold('1+3', 4, ASTAdd)
+        self.check_const_fold('"1"+"3"', '13', ASTAdd)
+        self.check_const_fold('1+3+6', 10, ASTAdd)
+
+        self.check_const_fold('1+3-6', -2)
+
+    def test_constant_tuples(self):
+        from parser import ASTTuple
+
+        self.check_const_fold('(1+1, 2)', (2, 2), ASTTuple)
+        self.check_const_fold('((1,) + (1, 2),)', ((1, 1, 2),), ASTTuple)
+
+        
+