Optimize product(), add tests

nikic · Jul 3, 2015 · 02177e8 · 02177e8
1 parent f444b6f
commit 02177e8
Show file tree

Hide file tree

Showing 3 changed files with 59 additions and 18 deletions.
diff --git a/src/iter.php b/src/iter.php
@@ -329,9 +329,10 @@ function chain(/* ...$iterables */) {
 /**
  * Returns the cartesian product of iterables that were passed as arguments.
  *
- * The resulting iterator will contain all the possible tuples of keys and values.
+ * The resulting iterator will contain all the possible tuples of keys and
+ * values.
  *
- * Please note that the iterables after the first must be rewindable
+ * Please note that the iterables after the first must be rewindable.
  *
  * Examples:
  *
@@ -343,22 +344,38 @@ function chain(/* ...$iterables */) {
  * @return \Iterator
  */
 function product(/* ...$iterables */) {
-    $iterables = func_get_args();
-
-    if (!$iterables) {
+    /** @var \Iterator[] $iterators */
+    $iterators = array_map('iter\\toIter', func_get_args());
+    $numIterators = count($iterators);
+    if (!$numIterators) {
         yield [] => [];
         return;
     }
 
-    $head = array_shift($iterables);
-    $tailProduct = call_user_func_array("iter\\rewindable\\product", $iterables);
+    $keyTuple = $valueTuple = array_fill(0, $numIterators, null);
 
-    foreach ($head as $headKey => $headValue) {
-        foreach ($tailProduct as $tailKey => $tailValue) {
-            array_unshift($tailValue, $headValue);
-            array_unshift($tailKey, $headKey);
-            yield $tailKey => $tailValue;
+    $i = -1;
+    while (true) {
+        while (++$i < $numIterators - 1) {
+            $iterators[$i]->rewind();
+            if (!$iterators[$i]->valid()) {
+                return;
+            }
+            $keyTuple[$i] = $iterators[$i]->key();
+            $valueTuple[$i] = $iterators[$i]->current();
+        }
+        foreach ($iterators[$i] as $keyTuple[$i] => $valueTuple[$i]) {
+            yield $keyTuple => $valueTuple;
+        }
+        while (--$i >= 0) {
+            $iterators[$i]->next();
+            if ($iterators[$i]->valid()) {
+                $keyTuple[$i] = $iterators[$i]->key();
+                $valueTuple[$i] = $iterators[$i]->current();
+                continue 2;
+            }
         }
+        return;
     }
 }
 
@@ -779,11 +796,7 @@ function toIter($iterable) {
     if ($iterable instanceof \IteratorAggregate) {
         return $iterable->getIterator();
     }
-    return call_user_func(function() use ($iterable) {
-        foreach ($iterable as $key => $value) {
-            yield $key => $value;
-        }
-    });
+    return new \ArrayIterator($iterable);
 }
 
 /**

diff --git a/test/IterRewindableTest.php b/test/IterRewindableTest.php
@@ -94,12 +94,15 @@ public function testRewindableVariants() {
             rewindable\flip(['a' => 1, 'b' => 2, 'c' => 3]),
             true
         );
-
         $this->assertRewindableEquals(
             [['a' => 1, 'b' => 2], ['c' => 3, 'd' => 4], ['e' => 5]],
             rewindable\chunk(['a' => 1, 'b' => 2, 'c' => 3, 'd' => 4, 'e' => 5], 2),
             true
         );
+        $this->assertRewindableEquals(
+            [[1,3,5],[1,3,6],[1,4,5],[1,4,6],[2,3,5],[2,3,6],[2,4,5],[2,4,6]],
+            rewindable\product([1,2], [3,4], [5,6])
+        );
     }
 
     public function testMakeRewindable() {

diff --git a/test/iterTest.php b/test/iterTest.php
@@ -282,6 +282,31 @@ public function testChunk() {
         $this->assertSame([[0, 1, 2]], toArray(chunk([0, 1, 2], 100000)));
         $this->assertSame([], toArray(chunk([], 100000)));
     }
+
+    public function testProduct() {
+        $this->assertKeysValues([[]], [[]], function() { return product(); });
+
+        $this->assertKeysValues(
+            [[0],[1]], [[1],[2]], function() { return product([1,2]); });
+
+        $this->assertKeysValues(
+            [[0,0],[0,1],[1,0],[1,1]],
+            [[1,3],[1,4],[2,3],[2,4]],
+            function() { return product([1,2],[3,4]); });
+
+        $this->assertKeysValues(
+            [[0,0,0],[0,0,1],[0,1,0],[0,1,1],[1,0,0],[1,0,1],[1,1,0],[1,1,1]],
+            [[1,1,1],[1,1,2],[1,2,1],[1,2,2],[2,1,1],[2,1,2],[2,2,1],[2,2,2]],
+            function() {
+                return product(range(1,2), [1,2], new \ArrayIterator([1,2]));
+            }
+        );
+    }
+
+    private function assertKeysValues(array $keys, array $values, callable $fn) {
+        $this->assertSame($keys, toArray(keys($fn())));
+        $this->assertSame($values, toArray(values($fn())));
+    }
 }
 
 class _CountableTestDummy implements \Countable {