Align the ShimMap iterator behavior with native Maps

src/compiler/core.ts

@@ -118,42 +118,105 @@ namespace ts {
     export const MapCtr = typeof Map !== "undefined" && "entries" in Map.prototype ? Map : shimMap();
 
     // Keep the class inside a function so it doesn't get compiled if it's not used.
-    function shimMap(): new <T>() => Map<T> {
+    export function shimMap(): new <T>() => Map<T> {
+
+        interface MapEntry<T> {
+            readonly key?: string;
+            value?: T;
+
+            // Linked list references for iterators.
+            nextEntry?: MapEntry<T>;
+            previousEntry?: MapEntry<T>;
+
+            /**
+             * Specifies if iterators should skip the next entry.
+             * This will be set when an entry is deleted.
+             * See https://github.com/Microsoft/TypeScript/pull/27292 for more information.
+             */
+            skipNext?: boolean;
+        }
 
         class MapIterator<T, U extends (string | T | [string, T])> {
-            private data: MapLike<T>;
-            private keys: ReadonlyArray<string>;
-            private index = 0;
-            private selector: (data: MapLike<T>, key: string) => U;
-            constructor(data: MapLike<T>, selector: (data: MapLike<T>, key: string) => U) {
-                this.data = data;
+            private currentEntry?: MapEntry<T>;
+            private selector: (key: string, value: T) => U;
+
+            constructor(currentEntry: MapEntry<T>, selector: (key: string, value: T) => U) {
+                this.currentEntry = currentEntry;
                 this.selector = selector;
-                this.keys = Object.keys(data);
             }
 
             public next(): { value: U, done: false } | { value: never, done: true } {
-                const index = this.index;
-                if (index < this.keys.length) {
-                    this.index++;
-                    return { value: this.selector(this.data, this.keys[index]), done: false };
+                // Navigate to the next entry.
+                while (this.currentEntry) {
+                    const skipNext = !!this.currentEntry.skipNext;
+                    this.currentEntry = this.currentEntry.nextEntry;
+
+                    if (!skipNext) {
+                        break;
+                    }
+                }
+
+                if (this.currentEntry) {
+                    return { value: this.selector(this.currentEntry.key!, this.currentEntry.value!), done: false };
+                }
+                else {
+                    return { value: undefined as never, done: true };
                 }
-                return { value: undefined as never, done: true };
             }
         }
 
         return class <T> implements Map<T> {
-            private data = createDictionaryObject<T>();
+            private data = createDictionaryObject<MapEntry<T>>();
             public size = 0;
 
+            // Linked list references for iterators.
+            // See https://github.com/Microsoft/TypeScript/pull/27292
+            // for more information.
+
+            /**
+             * The first entry in the linked list.
+             * Note that this is only a stub that serves as starting point
+             * for iterators and doesn't contain a key and a value.
+             */
+            private readonly firstEntry: MapEntry<T>;
+            private lastEntry: MapEntry<T>;
+
+            constructor() {
+                // Create a first (stub) map entry that will not contain a key
+                // and value but serves as starting point for iterators.
+                this.firstEntry = {};
+                // When the map is empty, the last entry is the same as the
+                // first one.
+                this.lastEntry = this.firstEntry;
+            }
+
             get(key: string): T | undefined {
-                return this.data[key];
+                const entry = this.data[key] as MapEntry<T> | undefined;
+                return entry && entry.value!;
             }
 
             set(key: string, value: T): this {
                 if (!this.has(key)) {
                     this.size++;
+
+                    // Create a new entry that will be appended at the
+                    // end of the linked list.
+                    const newEntry: MapEntry<T> = {
+                        key,
+                        value
+                    };
+                    this.data[key] = newEntry;
+
+                    // Adjust the references.
+                    const previousLastEntry = this.lastEntry;
+                    previousLastEntry.nextEntry = newEntry;
+                    newEntry.previousEntry = previousLastEntry;
+                    this.lastEntry = newEntry;
                 }
-                this.data[key] = value;
+                else {
+                    this.data[key].value = value;
+                }
+
                 return this;
             }
 
@@ -165,32 +228,81 @@ namespace ts {
             delete(key: string): boolean {
                 if (this.has(key)) {
                     this.size--;
+                    const entry = this.data[key];
                     delete this.data[key];
+
+                    // Adjust the linked list references of the neighbor entries.
+                    const previousEntry = entry.previousEntry!;
+                    previousEntry.nextEntry = entry.nextEntry;
+                    if (entry.nextEntry) {
+                        entry.nextEntry.previousEntry = previousEntry;
+                    }
+
+                    // When the deleted entry was the last one, we need to
+                    // adust the lastEntry reference.
+                    if (this.lastEntry === entry) {
+                        this.lastEntry = previousEntry;
+                    }
+
+                    // Adjust the forward reference of the deleted entry
+                    // in case an iterator still references it. This allows us
+                    // to throw away the entry, but when an active iterator
+                    // (which points to the current entry) continues, it will
+                    // navigate to the entry that originally came before the
+                    // current one and skip it.
+                    entry.previousEntry = undefined;
+                    entry.nextEntry = previousEntry;
+                    entry.skipNext = true;
+
                     return true;
                 }
                 return false;
             }
 
             clear(): void {
-                this.data = createDictionaryObject<T>();
+                this.data = createDictionaryObject<MapEntry<T>>();
                 this.size = 0;
+
+                // Reset the linked list. Note that we must adjust the forward
+                // references of the deleted entries to ensure iterators stuck
+                // in the middle of the list don't continue with deleted entries,
+                // but can continue with new entries added after the clear()
+                // operation.
+                const firstEntry = this.firstEntry;
+                let currentEntry = firstEntry.nextEntry;
+                while (currentEntry) {
+                    const nextEntry = currentEntry.nextEntry;
+                    currentEntry.previousEntry = undefined;
+                    currentEntry.nextEntry = firstEntry;
+                    currentEntry.skipNext = true;
+
+                    currentEntry = nextEntry;
+                }
+                firstEntry.nextEntry = undefined;
+                this.lastEntry = firstEntry;
             }
 
             keys(): Iterator<string> {
-                return new MapIterator(this.data, (_data, key) => key);
+                return new MapIterator(this.firstEntry, key => key);
             }
 
             values(): Iterator<T> {
-                return new MapIterator(this.data, (data, key) => data[key]);
+                return new MapIterator(this.firstEntry, (_key, value) => value);
             }
 
             entries(): Iterator<[string, T]> {
-                return new MapIterator(this.data, (data, key) => [key, data[key]] as [string, T]);
+                return new MapIterator(this.firstEntry, (key, value) => [key, value] as [string, T]);
             }
 
             forEach(action: (value: T, key: string) => void): void {
-                for (const key in this.data) {
-                    action(this.data[key], key);
+                const iterator = this.entries();
+                while (true) {
+                    const { value: entry, done } = iterator.next();
+                    if (done) {
+                        break;
+                    }
+
+                    action(entry[1], entry[0]);
                 }
             }
         };

src/testRunner/tsconfig.json

@@ -58,6 +58,7 @@
         "unittests/publicApi.ts",
         "unittests/reuseProgramStructure.ts",
         "unittests/semver.ts",
+        "unittests/shimMap.ts",
         "unittests/transform.ts",
         "unittests/tsbuild.ts",
         "unittests/tsbuildWatchMode.ts",

src/testRunner/unittests/shimMap.ts

@@ -0,0 +1,106 @@
+namespace ts {
+    describe("unittests:: shimMap", () => {
+
+        function testMapIterationAddedValues(map: Map<string>, useForEach: boolean): string {
+            let resultString = "";
+
+            map.set("1", "1");
+            map.set("3", "3");
+            map.set("2", "2");
+            map.set("4", "4");
+
+            let addedThree = false;
+            const doForEach = (value: string, key: string) => {
+                resultString += `${key}:${value};`;
+
+                // Add a new key ("0") - the map should provide this
+                // one in the next iteration.
+                if (key === "1") {
+                    map.set("1", "X1");
+                    map.set("0", "X0");
+                    map.set("4", "X4");
+                }
+                else if (key === "3") {
+                    if (!addedThree) {
+                        addedThree = true;
+
+                        // Remove and re-add key "3"; the map should
+                        // visit it after "0".
+                        map.delete("3");
+                        map.set("3", "Y3");
+
+                        // Change the value of "2"; the map should provide
+                        // it when visiting the key.
+                        map.set("2", "Y2");
+                    }
+                    else {
+                        // Check that an entry added when we visit the
+                        // currently last entry will still be visited.
+                        map.set("999", "999");
+                    }
+                }
+                else if (key === "999") {
+                    // Ensure that clear() behaves correctly same as removing all keys.
+                    map.set("A", "A");
+                    map.set("B", "B");
+                    map.set("C", "C");
+                }
+                else if (key === "A") {
+                    map.clear();
+                    map.set("Z", "Z");
+                }
+                else if (key === "Z") {
+                    // Check that the map behaves correctly when two items are
+                    // added and removed immediately.
+                    map.set("X", "X");
+                    map.set("X1", "X1");
+                    map.set("X2", "X2");
+                    map.delete("X1");
+                    map.delete("X2");
+                    map.set("Y", "Y");
+                }
+            };
+
+            if (useForEach) {
+                map.forEach(doForEach);
+            }
+            else {
+                // Use an iterator.
+                const iterator = map.entries();
+                while (true) {
+                    const { value: tuple, done } = iterator.next();
+                    if (done) {
+                        break;
+                    }
+
+                    doForEach(tuple[1], tuple[0]);
+                }
+            }
+
+            return resultString;
+        }
+
+        it("iterates values in insertion order and handles changes", () => {
+            const expectedResult = "1:1;3:3;2:Y2;4:X4;0:X0;3:Y3;999:999;A:A;Z:Z;X:X;Y:Y;";
+
+            // First, ensure the test actually has the same behavior as a native Map.
+            let nativeMap = createMap<string>();
+            const nativeMapForEachResult = testMapIterationAddedValues(nativeMap, /* useForEach */ true);
+            assert.equal(nativeMapForEachResult, expectedResult, "nativeMap-forEach");
+
+            nativeMap = createMap<string>();
+            const nativeMapIteratorResult = testMapIterationAddedValues(nativeMap, /* useForEach */ false);
+            assert.equal(nativeMapIteratorResult, expectedResult, "nativeMap-iterator");
+
+            // Then, test the shimMap.
+            let localShimMap = new (shimMap())<string>();
+            const shimMapForEachResult = testMapIterationAddedValues(localShimMap, /* useForEach */ true);
+            assert.equal(shimMapForEachResult, expectedResult, "shimMap-forEach");
+
+            localShimMap = new (shimMap())<string>();
+            const shimMapIteratorResult = testMapIterationAddedValues(localShimMap, /* useForEach */ false);
+            assert.equal(shimMapIteratorResult, expectedResult, "shimMap-iterator");
+
+        });
+    });
+}