[Python-checkins] r62458 - in python/branches/tlee-ast-optimize: Lib/test/test_optimizer.py Python/optimize.c

[thomas.lee]

Author: thomas.lee
Date: Tue Apr 22 14:36:34 2008
New Revision: 62458

Log:
Added a small optimization I missed in the peephole optimizer ('return None' -> 'return'). Port most of the old peephole binop tests to the AST optimizer.

Modified:
   python/branches/tlee-ast-optimize/Lib/test/test_optimizer.py
   python/branches/tlee-ast-optimize/Python/optimize.c

Modified: python/branches/tlee-ast-optimize/Lib/test/test_optimizer.py
==============================================================================
--- python/branches/tlee-ast-optimize/Lib/test/test_optimizer.py	(original)
+++ python/branches/tlee-ast-optimize/Lib/test/test_optimizer.py	Tue Apr 22 14:36:34 2008
@@ -73,12 +73,34 @@
         return self.assertAstNode(_ast.Str, 's', expected, code)
 
     def test_binop_fold_num(self):
-        # check binary constant folding for numeric values
-        self.assertNum(3, "1 + 2")
-        self.assertNum(18, "2 * 8 + 4 - 2")
-        self.assertNum(16, "2 ** 4")
-        self.assertNum(4, "1 << 2")
-        self.assertAstNode(_ast.BinOp, None, None, "10 / 5")
+        tests = (
+            ("1 + 2",         3),
+            ("'@'*4",         "@@@@"),
+            ("'abc' + 'def'", "abcdef"),
+            ("3**4",          81),
+            ("3*4",           12),
+            ("13//4",         3),
+            ("14%4",          2),
+            ("2+3",           5),
+            ("13-4",          9),
+            ("13<<2",         52),
+            ("13>>2",         3),
+            ("13&7",          5),
+            ("13^7",          10),
+            ("13|7",          15)
+        )
+
+        for code, expected in tests:
+            ast = compile(code, "<string>", "exec", _ast.PyCF_ONLY_AST)
+
+            if type(expected) is int:
+                actual = ast.body[0].value.n
+            elif type(expected) is str:
+                actual = ast.body[0].value.s
+            else:
+                raise Exception("Unexpected value: %s" % (expected,))
+
+            self.assertEqual(expected, actual)
 
     def test_binop_fold_num_with_variable(self):
         # check binary constant folding occurs even where
@@ -101,11 +123,6 @@
         self.assertEqual(_ast.Num, ast.body[1].value.right.__class__)
         self.assertEqual(6, ast.body[1].value.right.n)
 
-    def test_binop_fold_str(self):
-        # check binary constant folding for string values
-        self.assertStr("hello there", "'hello' + ' ' + 'there'")
-        self.assertStr("testtesttest", "'test' * 3")
-
     def test_unary_fold_num(self):
         # check unary constant folding for numeric values
         self.assertNum(-5, "-5")
@@ -114,6 +131,15 @@
         self.assertNum(False, "not True")
         self.assertNum(True, "not None")
 
+    def test_return_none_becomes_return(self):
+        code = """
+def foo():
+    return None
+"""
+        ast = self.compileast(code)
+        self.assertEqual(_ast.Return, ast.body[0].body[0].__class__)
+        self.assertEqual(None, ast.body[0].body[0].value)
+
     def test_eliminate_code_after_return(self):
         # ensure code following a "return" is erased from the AST
         code = """

Modified: python/branches/tlee-ast-optimize/Python/optimize.c
==============================================================================
--- python/branches/tlee-ast-optimize/Python/optimize.c	(original)
+++ python/branches/tlee-ast-optimize/Python/optimize.c	Tue Apr 22 14:36:34 2008
@@ -862,9 +862,17 @@
 optimize_return(stmt_ty* stmt_ptr, PyArena* arena)
 {
     stmt_ty stmt = *stmt_ptr;
-    if (stmt->v.Return.value != NULL)
+    if (stmt->v.Return.value != NULL) {
         if (!optimize_expr(&stmt->v.Return.value, arena))
             return 0;
+
+        /* "return None" becomes "return" */
+        if (stmt->v.Return.value->kind == Name_kind) {
+            PyObject* id = stmt->v.Return.value->v.Name.id;
+            if (strcmp(PyString_AS_STRING(id), "None") == 0)
+                stmt->v.Return.value = NULL;
+        }
+    }
     return 1;
 }