Adopt ResourceMap in more places (#151475)

This changes switches to use the new `ResourceMap` type in more places in the markdown extension where we need to have a map/set with uris as the key

4 files changed, 49 insertions, 52 deletions

diff --git a/extensions/markdown-language-features/src/languageFeatures/diagnostics.ts b/extensions/markdown-language-features/src/languageFeatures/diagnostics.ts
index 4db32d614f7..ed9b174ffb3 100644
--- a/extensions/markdown-language-features/src/languageFeatures/diagnostics.ts
+++ b/extensions/markdown-language-features/src/languageFeatures/diagnostics.ts
@@ -16,6 +16,7 @@ import { MdWorkspaceContents, SkinnyTextDocument } from '../workspaceContents';
 import { InternalHref, LinkDefinitionSet, MdLink, MdLinkProvider, MdLinkSource } from './documentLinkProvider';
 import { tryFindMdDocumentForLink } from './references';
 import { CommandManager } from '../commandManager';
+import { ResourceMap } from '../util/resourceMap';
 
 const localize = nls.loadMessageBundle();
 
@@ -85,30 +86,28 @@ class VSCodeDiagnosticConfiguration extends Disposable implements DiagnosticConf
 
 class InflightDiagnosticRequests {
 
-	private readonly inFlightRequests = new Map<string, { readonly cts: vscode.CancellationTokenSource }>();
+	private readonly inFlightRequests = new ResourceMap<{ readonly cts: vscode.CancellationTokenSource }>();
 
 	public trigger(resource: vscode.Uri, compute: (token: vscode.CancellationToken) => Promise<void>) {
 		this.cancel(resource);
 
-		const key = this.getResourceKey(resource);
 		const cts = new vscode.CancellationTokenSource();
 		const entry = { cts };
-		this.inFlightRequests.set(key, entry);
+		this.inFlightRequests.set(resource, entry);
 
 		compute(cts.token).finally(() => {
-			if (this.inFlightRequests.get(key) === entry) {
-				this.inFlightRequests.delete(key);
+			if (this.inFlightRequests.get(resource) === entry) {
+				this.inFlightRequests.delete(resource);
 			}
 			cts.dispose();
 		});
 	}
 
 	public cancel(resource: vscode.Uri) {
-		const key = this.getResourceKey(resource);
-		const existing = this.inFlightRequests.get(key);
+		const existing = this.inFlightRequests.get(resource);
 		if (existing) {
 			existing.cts.cancel();
-			this.inFlightRequests.delete(key);
+			this.inFlightRequests.delete(resource);
 		}
 	}
 
@@ -122,10 +121,6 @@ class InflightDiagnosticRequests {
 		}
 		this.inFlightRequests.clear();
 	}
-
-	private getResourceKey(resource: vscode.Uri): string {
-		return resource.toString();
-	}
 }
 
 class LinkWatcher extends Disposable {
@@ -314,16 +309,16 @@ export class DiagnosticManager extends Disposable {
 		const allOpenedTabResources = this.getAllTabResources();
 		await Promise.all(
 			vscode.workspace.textDocuments
-				.filter(doc => allOpenedTabResources.has(doc.uri.toString()) && isMarkdownFile(doc))
+				.filter(doc => allOpenedTabResources.has(doc.uri) && isMarkdownFile(doc))
 				.map(doc => this.triggerDiagnostics(doc)));
 	}
 
-	private getAllTabResources() {
-		const openedTabDocs = new Map<string, vscode.Uri>();
+	private getAllTabResources(): ResourceMap<void> {
+		const openedTabDocs = new ResourceMap<void>();
 		for (const group of vscode.window.tabGroups.all) {
 			for (const tab of group.tabs) {
 				if (tab.input instanceof vscode.TabInputText) {
-					openedTabDocs.set(tab.input.uri.toString(), tab.input.uri);
+					openedTabDocs.set(tab.input.uri);
 				}
 			}
 		}
@@ -354,7 +349,7 @@ interface FileLinksData {
  */
 class FileLinkMap {
 
-	private readonly _filesToLinksMap = new Map<string, FileLinksData>();
+	private readonly _filesToLinksMap = new ResourceMap<FileLinksData>();
 
 	constructor(links: Iterable<MdLink>) {
 		for (const link of links) {
@@ -362,13 +357,12 @@ class FileLinkMap {
 				continue;
 			}
 
-			const fileKey = link.href.path.toString();
-			const existingFileEntry = this._filesToLinksMap.get(fileKey);
+			const existingFileEntry = this._filesToLinksMap.get(link.href.path);
 			const linkData = { source: link.source, fragment: link.href.fragment };
 			if (existingFileEntry) {
 				existingFileEntry.links.push(linkData);
 			} else {
-				this._filesToLinksMap.set(fileKey, { path: link.href.path, links: [linkData] });
+				this._filesToLinksMap.set(link.href.path, { path: link.href.path, links: [linkData] });
 			}
 		}
 	}
diff --git a/extensions/markdown-language-features/src/preview/topmostLineMonitor.ts b/extensions/markdown-language-features/src/preview/topmostLineMonitor.ts
index bc8364c8971..0f659f9e924 100644
--- a/extensions/markdown-language-features/src/preview/topmostLineMonitor.ts
+++ b/extensions/markdown-language-features/src/preview/topmostLineMonitor.ts
@@ -6,6 +6,7 @@
 import * as vscode from 'vscode';
 import { Disposable } from '../util/dispose';
 import { isMarkdownFile } from '../util/file';
+import { ResourceMap } from '../util/resourceMap';
 
 export interface LastScrollLocation {
 	readonly line: number;
@@ -14,10 +15,10 @@ export interface LastScrollLocation {
 
 export class TopmostLineMonitor extends Disposable {
 
-	private readonly pendingUpdates = new Map<string, number>();
+	private readonly pendingUpdates = new ResourceMap<number>();
 	private readonly throttle = 50;
-	private previousTextEditorInfo = new Map<string, LastScrollLocation>();
-	private previousStaticEditorInfo = new Map<string, LastScrollLocation>();
+	private previousTextEditorInfo = new ResourceMap<LastScrollLocation>();
+	private previousStaticEditorInfo = new ResourceMap<LastScrollLocation>();
 
 	constructor() {
 		super();
@@ -42,28 +43,28 @@ export class TopmostLineMonitor extends Disposable {
 	public readonly onDidChanged = this._onChanged.event;
 
 	public setPreviousStaticEditorLine(scrollLocation: LastScrollLocation): void {
-		this.previousStaticEditorInfo.set(scrollLocation.uri.toString(), scrollLocation);
+		this.previousStaticEditorInfo.set(scrollLocation.uri, scrollLocation);
 	}
 
 	public getPreviousStaticEditorLineByUri(resource: vscode.Uri): number | undefined {
-		const scrollLoc = this.previousStaticEditorInfo.get(resource.toString());
-		this.previousStaticEditorInfo.delete(resource.toString());
+		const scrollLoc = this.previousStaticEditorInfo.get(resource);
+		this.previousStaticEditorInfo.delete(resource);
 		return scrollLoc?.line;
 	}
 
 
 	public setPreviousTextEditorLine(scrollLocation: LastScrollLocation): void {
-		this.previousTextEditorInfo.set(scrollLocation.uri.toString(), scrollLocation);
+		this.previousTextEditorInfo.set(scrollLocation.uri, scrollLocation);
 	}
 
 	public getPreviousTextEditorLineByUri(resource: vscode.Uri): number | undefined {
-		const scrollLoc = this.previousTextEditorInfo.get(resource.toString());
-		this.previousTextEditorInfo.delete(resource.toString());
+		const scrollLoc = this.previousTextEditorInfo.get(resource);
+		this.previousTextEditorInfo.delete(resource);
 		return scrollLoc?.line;
 	}
 
 	public getPreviousStaticTextEditorLineByUri(resource: vscode.Uri): number | undefined {
-		const state = this.previousStaticEditorInfo.get(resource.toString());
+		const state = this.previousStaticEditorInfo.get(resource);
 		return state?.line;
 	}
 
@@ -71,21 +72,20 @@ export class TopmostLineMonitor extends Disposable {
 		resource: vscode.Uri,
 		line: number
 	) {
-		const key = resource.toString();
-		if (!this.pendingUpdates.has(key)) {
+		if (!this.pendingUpdates.has(resource)) {
 			// schedule update
 			setTimeout(() => {
-				if (this.pendingUpdates.has(key)) {
+				if (this.pendingUpdates.has(resource)) {
 					this._onChanged.fire({
 						resource,
-						line: this.pendingUpdates.get(key) as number
+						line: this.pendingUpdates.get(resource) as number
 					});
-					this.pendingUpdates.delete(key);
+					this.pendingUpdates.delete(resource);
 				}
 			}, this.throttle);
 		}
 
-		this.pendingUpdates.set(key, line);
+		this.pendingUpdates.set(resource, line);
 	}
 }
 
diff --git a/extensions/markdown-language-features/src/test/inMemoryWorkspace.ts b/extensions/markdown-language-features/src/test/inMemoryWorkspace.ts
index 4924578701d..53dab1a3e6d 100644
--- a/extensions/markdown-language-features/src/test/inMemoryWorkspace.ts
+++ b/extensions/markdown-language-features/src/test/inMemoryWorkspace.ts
@@ -5,15 +5,16 @@
 
 import * as assert from 'assert';
 import * as vscode from 'vscode';
+import { ResourceMap } from '../util/resourceMap';
 import { MdWorkspaceContents, SkinnyTextDocument } from '../workspaceContents';
 
 
 export class InMemoryWorkspaceMarkdownDocuments implements MdWorkspaceContents {
-	private readonly _documents = new Map<string, SkinnyTextDocument>();
+	private readonly _documents = new ResourceMap<SkinnyTextDocument>(uri => uri.fsPath);
 
 	constructor(documents: SkinnyTextDocument[]) {
 		for (const doc of documents) {
-			this._documents.set(this.getKey(doc.uri), doc);
+			this._documents.set(doc.uri, doc);
 		}
 	}
 
@@ -22,11 +23,11 @@ export class InMemoryWorkspaceMarkdownDocuments implements MdWorkspaceContents {
 	}
 
 	public async getMarkdownDocument(resource: vscode.Uri): Promise<SkinnyTextDocument | undefined> {
-		return this._documents.get(this.getKey(resource));
+		return this._documents.get(resource);
 	}
 
 	public async pathExists(resource: vscode.Uri): Promise<boolean> {
-		return this._documents.has(this.getKey(resource));
+		return this._documents.has(resource);
 	}
 
 	private readonly _onDidChangeMarkdownDocumentEmitter = new vscode.EventEmitter<SkinnyTextDocument>();
@@ -39,23 +40,19 @@ export class InMemoryWorkspaceMarkdownDocuments implements MdWorkspaceContents {
 	public onDidDeleteMarkdownDocument = this._onDidDeleteMarkdownDocumentEmitter.event;
 
 	public updateDocument(document: SkinnyTextDocument) {
-		this._documents.set(this.getKey(document.uri), document);
+		this._documents.set(document.uri, document);
 		this._onDidChangeMarkdownDocumentEmitter.fire(document);
 	}
 
 	public createDocument(document: SkinnyTextDocument) {
-		assert.ok(!this._documents.has(this.getKey(document.uri)));
+		assert.ok(!this._documents.has(document.uri));
 
-		this._documents.set(this.getKey(document.uri), document);
+		this._documents.set(document.uri, document);
 		this._onDidCreateMarkdownDocumentEmitter.fire(document);
 	}
 
 	public deleteDocument(resource: vscode.Uri) {
-		this._documents.delete(this.getKey(resource));
+		this._documents.delete(resource);
 		this._onDidDeleteMarkdownDocumentEmitter.fire(resource);
 	}
-
-	private getKey(resource: vscode.Uri): string {
-		return resource.fsPath;
-	}
 }
diff --git a/extensions/markdown-language-features/src/util/resourceMap.ts b/extensions/markdown-language-features/src/util/resourceMap.ts
index 35935314e55..7fe1f801bc6 100644
--- a/extensions/markdown-language-features/src/util/resourceMap.ts
+++ b/extensions/markdown-language-features/src/util/resourceMap.ts
@@ -5,10 +5,20 @@
 
 import * as vscode from 'vscode';
 
+type ResourceToKey = (uri: vscode.Uri) => string;
+
+const defaultResourceToKey = (resource: vscode.Uri): string => resource.toString();
+
 export class ResourceMap<T> {
 
 	private readonly map = new Map<string, { readonly uri: vscode.Uri; readonly value: T }>();
 
+	private readonly toKey: ResourceToKey;
+
+	constructor(toKey: ResourceToKey = defaultResourceToKey) {
+		this.toKey = toKey;
+	}
+
 	public set(uri: vscode.Uri, value: T): this {
 		this.map.set(this.toKey(uri), { uri, value });
 		return this;
@@ -55,8 +65,4 @@ export class ResourceMap<T> {
 	public [Symbol.iterator](): IterableIterator<[vscode.Uri, T]> {
 		return this.entries();
 	}
-
-	private toKey(resource: vscode.Uri) {
-		return resource.toString();
-	}
 }