map TernarySearchTree iterable

src/vs/base/common/map.ts

@@ -402,57 +402,44 @@ export class TernarySearchTree<K, V> {
 				if (!node.mid) {
 					return undefined;
 				} else {
-					return this._nodeIterator(node.mid);
+					return this._values(node.mid);
 				}
 			}
 		}
 		return undefined;
 	}
 
-	private _nodeIterator(node: TernarySearchTreeNode<K, V>): Iterator<V> {
-		let res: { done: false; value: V; };
-		let idx: number;
-		let data: V[];
-		const next = (): IteratorResult<V> => {
-			if (!data) {
-				// lazy till first invocation
-				data = [];
-				idx = 0;
-				this._forEach(node, value => data.push(value));
-			}
-			if (idx >= data.length) {
-				return { done: true, value: undefined };
-			}
-
-			if (!res) {
-				res = { done: false, value: data[idx++] };
-			} else {
-				res.value = data[idx++];
-			}
-			return res;
-		};
-		return { next };
+	forEach(callback: (value: V, index: K) => any): void {
+		for (const [key, value] of this) {
+			callback(value, key);
+		}
 	}
 
-	forEach(callback: (value: V, index: K) => any) {
-		this._forEach(this._root, callback);
+	*[Symbol.iterator](): IterableIterator<[K, V]> {
+		yield* this._entries(this._root);
 	}
 
-	private _forEach(node: TernarySearchTreeNode<K, V> | undefined, callback: (value: V, index: K) => any) {
+	private *_values(node: TernarySearchTreeNode<K, V>): IterableIterator<V> {
+		for (const [, value] of this._entries(node)) {
+			yield value;
+		}
+	}
+
+	private *_entries(node: TernarySearchTreeNode<K, V> | undefined): IterableIterator<[K, V]> {
 		if (node) {
 			// left
-			this._forEach(node.left, callback);
+			yield* this._entries(node.left);
 
 			// node
 			if (node.value) {
 				// callback(node.value, this._iter.join(parts));
-				callback(node.value, node.key);
+				yield [node.key, node.value];
 			}
 			// mid
-			this._forEach(node.mid, callback);
+			yield* this._entries(node.mid);
 
 			// right
-			this._forEach(node.right, callback);
+			yield* this._entries(node.right);
 		}
 	}
 }

src/vs/base/test/common/map.test.ts

@@ -434,11 +434,22 @@ suite('Map', () => {
 		map.forEach((value, key) => {
 			assert.equal(trie.get(key), value);
 		});
+
+		// forEach
+		let forEachCount = 0;
 		trie.forEach((element, key) => {
 			assert.equal(element, map.get(key));
-			map.delete(key);
+			forEachCount++;
 		});
-		assert.equal(map.size, 0);
+		assert.equal(map.size, forEachCount);
+
+		// iterator
+		let iterCount = 0;
+		for (let [key, value] of trie) {
+			assert.equal(value, map.get(key));
+			iterCount++;
+		}
+		assert.equal(map.size, iterCount);
 	}
 
 	test('TernarySearchTree - set', function () {