ResultSet: Implement interface Countable

src/ResultSet.php

@@ -3,15 +3,17 @@
 namespace ipl\Orm;
 
 use ArrayIterator;
+use Countable;
 use Generator;
 use Iterator;
+use RuntimeException;
 use Traversable;
 
 /**
  * @template TRow of Model
  * @implements Iterator<int, TRow>
  */
-class ResultSet implements Iterator
+class ResultSet implements Iterator, Countable
 {
     /** @var ArrayIterator<int, TRow> */
     protected ArrayIterator $cache;
@@ -29,6 +31,8 @@ class ResultSet implements Iterator
 
     protected ?int $position = null;
 
+    protected ?int $count = null;
+
     /**
      * Create a new result set from the given traversable
      *
@@ -98,6 +102,10 @@ public function next(): void
         }
 
         if ($this->isCacheDisabled || ! $this->cache->valid()) {
+            // Raise count during the first loop only after each iteration, so
+            // that it is synchronized with how many times a loop has been run.
+            $this->count += 1;
+
             $this->generator->next();
             $this->advance();
         } else {
@@ -131,11 +139,28 @@ public function rewind(): void
 
         if ($this->position === null) {
             $this->advance();
+            $this->count = 0;
         } else {
             $this->position = 0;
         }
     }
 
+    public function count(): int
+    {
+        if (! $this->isCacheDisabled && $this->count === null && $this->cache->count() === 0) {
+            foreach ($this as $_) {
+                // exhaust the generator and establish the cache
+            }
+        } elseif (
+            $this->count === null
+            || ($this->limit === null || $this->count < $this->limit) && $this->hasMore()
+        ) {
+            throw new RuntimeException('Cannot count result set while it is not fully iterated');
+        }
+
+        return $this->count;
+    }
+
     protected function advance(): void
     {
         if (! $this->generator->valid()) {

tests/ResultSetTest.php

@@ -39,8 +39,8 @@ public function testResultWithCacheDisabled()
         }
 
         $this->assertEquals(
-            $items,
-            ['a', 'b', 'c']
+            ['a', 'b', 'c'],
+            $items
         );
     }
 
@@ -58,8 +58,8 @@ public function testResultWithCacheEnabled()
         }
 
         $this->assertEquals(
-            $items,
-            ['a', 'b', 'c', 'a', 'b', 'c']
+            ['a', 'b', 'c', 'a', 'b', 'c'],
+            $items
         );
     }
 
@@ -77,8 +77,172 @@ public function testResultWithCacheEnabledWithLimit()
         }
 
         $this->assertEquals(
-            $items,
-            ['a', 'b', 'a', 'b']
+            ['a', 'b', 'a', 'b'],
+            $items
         );
     }
+
+    public function testCountWithCacheDisabled(): void
+    {
+        $set = (new ResultSet(new ArrayIterator(['a', 'b', 'c'])))->disableCache();
+
+        foreach ($set as $item) {
+            // pass
+        }
+
+        $this->assertSame(3, $set->count());
+
+        $limitedSet = (new ResultSet(new ArrayIterator(['a', 'b', 'c']), 2))->disableCache();
+
+        foreach ($limitedSet as $item) {
+            // pass
+        }
+
+        $this->assertSame(2, $limitedSet->count());
+        $this->assertTrue($limitedSet->hasMore());
+
+        $partialSet = (new ResultSet(new ArrayIterator(['a', 'b']), 3))->disableCache();
+
+        foreach ($partialSet as $item) {
+            // pass
+        }
+
+        $this->assertSame(2, $partialSet->count());
+        $this->assertFalse($partialSet->hasMore());
+    }
+
+    public function testCountWithCacheEnabled(): void
+    {
+        $set = new ResultSet(new ArrayIterator(['a', 'b', 'c']));
+
+        foreach ($set as $item) {
+            // pass
+        }
+
+        $this->assertSame(3, $set->count());
+
+        // During a subsequent iteration, count should be allowed
+        foreach ($set as $item) {
+            $this->assertSame(3, $set->count());
+        }
+
+        $limitedSet = new ResultSet(new ArrayIterator(['a', 'b', 'c']), 2);
+
+        foreach ($limitedSet as $item) {
+            // pass
+        }
+
+        $this->assertSame(2, $limitedSet->count());
+        $this->assertTrue($limitedSet->hasMore());
+
+        $partialSet = new ResultSet(new ArrayIterator(['a', 'b']), 3);
+
+        foreach ($partialSet as $item) {
+            // pass
+        }
+
+        $this->assertSame(2, $partialSet->count());
+        $this->assertFalse($partialSet->hasMore());
+    }
+
+    public function testCountWithCacheDisabledBeforeIteration(): void
+    {
+        $set = (new ResultSet(new ArrayIterator(['a', 'b', 'c'])))->disableCache();
+
+        $this->expectException(\RuntimeException::class);
+        $this->expectExceptionMessage('Cannot count result set while it is not fully iterated');
+
+        $set->count();
+    }
+
+    public function testLimitedCountWithCacheDisabledBeforeIteration(): void
+    {
+        $set = (new ResultSet(new ArrayIterator(['a', 'b', 'c']), 2))->disableCache();
+
+        $this->expectException(\RuntimeException::class);
+        $this->expectExceptionMessage('Cannot count result set while it is not fully iterated');
+
+        $set->count();
+    }
+
+    public function testCountWithCacheDisabledDuringIteration(): void
+    {
+        $set = (new ResultSet(new ArrayIterator(['a', 'b', 'c'])))->disableCache();
+
+        $this->expectException(\RuntimeException::class);
+        $this->expectExceptionMessage('Cannot count result set while it is not fully iterated');
+
+        foreach ($set as $item) {
+            $set->count();
+        }
+    }
+
+    public function testLimitedCountWithCacheDisabledDuringIteration(): void
+    {
+        $set = (new ResultSet(new ArrayIterator(['a', 'b', 'c']), 2))->disableCache();
+
+        $this->expectException(\RuntimeException::class);
+        $this->expectExceptionMessage('Cannot count result set while it is not fully iterated');
+
+        foreach ($set as $item) {
+            $set->count();
+        }
+    }
+
+    public function testCountWithCacheBeforeIteration(): void
+    {
+        $set = new ResultSet(new ArrayIterator(['a', 'b', 'c']));
+
+        $this->assertSame(3, $set->count());
+
+        $result = [];
+        foreach ($set as $item) {
+            $result[] = $item;
+        }
+
+        $this->assertSame(['a', 'b', 'c'], $result);
+    }
+
+    public function testLimitedCountWithCacheBeforeIteration(): void
+    {
+        $set = new ResultSet(new ArrayIterator(['a', 'b', 'c']), 2);
+
+        $this->assertSame(2, $set->count());
+
+        $result = [];
+        foreach ($set as $item) {
+            $result[] = $item;
+        }
+
+        $this->assertSame(['a', 'b'], $result);
+    }
+
+    public function testCountWithCacheDuringIteration(): void
+    {
+        $set = new ResultSet(new ArrayIterator(['a', 'b', 'c']));
+
+        $this->expectException(\RuntimeException::class);
+        $this->expectExceptionMessage('Cannot count result set while it is not fully iterated');
+
+        foreach ($set as $item) {
+            $set->count();
+        }
+    }
+
+    public function testLimitedCountWithCacheDuringIteration(): void
+    {
+        $set = new ResultSet(new ArrayIterator(['a', 'b', 'c']), 2);
+
+        $this->expectException(\RuntimeException::class);
+        $this->expectExceptionMessage('Cannot count result set while it is not fully iterated');
+
+        foreach ($set as $item) {
+            $set->count();
+        }
+    }
+
+    public function testEmptySetHasCountZero(): void
+    {
+        $this->assertSame(0, (new ResultSet(new ArrayIterator([])))->count());
+    }
 }