From 08451feeabaf24a57fcdb0a36362b6a242cd1f11 Mon Sep 17 00:00:00 2001
From: Jeremy Mikola <jmikola@gmail.com>
Date: Mon, 22 May 2017 18:09:42 -0400
Subject: [PATCH] Unpack nested IteratorAggregate objects for Count

This allows generators (#2149) and internal Traversables (#2642) to be properly detected before attempting to restore the Iterator's position after counting (#1125).
---
 src/Framework/Constraint/Count.php       |  8 ++++----
 tests/Framework/Constraint/CountTest.php | 19 +++++++++++++++++++
 tests/_files/TestIteratorAggregate2.php  | 19 +++++++++++++++++++
 3 files changed, 42 insertions(+), 4 deletions(-)
 create mode 100644 tests/_files/TestIteratorAggregate2.php

diff --git a/src/Framework/Constraint/Count.php b/src/Framework/Constraint/Count.php
index 24e848ca3ef..0fbc54c9c98 100644
--- a/src/Framework/Constraint/Count.php
+++ b/src/Framework/Constraint/Count.php
@@ -56,12 +56,12 @@ protected function getCountOf($other)
         }
 
         if ($other instanceof Traversable) {
-            if ($other instanceof IteratorAggregate) {
-                $iterator = $other->getIterator();
-            } else {
-                $iterator = $other;
+            while ($other instanceof IteratorAggregate) {
+                $other = $other->getIterator();
             }
 
+            $iterator = $other;
+
             if ($iterator instanceof Generator) {
                 return $this->getCountOfGenerator($iterator);
             }
diff --git a/tests/Framework/Constraint/CountTest.php b/tests/Framework/Constraint/CountTest.php
index fabad89cf95..cad90c399b3 100644
--- a/tests/Framework/Constraint/CountTest.php
+++ b/tests/Framework/Constraint/CountTest.php
@@ -25,9 +25,11 @@ public function testCount()
         $countConstraint = new Count(2);
         $it              = new \TestIterator([1, 2]);
         $ia              = new \TestIteratorAggregate($it);
+        $ia2             = new \TestIteratorAggregate2($ia);
 
         $this->assertTrue($countConstraint->evaluate($it, '', true));
         $this->assertTrue($countConstraint->evaluate($ia, '', true));
+        $this->assertTrue($countConstraint->evaluate($ia2, '', true));
     }
 
     public function testCountDoesNotChangeIteratorKey()
@@ -78,6 +80,23 @@ public function testCountDoesNotChangeIteratorKey()
         $it->next();
         $countConstraint->evaluate($ia, '', true);
         $this->assertFalse($it->valid());
+
+        // test with nested IteratorAggregate
+        $it = new \TestIterator([1, 2]);
+        $ia = new \TestIteratorAggregate($it);
+        $ia2 = new \TestIteratorAggregate2($ia);
+
+        $countConstraint = new Count(2);
+        $countConstraint->evaluate($ia2, '', true);
+        $this->assertEquals(1, $it->current());
+
+        $it->next();
+        $countConstraint->evaluate($ia2, '', true);
+        $this->assertEquals(2, $it->current());
+
+        $it->next();
+        $countConstraint->evaluate($ia2, '', true);
+        $this->assertFalse($it->valid());
     }
 
     public function testCountGeneratorsDoNotRewind()
diff --git a/tests/_files/TestIteratorAggregate2.php b/tests/_files/TestIteratorAggregate2.php
new file mode 100644
index 00000000000..e02b89a03dd
--- /dev/null
+++ b/tests/_files/TestIteratorAggregate2.php
@@ -0,0 +1,19 @@
+<?php
+
+/* This class is used for testing a chain of IteratorAggregate objects, since
+ * PHP does allow IteratorAggregate::getIterator() to return an instance of the
+ * same class. */
+class TestIteratorAggregate2 implements IteratorAggregate
+{
+    private $traversable;
+
+    public function __construct(\Traversable $traversable)
+    {
+        $this->traversable = $traversable;
+    }
+
+    public function getIterator()
+    {
+        return $this->traversable;
+    }
+}