From f91e514f4452ac1e6d6efd9b0982723eef2f63ff Mon Sep 17 00:00:00 2001
From: Mysterie <kajusska@gmail.com>
Date: Sat, 22 Sep 2012 15:19:07 +0200
Subject: [PATCH] Improvement of loop handling

---
 test/test_loops.py      |  9 +++++++++
 uncompyle2/Scanner25.py | 19 ++++++++++++++++---
 uncompyle2/Scanner26.py | 11 ++++-------
 3 files changed, 29 insertions(+), 10 deletions(-)

diff --git a/test/test_loops.py b/test/test_loops.py
index 92a786ec..462ef419 100755
--- a/test/test_loops.py
+++ b/test/test_loops.py
@@ -36,3 +36,12 @@ while i < 10:
     print i,
 else:
     print 'Else'
+    
+for x, y in [(1,2),(3,4)]:
+    if x in ['==', '>=', '>']:
+        if '0' in y:
+            print
+
+for x in (1, 2, 3):
+    if x == 1:
+        print x
diff --git a/uncompyle2/Scanner25.py b/uncompyle2/Scanner25.py
index c9375bdb..3f22697f 100755
--- a/uncompyle2/Scanner25.py
+++ b/uncompyle2/Scanner25.py
@@ -85,6 +85,9 @@ class Scanner:
             ret = self.getOpcodeToDel(i)
             if ret != None:
                 toDel += ret
+
+        #self.print_bytecode()
+
         if toDel:
             toDel = sorted(list(set(toDel)))
             delta = 0
@@ -278,9 +281,10 @@ class Scanner:
                     destFor = self.get_target(jmpabs1target)
                     if destFor == i+opsize+4:
                         setupLoop = self.last_instr(0, jmpabs1target, SETUP_LOOP)
-                        assert self.get_target(setupLoop) >= i+opsize+4+self.op_size(POP_BLOCK)
-                        self.restructJump(jmpabs1target, destFor+self.op_size(POP_BLOCK))
-                        toDel += [setupLoop, i+opsize+1, i+opsize+4]
+                        standarFor =  self.last_instr(setupLoop, jmpabs1target, GET_ITER)
+                        if standarFor == None:
+                            self.restructJump(jmpabs1target, destFor+self.op_size(POP_BLOCK))
+                            toDel += [setupLoop, i+opsize+1, i+opsize+4]
             if len(toDel) > 0:
                 return toDel
             return None
@@ -455,6 +459,15 @@ class Scanner:
         target = self.code[pos+1] + self.code[pos+2] * 256
         return target
 
+    def print_bytecode(self):
+        for i in self.op_range(0, len(self.code)):
+            op = self.code[i]
+            if op in dis.hasjabs+dis.hasjrel:
+                dest = self.get_target(i, op)
+                print '%i\t%s\t%i' % (i, dis.opname[op], dest)
+            else:
+                print '%i\t%s\t' % (i, dis.opname[op])
+        
     def first_instr(self, start, end, instr, target=None, exact=True):
         """
         Find the first <instr> in the block from start to end.
diff --git a/uncompyle2/Scanner26.py b/uncompyle2/Scanner26.py
index e951610b..a9a31585 100755
--- a/uncompyle2/Scanner26.py
+++ b/uncompyle2/Scanner26.py
@@ -285,9 +285,10 @@ class Scanner:
                     destFor = self.get_target(jmpabs1target)
                     if destFor == i+opsize+4:
                         setupLoop = self.last_instr(0, jmpabs1target, SETUP_LOOP)
-                        assert self.get_target(setupLoop) >= i+opsize+4+self.op_size(POP_BLOCK)
-                        self.restructJump(jmpabs1target, destFor+self.op_size(POP_BLOCK))
-                        toDel += [setupLoop, i+opsize+1, i+opsize+4]
+                        standarFor =  self.last_instr(setupLoop, jmpabs1target, GET_ITER)
+                        if standarFor == None:
+                            self.restructJump(jmpabs1target, destFor+self.op_size(POP_BLOCK))
+                            toDel += [setupLoop, i+opsize+1, i+opsize+4]
             if len(toDel) > 0:
                 return toDel
             return None
@@ -384,8 +385,6 @@ class Scanner:
             if(self.code[i] in (PJIF,PJIT)):
                 target = self.get_argument(i)
                 target += i + 3
-                if target > 0xFFFF:
-                    raise 'TODO'
                 self.restructJump(i, target)
 
         for i in self.op_range(0, len(self.code)):
@@ -393,8 +392,6 @@ class Scanner:
                 target = self.get_target(i)
                 if self.code[target] == JA:
                     target = self.get_target(target)
-                    if target > 0xFFFF:
-                        raise 'TODO'
                     self.restructJump(i, target)
 
     def restructCode(self, listDel):