pyassem.py:

    Fix annoying bugs in flow graph layout code.  In some cases the
    implicit control transfers weren't honored.  In other cases,
    JUMP_FORWARD instructions jumped backwards.

    Remove unused arg from nextBlock().

pycodegen.py

    Add optional force kwarg to set_lineno() that will emit a
    SET_LINENO even if it is the same as the previous lineno.

    Use explicit LOAD_FAST and STORE_FAST to access list comp implicit
    variables.  (The symbol table doesn't know about them.)

2 files changed, 111 insertions, 14 deletions

diff --git a/Lib/compiler/pyassem.py b/Lib/compiler/pyassem.py
index 15e91f4..e8810ae 100644
--- a/Lib/compiler/pyassem.py
+++ b/Lib/compiler/pyassem.py
@@ -27,11 +27,12 @@ class FlowGraph:
         if self._debug:
             if self.current:
                 print "end", repr(self.current)
+                print "    next", self.current.next
                 print "   ", self.current.get_children()
             print repr(block)
         self.current = block
 
-    def nextBlock(self, block=None, force=0):
+    def nextBlock(self, block=None):
         # XXX think we need to specify when there is implicit transfer
         # from one block to the next.  might be better to represent this
         # with explicit JUMP_ABSOLUTE instructions that are optimized
@@ -95,12 +96,110 @@ class FlowGraph:
                 b.addNext(self.exit)
         order = dfs_postorder(self.entry, {})
         order.reverse()
+        self.fixupOrder(order, self.exit)
         # hack alert
         if not self.exit in order:
             order.append(self.exit)
 
         return order
 
+    def fixupOrder(self, blocks, default_next):
+        """Fixup bad order introduced by DFS."""
+
+        # XXX This is a total mess.  There must be a better way to get
+        # the code blocks in the right order.
+        
+        self.fixupOrderHonorNext(blocks, default_next)
+        self.fixupOrderForward(blocks, default_next)
+
+    def fixupOrderHonorNext(self, blocks, default_next):
+        """Fix one problem with DFS.
+        
+        The DFS uses child block, but doesn't know about the special
+        "next" block.  As a result, the DFS can order blocks so that a
+        block isn't next to the right block for implicit control
+        transfers.
+        """
+        index = {}
+        for i in range(len(blocks)):
+            index[blocks[i]] = i
+
+        for i in range(0, len(blocks) - 1):
+            b = blocks[i]
+            n = blocks[i + 1]
+            if not b.next or b.next[0] == default_next or b.next[0] == n:
+                continue
+            # The blocks are in the wrong order.  Find the chain of
+            # blocks to insert where they belong.
+            cur = b
+            chain = []
+            elt = cur
+            while elt.next and elt.next[0] != default_next:
+                chain.append(elt.next[0])
+                elt = elt.next[0]
+            # Now remove the blocks in the chain from the current
+            # block list, so that they can be re-inserted.
+            l = []
+            for b in chain:
+                assert index[b] > i
+                l.append((index[b], b))
+            l.sort()
+            l.reverse()
+            for j, b in l:
+                del blocks[index[b]]
+            # Insert the chain in the proper location
+            blocks[i:i + 1] = [cur] + chain
+            # Finally, re-compute the block indexes
+            for i in range(len(blocks)):
+                index[blocks[i]] = i
+
+    def fixupOrderForward(self, blocks, default_next):
+        """Make sure all JUMP_FORWARDs jump forward"""
+        index = {}
+        chains = []
+        cur = []
+        for b in blocks:
+            index[b] = len(chains)
+            cur.append(b)
+            if b.next and b.next[0] == default_next:
+                chains.append(cur)
+                cur = []
+        chains.append(cur)
+
+        while 1:
+            constraints = []
+
+            for i in range(len(chains)):
+                l = chains[i]
+                for b in l:
+                    for c in b.get_children():
+                        if index[c] < i:
+                            forward_p = 0
+                            for inst in b.insts:
+                                if inst[0] == 'JUMP_FORWARD':
+                                    if inst[1] == c:
+                                        forward_p = 1
+                            if not forward_p:
+                                continue
+                            constraints.append((index[c], i))
+
+            if not constraints:
+                break
+
+            # XXX just do one for now
+            # do swaps to get things in the right order
+            goes_before, a_chain = constraints[0]
+            assert a_chain > goes_before
+            c = chains[a_chain]
+            chains.remove(c)
+            chains.insert(goes_before, c)
+
+
+        del blocks[:]
+        for c in chains:
+            for b in c:
+                blocks.append(b)
+            
     def getBlocks(self):
         return self.blocks.elements()
 
diff --git a/Lib/compiler/pycodegen.py b/Lib/compiler/pycodegen.py
index 691232f..63b0a57 100644
--- a/Lib/compiler/pycodegen.py
+++ b/Lib/compiler/pycodegen.py
@@ -193,7 +193,7 @@ class CodeGenerator:
         else:
             self.emit(prefix + '_GLOBAL', name)
 
-    def set_lineno(self, node):
+    def set_lineno(self, node, force=0):
         """Emit SET_LINENO if node has lineno attribute and it is 
         different than the last lineno emitted.
 
@@ -205,7 +205,8 @@ class CodeGenerator:
         then, this method works around missing line numbers.
         """
         lineno = getattr(node, 'lineno', None)
-        if lineno is not None and lineno != self.last_lineno:
+        if lineno is not None and (lineno != self.last_lineno
+                                   or force):
             self.emit('SET_LINENO', lineno)
             self.last_lineno = lineno
             return 1
@@ -413,9 +414,6 @@ class CodeGenerator:
     __list_count = 0
     
     def visitListComp(self, node):
-        # XXX would it be easier to transform the AST into the form it
-        # would have if the list comp were expressed as a series of
-        # for and if stmts and an explicit append?
         self.set_lineno(node)
         # setup list
         append = "$append%d" % self.__list_count
@@ -423,10 +421,10 @@ class CodeGenerator:
         self.emit('BUILD_LIST', 0)
         self.emit('DUP_TOP')
         self.emit('LOAD_ATTR', 'append')
-        self.storeName(append)
-        l = len(node.quals)
+        self.emit('STORE_FAST', append)
+        
         stack = []
-        for i, for_ in zip(range(l), node.quals):
+        for i, for_ in zip(range(len(node.quals)), node.quals):
             start, anchor = self.visit(for_)
             cont = None
             for if_ in for_.ifs:
@@ -434,8 +432,8 @@ class CodeGenerator:
                     cont = self.newBlock()
                 self.visit(if_, cont)
             stack.insert(0, (start, cont, anchor))
-            
-        self.loadName(append)
+
+        self.emit('LOAD_FAST', append)
         self.visit(node.expr)
         self.emit('CALL_FUNCTION', 1)
         self.emit('POP_TOP')
@@ -449,12 +447,11 @@ class CodeGenerator:
                 self.nextBlock(skip_one)
             self.emit('JUMP_ABSOLUTE', start)
             self.startBlock(anchor)
-        self.delName(append)
+        self.emit('DELETE_FAST', append)
         
         self.__list_count = self.__list_count - 1
 
     def visitListCompFor(self, node):
-        self.set_lineno(node)
         start = self.newBlock()
         anchor = self.newBlock()
 
@@ -468,7 +465,7 @@ class CodeGenerator:
         return start, anchor
 
     def visitListCompIf(self, node, branch):
-        self.set_lineno(node)
+        self.set_lineno(node, force=1)
         self.visit(node.test)
         self.emit('JUMP_IF_FALSE', branch)
         self.newBlock()
@@ -672,6 +669,7 @@ class CodeGenerator:
             print node
 
     def _visitAssSequence(self, node, op='UNPACK_SEQUENCE'):
+##        print >> sys.stderr, "AssSequence", op, findOp(node), node
         if findOp(node) != 'OP_DELETE':
             self.emit(op, len(node.nodes))
         for child in node.nodes: