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/set.spec.ts b/test/unit/builtins/set.spec.ts index a03f27ab2..cbf4ff418 100644 --- a/test/unit/builtins/set.spec.ts +++ b/test/unit/builtins/set.spec.ts @@ -195,6 +195,19 @@ describe.each(iterationMethods)("set.%s() preserves insertion order", iterationM }); }); +test("set iterator persists after delete", () => { + util.testFunction` + const set1 = new Set(); + set1.add(42); + set1.add("forty two"); + + const iterator1 = set1.values(); + set1.delete(42); + + return iterator1.next().value; + `.expectToMatchJsResult(); +}); + test("instanceof Set without creating set", () => { util.testFunction` const myset = 3 as any;