PHPLIB-1184: Improve CallbackIterator

src/Model/CallbackIterator.php

@@ -17,44 +17,54 @@
 
 namespace MongoDB\Model;
 
-use Closure;
 use Iterator;
 use IteratorIterator;
 use ReturnTypeWillChange;
 use Traversable;
 
+use function call_user_func;
+
 /**
  * Iterator to apply a callback before returning an element
  *
  * @internal
+ *
+ * @template TKey
+ * @template TValue
+ * @template TCallbackValue
+ * @template-implements Iterator<TKey, TCallbackValue>
  */
 class CallbackIterator implements Iterator
 {
-    /** @var Closure */
-    private $callback;
+    /** @var callable(TValue, TKey): TCallbackValue */
+    private $callable;
 
-    /** @var Iterator */
+    /** @var Iterator<TKey, TValue> */
     private $iterator;
 
-    public function __construct(Traversable $traversable, Closure $callback)
+    /**
+     * @param Traversable<TKey, TValue>              $traversable
+     * @param callable(TValue, TKey): TCallbackValue $callable
+     */
+    public function __construct(Traversable $traversable, callable $callable)
     {
         $this->iterator = $traversable instanceof Iterator ? $traversable : new IteratorIterator($traversable);
-        $this->callback = $callback;
+        $this->callable = $callable;
     }
 
     /**
      * @see https://php.net/iterator.current
-     * @return mixed
+     * @return TCallbackValue
      */
     #[ReturnTypeWillChange]
     public function current()
     {
-        return ($this->callback)($this->iterator->current());
+        return call_user_func($this->callable, $this->iterator->current(), $this->iterator->key());
     }
 
     /**
      * @see https://php.net/iterator.key
-     * @return mixed
+     * @return TKey
      */
     #[ReturnTypeWillChange]
     public function key()

tests/Model/CallbackIteratorTest.php

@@ -0,0 +1,70 @@
+<?php
+
+namespace MongoDB\Tests\Model;
+
+use ArrayIterator;
+use MongoDB\Model\CallbackIterator;
+use MongoDB\Tests\TestCase;
+
+use function array_keys;
+use function iterator_to_array;
+use function strrev;
+
+class CallbackIteratorTest extends TestCase
+{
+    public function testArrayIteration(): void
+    {
+        $expectedKey = 0;
+
+        $original = [1, 2, 3];
+
+        $callbackIterator = new CallbackIterator(
+            new ArrayIterator($original),
+            function ($value, $key) use (&$expectedKey) {
+                $this->assertSame($expectedKey, $key);
+                $expectedKey++;
+
+                return $value * 2;
+            }
+        );
+
+        $this->assertSame([2, 4, 6], iterator_to_array($callbackIterator));
+    }
+
+    public function testHashIteration(): void
+    {
+        $expectedKey = 0;
+
+        $original = ['a' => 1, 'b' => 2, 'c' => 3];
+        $expectedKeys = array_keys($original);
+
+        $callbackIterator = new CallbackIterator(
+            new ArrayIterator($original),
+            function ($value, $key) use (&$expectedKey, $expectedKeys) {
+                $this->assertSame($expectedKeys[$expectedKey], $key);
+                $expectedKey++;
+
+                return $value * 2;
+            }
+        );
+
+        $this->assertSame(['a' => 2, 'b' => 4, 'c' => 6], iterator_to_array($callbackIterator));
+    }
+
+    public function testWithCallable(): void
+    {
+        $original = ['foo', 'bar', 'baz'];
+
+        $callbackIterator = new CallbackIterator(
+            new ArrayIterator($original),
+            [self::class, 'reverseValue']
+        );
+
+        $this->assertSame(['oof', 'rab', 'zab'], iterator_to_array($callbackIterator));
+    }
+
+    public static function reverseValue($value, $key)
+    {
+        return strrev($value);
+    }
+}