http://bugs.python.org/issue2459 -- speed up loops

Lib/compiler/pycodegen.py

 import imp
 import os
 import marshal
 import struct
 import sys
 from cStringIO import StringIO
 
 from compiler import ast, parse, walk, syntax
 from compiler import pyassem, misc, future, symbols
 from compiler.consts import SC_LOCAL, SC_GLOBAL_IMPLICIT, SC_GLOBAL_EXPLICT, \
@@ skipping 395 matching lines @@
         self.emit('BUILD_CLASS')
         self.storeName(node.name)
 
     # The rest are standard visitor methods
 
     # The next few implement control-flow statements
 
     def visitIf(self, node):
         end = self.newBlock()
         numtests = len(node.tests)
-        for i in range(numtests):
+        for i in range(numtests - 1):
             test, suite = node.tests[i]
             if is_constant_false(test):
                 # XXX will need to check generator stuff here
                 continue
             self.set_lineno(test)
             self.visit(test)
             nextTest = self.newBlock()
             self.emit('POP_JUMP_IF_FALSE', nextTest)
             self.nextBlock()
             self.visit(suite)
             self.emit('JUMP_FORWARD', end)
             self.startBlock(nextTest)
-        if node.else_:
-            self.visit(node.else_)
+        # The last test can optimize a few things
+        test, suite = node.tests[numtests - 1]
+        if is_constant_false(test):
+            if node.else_:
+                self.visit(node.else_)
+        else:
+            self.set_lineno(test)
+            self.visit(test)
+            else_block = self.newBlock() if node.else_ else end
+            self.emit('POP_JUMP_IF_FALSE', else_block)
+            self.nextBlock()
+            self.visit(suite)
+            if node.else_:
+                self.emit('JUMP_FORWARD', end)
+                self.startBlock(else_block)
+                self.visit(node.else_)
         self.nextBlock(end)
 
     def visitWhile(self, node):
         self.set_lineno(node)
 
         loop = self.newBlock()
         else_ = self.newBlock()
+        body = self.newBlock()
+        tail = self.newBlock()
+        after = self.newBlock()
 
-        after = self.newBlock()
         self.emit('SETUP_LOOP', after)
 
         self.nextBlock(loop)
         self.setups.push((LOOP, loop))
+        self.emit('JUMP_FORWARD', tail)
+        self.startBlock(body)
 
         self.set_lineno(node, force=True)
+
+        self.visit(node.body)
+
+        self.nextBlock(tail)
         self.visit(node.test)
-        self.emit('POP_JUMP_IF_FALSE', else_ or after)
+        self.emit('POP_JUMP_IF_TRUE', body)
 
-        self.nextBlock()
-        self.visit(node.body)
-        self.emit('JUMP_ABSOLUTE', loop)
-
-        self.startBlock(else_) # or just the POPs if not else clause
+        self.nextBlock(else_) # or just the POPs if not else clause
         self.emit('POP_BLOCK')
         self.setups.pop()
         if node.else_:
             self.visit(node.else_)
         self.nextBlock(after)
 
     def visitFor(self, node):
         start = self.newBlock()
+        body = self.newBlock()
+        tail = self.newBlock()
         anchor = self.newBlock()
         after = self.newBlock()
         self.setups.push((LOOP, start))
 
         self.set_lineno(node)
         self.emit('SETUP_LOOP', after)
         self.visit(node.list)
         self.emit('GET_ITER')
 
         self.nextBlock(start)
         self.set_lineno(node, force=1)
-        self.emit('FOR_ITER', anchor)
+        self.emit('JUMP_FORWARD', tail)
+        self.startBlock(body)
         self.visit(node.assign)
         self.visit(node.body)
-        self.emit('JUMP_ABSOLUTE', start)
+        self.nextBlock(tail)
+        self.emit('FOR_ITER', body)
         self.nextBlock(anchor)
         self.emit('POP_BLOCK')
         self.setups.pop()
         if node.else_:
             self.visit(node.else_)
         self.nextBlock(after)
 
     def visitBreak(self, node):
         if not self.setups:
             raise SyntaxError, "'break' outside loop (%s, %d)" % \
@@ skipping 77 matching lines @@
             self.emit('POP_TOP')
             self.nextBlock(end)
 
     # list comprehensions
     def visitListComp(self, node):
         self.set_lineno(node)
         # setup list
         self.emit('BUILD_LIST', 0)
 
         stack = []
-        for i, for_ in zip(range(len(node.quals)), node.quals):
-            start, anchor = self.visit(for_)
-            cont = None
+        for for_ in node.quals:
+            body, tail = self.visit(for_)
             for if_ in for_.ifs:
-                if cont is None:
-                    cont = self.newBlock()
-                self.visit(if_, cont)
-            stack.insert(0, (start, cont, anchor))
+                self.visit(if_, tail)
+            stack.insert(0, (for_, body, tail))
 
         self.visit(node.expr)
         self.emit('LIST_APPEND', len(node.quals) + 1)
 
-        for start, cont, anchor in stack:
-            if cont:
-                self.nextBlock(cont)
-            self.emit('JUMP_ABSOLUTE', start)
-            self.startBlock(anchor)
+        for for_, body, tail in stack:
+            self.nextBlock(tail)
+            self.set_lineno(for_, force=True)
+            self.emit('FOR_ITER', body)
+            self.nextBlock()
 
     def visitListCompFor(self, node):
-        start = self.newBlock()
-        anchor = self.newBlock()
+        body = self.newBlock()
+        tail = self.newBlock()
 
         self.visit(node.list)
         self.emit('GET_ITER')
-        self.nextBlock(start)
-        self.set_lineno(node, force=True)
-        self.emit('FOR_ITER', anchor)
         self.nextBlock()
+        self.emit('JUMP_FORWARD', tail)
+        self.nextBlock(body)
         self.visit(node.assign)
-        return start, anchor
+        return body, tail
 
     def visitListCompIf(self, node, branch):
         self.set_lineno(node, force=True)
         self.visit(node.test)
         self.emit('POP_JUMP_IF_FALSE', branch)
-        self.newBlock()
 
     def _makeClosure(self, gen, args):
         frees = gen.scope.get_free_vars()
         if frees:
             for name in frees:
                 self.emit('LOAD_CLOSURE', name)
             self.emit('BUILD_TUPLE', len(frees))
             self.emit('LOAD_CONST', gen)
             self.emit('MAKE_CLOSURE', args)
         else:
@@ skipping 10 matching lines @@
         # precomputation of outmost iterable
         self.visit(node.code.quals[0].iter)
         self.emit('GET_ITER')
         self.emit('CALL_FUNCTION', 1)
 
     def visitGenExprInner(self, node):
         self.set_lineno(node)
         # setup list
 
         stack = []
-        for i, for_ in zip(range(len(node.quals)), node.quals):
-            start, anchor, end = self.visit(for_)
-            cont = None
+        for for_ in node.quals:
+            body, tail, end = self.visit(for_)
             for if_ in for_.ifs:
-                if cont is None:
-                    cont = self.newBlock()
-                self.visit(if_, cont)
-            stack.insert(0, (start, cont, anchor, end))
+                self.visit(if_, tail)
+            stack.insert(0, (for_, body, tail, end))
 
         self.visit(node.expr)
         self.emit('YIELD_VALUE')
         self.emit('POP_TOP')
 
-        for start, cont, anchor, end in stack:
-            if cont:
-                self.nextBlock(cont)
-            self.emit('JUMP_ABSOLUTE', start)
-            self.startBlock(anchor)
+        for for_, body, tail, end in stack:
+            self.nextBlock(tail)
+            self.set_lineno(for_, force=True)
+            self.emit('FOR_ITER', body)
+            self.nextBlock()
             self.emit('POP_BLOCK')
             self.setups.pop()
             self.nextBlock(end)
 
         self.emit('LOAD_CONST', None)
 
     def visitGenExprFor(self, node):
         start = self.newBlock()
-        anchor = self.newBlock()
+        body = self.newBlock()
+        tail = self.newBlock()
         end = self.newBlock()
 
         self.setups.push((LOOP, start))
         self.emit('SETUP_LOOP', end)
 
         if node.is_outmost:
             self.loadName('.0')
         else:
             self.visit(node.iter)
             self.emit('GET_ITER')
 
         self.nextBlock(start)
-        self.set_lineno(node, force=True)
-        self.emit('FOR_ITER', anchor)
-        self.nextBlock()
+        self.emit('JUMP_FORWARD', tail)
+        self.nextBlock(body)
         self.visit(node.assign)
-        return start, anchor, end
+        return body, tail, end
 
     def visitGenExprIf(self, node, branch):
         self.set_lineno(node, force=True)
         self.visit(node.test)
         self.emit('POP_JUMP_IF_FALSE', branch)
         self.newBlock()
 
     # exception related
 
     def visitAssert(self, node):
@@ skipping 797 matching lines @@
     ast.Slice: AugSlice,
     ast.Subscript: AugSubscript,
     }
 
 def wrap_aug(node):
     return wrapper[node.__class__](node)
 
 if __name__ == "__main__":
     for file in sys.argv[1:]:
         compileFile(file)