Fix Map iterator mutation and add tests for mutation and exhaustion behavior

RealColdFry · RealColdFry · commit e52e21620960 · 2026-03-23T00:42:31.000Z
diff --git a/src/lualib/Map.ts b/src/lualib/Map.ts
@@ -113,46 +113,64 @@ export class Map<K extends AnyNotNil, V> {
     }
 
     public entries(): IterableIterator<[K, V]> {
+        const getFirstKey = () => this.firstKey;
         const { items, nextKey } = this;
-        let key = this.firstKey;
+        let key: K | undefined;
+        let started = false;
         return {
             [Symbol.iterator](): IterableIterator<[K, V]> {
                 return this;
             },
             next(): IteratorResult<[K, V]> {
-                const result = { done: !key, value: [key, items.get(key!)] as [K, V] };
-                key = nextKey.get(key!);
-                return result;
+                if (!started) {
+                    started = true;
+                    key = getFirstKey();
+                } else {
+                    key = nextKey.get(key!);
+                }
+                return { done: !key, value: [key!, items.get(key!)] as [K, V] };
             },
         };
     }
 
     public keys(): IterableIterator<K> {
+        const getFirstKey = () => this.firstKey;
         const nextKey = this.nextKey;
-        let key = this.firstKey;
+        let key: K | undefined;
+        let started = false;
         return {
             [Symbol.iterator](): IterableIterator<K> {
                 return this;
             },
             next(): IteratorResult<K> {
-                const result = { done: !key, value: key };
-                key = nextKey.get(key!);
-                return result as IteratorResult<K>;
+                if (!started) {
+                    started = true;
+                    key = getFirstKey();
+                } else {
+                    key = nextKey.get(key!);
+                }
+                return { done: !key, value: key! };
             },
         };
     }
 
     public values(): IterableIterator<V> {
+        const getFirstKey = () => this.firstKey;
         const { items, nextKey } = this;
-        let key = this.firstKey;
+        let key: K | undefined;
+        let started = false;
         return {
             [Symbol.iterator](): IterableIterator<V> {
                 return this;
             },
             next(): IteratorResult<V> {
-                const result = { done: !key, value: items.get(key!) };
-                key = nextKey.get(key!);
-                return result;
+                if (!started) {
+                    started = true;
+                    key = getFirstKey();
+                } else {
+                    key = nextKey.get(key!);
+                }
+                return { done: !key, value: items.get(key!) };
             },
         };
     }
diff --git a/test/unit/builtins/map.spec.ts b/test/unit/builtins/map.spec.ts
@@ -210,6 +210,44 @@ describe.each(iterationMethods)("map.%s() preserves insertion order", iterationM
     });
 });
 
+describe.each(iterationMethods)("map.%s() handles mutation", iterationMethod => {
+    test("iterator persists after delete", () => {
+        util.testFunction`
+            const map = new Map<number, string>([[1, "a"], [2, "b"]]);
+            const iter = map.${iterationMethod}();
+            map.delete(1);
+            return iter.next().value;
+        `.expectToMatchJsResult();
+    });
+
+    test("iterator with delete and add between iterations", () => {
+        util.testFunction`
+            const map = new Map<number, string>([[1, "a"], [2, "b"], [3, "c"]]);
+            const iter = map.${iterationMethod}();
+            iter.next(); // 1
+            map.delete(2);
+            map.set(4, "d");
+            const results: IteratorResult<any>[] = [];
+            let r = iter.next();
+            while (!r.done) { results.push({ done: r.done, value: r.value }); r = iter.next(); }
+            return results;
+        `.expectToMatchJsResult();
+    });
+
+    test("iterator does not restart after exhaustion", () => {
+        util.testFunction`
+            const map = new Map<number, string>([[1, "a"], [2, "b"]]);
+            const iter = map.${iterationMethod}();
+            const results: boolean[] = [];
+            results.push(iter.next().done!);
+            results.push(iter.next().done!);
+            results.push(iter.next().done!); // should be done
+            results.push(iter.next().done!); // should still be done, not restart
+            return results;
+        `.expectToMatchJsResult();
+    });
+});
+
 describe("Map.groupBy", () => {
     test("empty", () => {
         util.testFunction`
