[pypy-svn] r52552 - in pypy/dist/pypy: config translator translator/test

[arigo at codespeak.net]

Author: arigo
Date: Sat Mar 15 10:05:53 2008
New Revision: 52552

Modified:
   pypy/dist/pypy/config/translationoption.py
   pypy/dist/pypy/translator/simplify.py
   pypy/dist/pypy/translator/test/test_simplify.py
Log:
Fix a case for list comprehension detection.


Modified: pypy/dist/pypy/config/translationoption.py
==============================================================================
--- pypy/dist/pypy/config/translationoption.py	(original)
+++ pypy/dist/pypy/config/translationoption.py	Sat Mar 15 10:05:53 2008
@@ -135,7 +135,7 @@
                "operations that results from a list comprehension and "
                "attempt to pre-allocate the list",
                default=False,
-               cmdline=None),
+               cmdline='--listcompr'),
     ChoiceOption("fork_before",
                  "(UNIX) Create restartable checkpoint before step",
                  ["annotate", "rtype", "backendopt", "database", "source",

Modified: pypy/dist/pypy/translator/simplify.py
==============================================================================
--- pypy/dist/pypy/translator/simplify.py	(original)
+++ pypy/dist/pypy/translator/simplify.py	Sat Mar 15 10:05:53 2008
@@ -744,12 +744,12 @@
         if (len(block.operations) == 1 and
             block.operations[0].opname == 'next' and
             block.exitswitch == c_last_exception and
-            len(block.exits) == 2 and
-            block.exits[1].exitcase is StopIteration and
-            block.exits[0].exitcase is None):
-            # it's a straightforward loop start block
-            loopnextblocks.append((block, block.operations[0].args[0]))
-            continue
+            len(block.exits) >= 2):
+            cases = [link.exitcase for link in block.exits]
+            if None in cases and StopIteration in cases:
+                # it's a straightforward loop start block
+                loopnextblocks.append((block, block.operations[0].args[0]))
+                continue
         for op in block.operations:
             if op.opname == 'newlist' and not op.args:
                 vlist = variable_families.find_rep(op.result)
@@ -958,9 +958,13 @@
 
             # ... and when the iterator is exhausted, we should no longer
             # reach 'append' at all.
-            assert loopnextblock.exits[1].exitcase is StopIteration
-            stopblock = loopnextblock.exits[1].target
-            if self.reachable(stopblock, appendblock, avoid=newlistblock):
+            stopblocks = [link.target for link in loopnextblock.exits
+                                      if link.exitcase is not None]
+            accepted = True
+            for stopblock1 in stopblocks:
+                if self.reachable(stopblock1, appendblock, avoid=newlistblock):
+                    accepted = False
+            if not accepted:
                 continue
 
             # now explicitly find the "loop body" blocks: they are the ones
@@ -1000,7 +1004,8 @@
         # Found a suitable loop, let's patch the graph:
         assert iterblock not in loopbody
         assert loopnextblock in loopbody
-        assert stopblock not in loopbody
+        for stopblock1 in stopblocks:
+            assert stopblock1 not in loopbody
 
         # at StopIteration, the new list is exactly of the same length as
         # the one we iterate over if it's not possible to skip the appendblock
@@ -1040,7 +1045,7 @@
                         continue  # list not passed along this link anyway
                     hints = {'fence': True}
                     if (exactlength and block is loopnextblock and
-                        link.target is stopblock):
+                        link.target in stopblocks):
                         hints['exactlength'] = True
                     chints = Constant(hints)
                     newblock = insert_empty_block(None, link)

Modified: pypy/dist/pypy/translator/test/test_simplify.py
==============================================================================
--- pypy/dist/pypy/translator/test/test_simplify.py	(original)
+++ pypy/dist/pypy/translator/test/test_simplify.py	Sat Mar 15 10:05:53 2008
@@ -220,6 +220,30 @@
         'hint': 2,
         }
 
+def test_detect_list_comprehension_with_exc():
+    def g(x):
+        return x * 17
+    def free_some_stuff():
+        pass
+    def f1(l):
+        try:
+            return [g(x) for x in l]
+        finally:
+            free_some_stuff()
+
+    t = TranslationContext(list_comprehension_operations=True)
+    graph = t.buildflowgraph(f1)
+    if conftest.option.view:
+        graph.show()
+    assert summary(graph) == {
+        'newlist': 1,
+        'iter': 1,
+        'next': 1,
+        'getattr': 1,
+        'simple_call': 4,
+        'hint': 2,
+        }
+
 class TestLLSpecializeListComprehension:
     typesystem = 'lltype'