From 9d5d765fe412547790fb35601506febad86b7925 Mon Sep 17 00:00:00 2001
From: Jake Bailey <5341706+jakebailey@users.noreply.github.com>
Date: Wed, 1 Mar 2023 14:27:42 -0800
Subject: [PATCH] Replace occurrences with highlights

---
 src/language/typescript/languageFeatures.ts   | 26 ++++++++++++-------
 .../typescript/monaco.contribution.ts         |  9 +++----
 src/language/typescript/tsMode.ts             |  2 +-
 src/language/typescript/tsWorker.ts           |  9 ++++---
 4 files changed, 26 insertions(+), 20 deletions(-)

diff --git a/src/language/typescript/languageFeatures.ts b/src/language/typescript/languageFeatures.ts
index 15b0e7fe..1084aa73 100644
--- a/src/language/typescript/languageFeatures.ts
+++ b/src/language/typescript/languageFeatures.ts
@@ -715,7 +715,10 @@ export class QuickInfoAdapter extends Adapter implements languages.HoverProvider
 
 // --- occurrences ------
 
-export class OccurrencesAdapter extends Adapter implements languages.DocumentHighlightProvider {
+export class DocumentHighlightAdapter
+	extends Adapter
+	implements languages.DocumentHighlightProvider
+{
 	public async provideDocumentHighlights(
 		model: editor.ITextModel,
 		position: Position,
@@ -729,19 +732,24 @@ export class OccurrencesAdapter extends Adapter implements languages.DocumentHig
 			return;
 		}
 
-		const entries = await worker.getOccurrencesAtPosition(resource.toString(), offset);
+		const entries = await worker.getDocumentHighlights(resource.toString(), offset, [
+			resource.toString()
+		]);
 
 		if (!entries || model.isDisposed()) {
 			return;
 		}
 
-		return entries.map((entry) => {
-			return <languages.DocumentHighlight>{
-				range: this._textSpanToRange(model, entry.textSpan),
-				kind: entry.isWriteAccess
-					? languages.DocumentHighlightKind.Write
-					: languages.DocumentHighlightKind.Text
-			};
+		return entries.flatMap((entry) => {
+			return entry.highlightSpans.map((highlightSpans) => {
+				return <languages.DocumentHighlight>{
+					range: this._textSpanToRange(model, highlightSpans.textSpan),
+					kind:
+						highlightSpans.kind === 'writtenReference'
+							? languages.DocumentHighlightKind.Write
+							: languages.DocumentHighlightKind.Text
+				};
+			});
 		});
 	}
 }
diff --git a/src/language/typescript/monaco.contribution.ts b/src/language/typescript/monaco.contribution.ts
index 9b6a6ae8..9bd8ef38 100644
--- a/src/language/typescript/monaco.contribution.ts
+++ b/src/language/typescript/monaco.contribution.ts
@@ -440,13 +440,10 @@ export interface TypeScriptWorker {
 	 */
 	getQuickInfoAtPosition(fileName: string, position: number): Promise<any | undefined>;
 
-	/**
-	 * Get other ranges which are related to the item at the given position in the file (often used for highlighting).
-	 * @returns `Promise<ReadonlyArray<typescript.ReferenceEntry> | undefined>`
-	 */
-	getOccurrencesAtPosition(
+	getDocumentHighlights(
 		fileName: string,
-		position: number
+		position: number,
+		filesToSearch: string[]
 	): Promise<ReadonlyArray<any> | undefined>;
 
 	/**
diff --git a/src/language/typescript/tsMode.ts b/src/language/typescript/tsMode.ts
index 8cdab6fa..0f6d4d21 100644
--- a/src/language/typescript/tsMode.ts
+++ b/src/language/typescript/tsMode.ts
@@ -86,7 +86,7 @@ function setupMode(
 			providers.push(
 				languages.registerDocumentHighlightProvider(
 					modeId,
-					new languageFeatures.OccurrencesAdapter(worker)
+					new languageFeatures.DocumentHighlightAdapter(worker)
 				)
 			);
 		}
diff --git a/src/language/typescript/tsWorker.ts b/src/language/typescript/tsWorker.ts
index 65601f84..7cb83991 100644
--- a/src/language/typescript/tsWorker.ts
+++ b/src/language/typescript/tsWorker.ts
@@ -299,14 +299,15 @@ export class TypeScriptWorker implements ts.LanguageServiceHost, ITypeScriptWork
 		return this._languageService.getQuickInfoAtPosition(fileName, position);
 	}
 
-	async getOccurrencesAtPosition(
+	async getDocumentHighlights(
 		fileName: string,
-		position: number
-	): Promise<ReadonlyArray<ts.ReferenceEntry> | undefined> {
+		position: number,
+		filesToSearch: string[]
+	): Promise<ReadonlyArray<ts.DocumentHighlights> | undefined> {
 		if (fileNameIsLib(fileName)) {
 			return undefined;
 		}
-		return this._languageService.getOccurrencesAtPosition(fileName, position);
+		return this._languageService.getDocumentHighlights(fileName, position, filesToSearch);
 	}
 
 	async getDefinitionAtPosition(