Add implementation of Countable for CachingIterator

lib/Doctrine/ODM/MongoDB/Iterator/CachingIterator.php

@@ -4,11 +4,13 @@
 
 namespace Doctrine\ODM\MongoDB\Iterator;
 
+use Countable;
 use Generator;
 use ReturnTypeWillChange;
 use RuntimeException;
 use Traversable;
 
+use function count;
 use function current;
 use function key;
 use function next;
@@ -26,7 +28,7 @@
  * @template TValue
  * @template-implements Iterator<TValue>
  */
-final class CachingIterator implements Iterator
+final class CachingIterator implements Countable, Iterator
 {
     /** @var array<mixed, TValue> */
     private $items = [];
@@ -55,6 +57,18 @@ public function __construct(Traversable $iterator)
         $this->storeCurrentItem();
     }
 
+    /**
+     * @see https://php.net/countable.count
+     * @return int
+     */
+    #[ReturnTypeWillChange]
+    public function count(): int
+    {
+        $this->exhaustIterator();
+
+        return count($this->items);
+    }
+
     public function __destruct()
     {
         $this->iterator = null;

tests/Doctrine/ODM/MongoDB/Tests/Functional/Iterator/CachingIteratorTest.php

@@ -5,6 +5,7 @@
 namespace Doctrine\ODM\MongoDB\Tests\Functional\Iterator;
 
 use Doctrine\ODM\MongoDB\Iterator\CachingIterator;
+use Iterator;
 use Exception;
 use Generator;
 use PHPUnit\Framework\TestCase;
@@ -110,6 +111,71 @@ public function testToArrayAfterPartialIteration(): void
         $this->assertSame([1, 2, 3], $iterator->toArray());
     }
 
+    public function testCount(): void
+    {
+        $iterator = new CachingIterator($this->getTraversable([1, 2, 3]));
+        $this->assertCount(3, $iterator);
+    }
+
+    public function testCountAfterPartialIteration(): void
+    {
+        $iterator = new CachingIterator($this->getTraversable([1, 2, 3]));
+
+        $iterator->rewind();
+        $this->assertTrue($iterator->valid());
+        $this->assertSame(0, $iterator->key());
+        $this->assertSame(1, $iterator->current());
+
+        $iterator->next();
+        $this->assertCount(3, $iterator);
+    }
+
+    public function testCountWithEmptySet(): void
+    {
+        $iterator = new CachingIterator($this->getTraversable([]));
+        $this->assertCount(0, $iterator);
+    }
+
+    /**
+     * This protects against iterators that return valid keys on invalid
+     * positions, which was the case in ext-mongodb until PHPC-1748 was fixed.
+     */
+    public function testWithWrongIterator(): void
+    {
+        $nestedIterator = new class implements Iterator {
+            /** @var int */
+            private $i = 0;
+
+            public function current(): int
+            {
+                return $this->i;
+            }
+
+            public function next(): void
+            {
+                $this->i++;
+            }
+
+            public function key(): int
+            {
+                return $this->i;
+            }
+
+            public function valid(): bool
+            {
+                return $this->i == 0;
+            }
+
+            public function rewind(): void
+            {
+                $this->i = 0;
+            }
+        };
+
+        $iterator = new CachingIterator($nestedIterator);
+        $this->assertCount(1, $iterator);
+    }
+
     private function getTraversable(array $items): Generator
     {
         foreach ($items as $item) {