From dd47f2492bdcf001d62243766778c79357a5ce22 Mon Sep 17 00:00:00 2001
From: Andy <anhans@microsoft.com>
Date: Wed, 21 Feb 2018 10:02:34 -0800
Subject: [PATCH] getSemanticDocumentHighlights: Use `toMultiMap` helper
 (#22059)

* getSemanticDocumentHighlights: Use `toMultiMap` helper

* Rename to arrayToMultiMap and follow pattern of arrayToMap and arrayToNumericMap
---
 src/compiler/core.ts               | 34 ++++++++++++++++++------------
 src/services/documentHighlights.ts | 20 ++++--------------
 2 files changed, 24 insertions(+), 30 deletions(-)

diff --git a/src/compiler/core.ts b/src/compiler/core.ts
index 698824fa9b0..ecf667237c5 100644
--- a/src/compiler/core.ts
+++ b/src/compiler/core.ts
@@ -1333,20 +1333,20 @@ namespace ts {
      */
     export function arrayToMap<T>(array: ReadonlyArray<T>, makeKey: (value: T) => string): Map<T>;
     export function arrayToMap<T, U>(array: ReadonlyArray<T>, makeKey: (value: T) => string, makeValue: (value: T) => U): Map<U>;
-    export function arrayToMap<T, U>(array: ReadonlyArray<T>, makeKey: (value: T) => string, makeValue?: (value: T) => U): Map<T | U> {
+    export function arrayToMap<T, U>(array: ReadonlyArray<T>, makeKey: (value: T) => string, makeValue: (value: T) => T | U = identity): Map<T | U> {
         const result = createMap<T | U>();
         for (const value of array) {
-            result.set(makeKey(value), makeValue ? makeValue(value) : value);
+            result.set(makeKey(value), makeValue(value));
         }
         return result;
     }
 
     export function arrayToNumericMap<T>(array: ReadonlyArray<T>, makeKey: (value: T) => number): T[];
-    export function arrayToNumericMap<T, V>(array: ReadonlyArray<T>, makeKey: (value: T) => number, makeValue: (value: T) => V): V[];
-    export function arrayToNumericMap<T, V>(array: ReadonlyArray<T>, makeKey: (value: T) => number, makeValue?: (value: T) => V): V[] {
-        const result: V[] = [];
+    export function arrayToNumericMap<T, U>(array: ReadonlyArray<T>, makeKey: (value: T) => number, makeValue: (value: T) => U): U[];
+    export function arrayToNumericMap<T, U>(array: ReadonlyArray<T>, makeKey: (value: T) => number, makeValue: (value: T) => T | U = identity): (T | U)[] {
+        const result: (T | U)[] = [];
         for (const value of array) {
-            result[makeKey(value)] = makeValue ? makeValue(value) : value as any as V;
+            result[makeKey(value)] = makeValue(value);
         }
         return result;
     }
@@ -1362,6 +1362,20 @@ namespace ts {
         return arrayToMap<any, true>(array, makeKey || (s => s), () => true);
     }
 
+    export function arrayToMultiMap<T>(values: ReadonlyArray<T>, makeKey: (value: T) => string): MultiMap<T>;
+    export function arrayToMultiMap<T, U>(values: ReadonlyArray<T>, makeKey: (value: T) => string, makeValue: (value: T) => U): MultiMap<U>;
+    export function arrayToMultiMap<T, U>(values: ReadonlyArray<T>, makeKey: (value: T) => string, makeValue: (value: T) => T | U = identity): MultiMap<T | U> {
+        const result = createMultiMap<T | U>();
+        for (const value of values) {
+            result.add(makeKey(value), makeValue(value));
+        }
+        return result;
+    }
+
+    export function group<T>(values: ReadonlyArray<T>, getGroupId: (value: T) => string): ReadonlyArray<ReadonlyArray<T>> {
+        return arrayFrom(arrayToMultiMap(values, getGroupId).values());
+    }
+
     export function cloneMap(map: SymbolTable): SymbolTable;
     export function cloneMap<T>(map: ReadonlyMap<T>): Map<T>;
     export function cloneMap<T>(map: ReadonlyUnderscoreEscapedMap<T>): UnderscoreEscapedMap<T>;
@@ -1438,14 +1452,6 @@ namespace ts {
         }
     }
 
-    export function group<T>(values: ReadonlyArray<T>, getGroupId: (value: T) => string): ReadonlyArray<ReadonlyArray<T>> {
-        const groupIdToGroup = createMultiMap<T>();
-        for (const value of values) {
-            groupIdToGroup.add(getGroupId(value), value);
-        }
-        return arrayFrom(groupIdToGroup.values());
-    }
-
     /**
      * Tests whether a value is an array.
      */
diff --git a/src/services/documentHighlights.ts b/src/services/documentHighlights.ts
index dd5281f15bc..6d3094fc190 100644
--- a/src/services/documentHighlights.ts
+++ b/src/services/documentHighlights.ts
@@ -21,23 +21,11 @@ namespace ts.DocumentHighlights {
         };
     }
 
-    function getSemanticDocumentHighlights(position: number, node: Node, program: Program, cancellationToken: CancellationToken, sourceFilesToSearch: ReadonlyArray<SourceFile>): DocumentHighlights[] {
+    function getSemanticDocumentHighlights(position: number, node: Node, program: Program, cancellationToken: CancellationToken, sourceFilesToSearch: ReadonlyArray<SourceFile>): DocumentHighlights[] | undefined {
         const referenceEntries = FindAllReferences.getReferenceEntriesForNode(position, node, program, sourceFilesToSearch, cancellationToken);
-        return referenceEntries && convertReferencedSymbols(referenceEntries);
-    }
-
-    function convertReferencedSymbols(referenceEntries: ReadonlyArray<FindAllReferences.Entry>): DocumentHighlights[] {
-        const fileNameToDocumentHighlights = createMap<HighlightSpan[]>();
-        for (const entry of referenceEntries) {
-            const { fileName, span } = FindAllReferences.toHighlightSpan(entry);
-            let highlightSpans = fileNameToDocumentHighlights.get(fileName);
-            if (!highlightSpans) {
-                fileNameToDocumentHighlights.set(fileName, highlightSpans = []);
-            }
-            highlightSpans.push(span);
-        }
-
-        return arrayFrom(fileNameToDocumentHighlights.entries(), ([fileName, highlightSpans ]) => ({ fileName, highlightSpans }));
+        if (!referenceEntries) return undefined;
+        const map = arrayToMultiMap(referenceEntries.map(FindAllReferences.toHighlightSpan), e => e.fileName, e => e.span);
+        return arrayFrom(map.entries(), ([fileName, highlightSpans]) => ({ fileName, highlightSpans }));
     }
 
     function getSyntacticDocumentHighlights(node: Node, sourceFile: SourceFile): DocumentHighlights[] {