added Iterables::memoize()

Author: dg

src/Utils/Iterables.php

@@ -156,6 +156,49 @@ public static function map(iterable $iterable, callable $transformer): \Generato
 	}
 
 
+	/**
+	 * Wraps around iterator and caches its keys and values during iteration.
+	 * This allows the data to be re-iterated multiple times.
+	 * @template K
+	 * @template V
+	 * @param  iterable<K, V>  $iterable
+	 * @return \IteratorAggregate<K, V>
+	 */
+	public static function memoize(iterable $iterable): iterable
+	{
+		return new class (self::toIterator($iterable)) implements \IteratorAggregate {
+			public function __construct(
+				private \Iterator $iterator,
+				private array $cache = [],
+			) {
+			}
+
+
+			public function getIterator(): \Generator
+			{
+				if (!$this->cache) {
+					$this->iterator->rewind();
+				}
+				$i = 0;
+				while (true) {
+					if (isset($this->cache[$i])) {
+						[$k, $v] = $this->cache[$i];
+					} elseif ($this->iterator->valid()) {
+						$k = $this->iterator->key();
+						$v = $this->iterator->current();
+						$this->iterator->next();
+						$this->cache[$i] = [$k, $v];
+					} else {
+						break;
+					}
+					yield $k => $v;
+					$i++;
+				}
+			}
+		};
+	}
+
+
 	/**
 	 * Creates an iterator from anything that is iterable.
 	 * @template K

tests/Utils/Iterables.memoize().phpt

@@ -0,0 +1,77 @@
+<?php
+
+/**
+ * Test: Nette\Utils\Iterables::memoize()
+ */
+
+declare(strict_types=1);
+
+use Nette\Utils\Iterables;
+use Tester\Assert;
+
+require __DIR__ . '/../bootstrap.php';
+
+
+function iterator(): Generator
+{
+	yield 'a' => 'apple';
+	yield ['b'] => ['banana'];
+	yield 'c' => 'cherry';
+}
+
+
+test('iteration', function () {
+	$iterator = Iterables::memoize(iterator());
+
+	$pairs = [];
+	foreach ($iterator as $key => $value) {
+		$pairs[] = [$key, $value];
+	}
+	Assert::same(
+		[
+			['a', 'apple'],
+			[['b'], ['banana']],
+			['c', 'cherry'],
+		],
+		$pairs,
+	);
+});
+
+
+test('re-iteration', function () {
+	$iterator = Iterables::memoize(iterator());
+
+	foreach ($iterator as $value);
+
+	$pairs = [];
+	foreach ($iterator as $key => $value) {
+		$pairs[] = [$key, $value];
+	}
+	Assert::same(
+		[
+			['a', 'apple'],
+			[['b'], ['banana']],
+			['c', 'cherry'],
+		],
+		$pairs,
+	);
+});
+
+
+test('nested re-iteration', function () {
+	$iterator = Iterables::memoize(iterator());
+
+	$pairs = [];
+	foreach ($iterator as $key => $value) {
+		$pairs[] = [$key, $value];
+		foreach ($iterator as $value);
+	}
+	Assert::same(
+		[
+			['a', 'apple'],
+			[['b'], ['banana']],
+			['c', 'cherry'],
+		],
+		$pairs,
+	);
+});