[pypy-commit] pypy default: Improve itertools.product()

[arigo]

Author: Armin Rigo <arigo at tunes.org>
Branch: 
Changeset: r68843:c051852e3f7d
Date: 2014-01-22 14:05 +0100
http://bitbucket.org/pypy/pypy/changeset/c051852e3f7d/

Log:	Improve itertools.product()

diff --git a/pypy/module/itertools/interp_itertools.py b/pypy/module/itertools/interp_itertools.py
--- a/pypy/module/itertools/interp_itertools.py
+++ b/pypy/module/itertools/interp_itertools.py
@@ -1083,58 +1083,64 @@
 class W_Product(W_Root):
     def __init__(self, space, args_w, w_repeat):
         self.gears = [
-            space.fixedview(arg_w) for arg_w in args_w
+            space.unpackiterable(arg_w) for arg_w in args_w
         ] * space.int_w(w_repeat)
-        self.num_gears = len(self.gears)
-        # initialization of indicies to loop over
-        self.indicies = [
-            (0, len(gear))
-            for gear in self.gears
-        ]
-        self.cont = True
-        for _, lim in self.indicies:
-            if lim <= 0:
-                self.cont = False
+        #
+        for gear in self.gears:
+            if len(gear) == 0:
+                self.lst = None
                 break
+        else:
+            self.indices = [0] * len(self.gears)
+            self.lst = [gear[0] for gear in self.gears]
 
-    def roll_gears(self):
-        if self.num_gears == 0:
-            self.cont = False
-            return
+    def _rotate_previous_gears(self):
+        lst = self.lst
+        x = len(self.gears) - 1
+        lst[x] = self.gears[x][0]
+        self.indices[x] = 0
+        x -= 1
+        # the outer loop runs as long as a we have a carry
+        while x >= 0:
+            gear = self.gears[x]
+            index = self.indices[x] + 1
+            if index < len(gear):
+                # no carry: done
+                lst[x] = gear[index]
+                self.indices[x] = index
+                return
+            lst[x] = gear[0]
+            self.indices[x] = 0
+            x -= 1
+        else:
+            self.lst = None
 
-        # Starting from the end of the gear indicies work to the front
-        # incrementing the gear until the limit is reached. When the limit
-        # is reached carry operation to the next gear
-        should_carry = True
-
-        for n in range(0, self.num_gears):
-            nth_gear = self.num_gears - n - 1
-            if should_carry:
-                count, lim = self.indicies[nth_gear]
-                count += 1
-                if count == lim and nth_gear == 0:
-                    self.cont = False
-                if count == lim:
-                    should_carry = True
-                    count = 0
-                else:
-                    should_carry = False
-                self.indicies[nth_gear] = (count, lim)
+    def fill_next_result(self):
+        # the last gear is done here, in a function with no loop,
+        # to allow the JIT to look inside
+        lst = self.lst
+        x = len(self.gears) - 1
+        if x >= 0:
+            gear = self.gears[x]
+            index = self.indices[x] + 1
+            if index < len(gear):
+                # no carry: done
+                lst[x] = gear[index]
+                self.indices[x] = index
             else:
-                break
+                self._rotate_previous_gears()
+        else:
+            self.lst = None
 
     def iter_w(self, space):
         return space.wrap(self)
 
     def next_w(self, space):
-        if not self.cont:
+        if self.lst is None:
             raise OperationError(space.w_StopIteration, space.w_None)
-        l = [None] * self.num_gears
-        for x in range(0, self.num_gears):
-            index, limit = self.indicies[x]
-            l[x] = self.gears[x][index]
-        self.roll_gears()
-        return space.newtuple(l)
+        w_result = space.newtuple(self.lst[:])
+        self.fill_next_result()
+        return w_result
 
 
 def W_Product__new__(space, w_subtype, __args__):