From 45a45fb9b6e79380a84b9925dae519662c66e335 Mon Sep 17 00:00:00 2001
From: Mihai Ciuraru <mihai.ciuraru@gmail.com>
Date: Thu, 17 Nov 2022 01:41:50 +0200
Subject: [PATCH] Allow functions returning an iterator to be used as input

---
 linq.js          | 16 ++++++++++++++++
 test/iterator.js | 19 +++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/linq.js b/linq.js
index 0be1f96..d75a4bf 100644
--- a/linq.js
+++ b/linq.js
@@ -300,6 +300,22 @@ Enumerable.from = function (obj) {
                 Functions.Blank);
         });
     }
+    if (typeof obj == Types.Function && Object.keys(obj).length == 0) {
+        return new Enumerable(function () {
+            var orig;
+
+            return new IEnumerator(
+                function () {
+                    orig = obj()[Symbol.iterator]();
+                },
+                function () {
+                    var next = orig.next();
+                    return (next.done ? false : (this.yieldReturn(next.value)));
+                },
+                Functions.Blank);
+        });
+    }
+
     if (typeof obj != Types.Function) {
         // array or array-like object
         if (typeof obj.length == Types.Number) {
diff --git a/test/iterator.js b/test/iterator.js
index f568a92..91396f3 100644
--- a/test/iterator.js
+++ b/test/iterator.js
@@ -35,3 +35,22 @@ test("from Iterable", function () {
     }
     deepEqual(actual, ["abc", "def"]);
 });
+
+test("reusable iterator", function () {
+    const set = new Set([1, 2, 3])
+
+    let a = Enumerable.from(set.entries());
+
+    deepEqual(a.toArray(), [[1, 1], [2, 2], [3, 3]]);
+    deepEqual(a.toArray(), []);
+
+    let b = Enumerable.from(() => set.entries());
+
+    deepEqual(b.toArray(), [[1, 1], [2, 2], [3, 3]]);
+    deepEqual(b.toArray(), [[1, 1], [2, 2], [3, 3]]);
+
+    let c = Enumerable.from(() => ['x', 'y', 'z']);
+
+    deepEqual(c.toArray(), ['x', 'y', 'z']);
+    deepEqual(c.toArray(), ['x', 'y', 'z']);
+});