diff --git a/src/lualib/Map.ts b/src/lualib/Map.ts index 16fa7e453..c2bcb43a4 100644 --- a/src/lualib/Map.ts +++ b/src/lualib/Map.ts @@ -113,46 +113,64 @@ export class Map { } 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 { + const getFirstKey = () => this.firstKey; const nextKey = this.nextKey; - let key = this.firstKey; + let key: K | undefined; + let started = false; return { [Symbol.iterator](): IterableIterator { return this; }, next(): IteratorResult { - const result = { done: !key, value: key }; - key = nextKey.get(key!); - return result as IteratorResult; + if (!started) { + started = true; + key = getFirstKey(); + } else { + key = nextKey.get(key!); + } + return { done: !key, value: key! }; }, }; } public values(): IterableIterator { + const getFirstKey = () => this.firstKey; const { items, nextKey } = this; - let key = this.firstKey; + let key: K | undefined; + let started = false; return { [Symbol.iterator](): IterableIterator { return this; }, next(): IteratorResult { - 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/src/lualib/Set.ts b/src/lualib/Set.ts index b46555edf..e302e4094 100644 --- a/src/lualib/Set.ts +++ b/src/lualib/Set.ts @@ -102,46 +102,64 @@ export class Set { } public entries(): IterableIterator<[T, T]> { + const getFirstKey = () => this.firstKey; const nextKey = this.nextKey; - let key: T = this.firstKey!; + let key: T | undefined; + let started = false; return { [Symbol.iterator](): IterableIterator<[T, T]> { return this; }, next(): IteratorResult<[T, T]> { - const result = { done: !key, value: [key, key] as [T, T] }; - key = nextKey.get(key); - return result; + if (!started) { + started = true; + key = getFirstKey(); + } else { + key = nextKey.get(key!); + } + return { done: !key, value: [key!, key!] as [T, T] }; }, }; } public keys(): IterableIterator { + const getFirstKey = () => this.firstKey; const nextKey = this.nextKey; - let key: T = this.firstKey!; + let key: T | undefined; + let started = false; return { [Symbol.iterator](): IterableIterator { return this; }, next(): IteratorResult { - const result = { done: !key, value: key }; - key = nextKey.get(key); - return result; + if (!started) { + started = true; + key = getFirstKey(); + } else { + key = nextKey.get(key!); + } + return { done: !key, value: key! }; }, }; } public values(): IterableIterator { + const getFirstKey = () => this.firstKey; const nextKey = this.nextKey; - let key: T = this.firstKey!; + let key: T | undefined; + let started = false; return { [Symbol.iterator](): IterableIterator { return this; }, next(): IteratorResult { - const result = { done: !key, value: key }; - key = nextKey.get(key); - return result; + if (!started) { + started = true; + key = getFirstKey(); + } else { + key = nextKey.get(key!); + } + return { done: !key, value: key! }; }, }; } diff --git a/test/unit/builtins/map.spec.ts b/test/unit/builtins/map.spec.ts index f2e1cbfbe..27b539589 100644 --- 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([[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([[1, "a"], [2, "b"], [3, "c"]]); + const iter = map.${iterationMethod}(); + iter.next(); // 1 + map.delete(2); + map.set(4, "d"); + const results: IteratorResult[] = []; + 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([[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` diff --git a/test/unit/builtins/set.spec.ts b/test/unit/builtins/set.spec.ts index a03f27ab2..65ee4aef0 100644 --- a/test/unit/builtins/set.spec.ts +++ b/test/unit/builtins/set.spec.ts @@ -195,6 +195,48 @@ describe.each(iterationMethods)("set.%s() preserves insertion order", iterationM }); }); +describe.each(iterationMethods)("set.%s() handles mutation", iterationMethod => { + test("iterator persists after delete", () => { + util.testFunction` + const set1 = new Set(); + set1.add(42); + set1.add("forty two"); + + const iterator1 = set1.${iterationMethod}(); + set1.delete(42); + + return iterator1.next().value; + `.expectToMatchJsResult(); + }); + + test("iterator with delete and add between iterations", () => { + util.testFunction` + const set = new Set([1, 2, 3]); + const iter = set.${iterationMethod}(); + iter.next(); // 1 + set.delete(2); + set.add(4); + const results: IteratorResult[] = []; + 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 set = new Set([1, 2]); + const iter = set.${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(); + }); +}); + test("instanceof Set without creating set", () => { util.testFunction` const myset = 3 as any;