Add resourceMap helper for markdown extension (#151471)

This change introduces a `ResoruceMap` map type that is essentially `Map<vscode.Uri, T>`

It also fixes a potential race condition with `MdWorkspaceCache` where two quick calls would both trigger init

4 files changed, 90 insertions, 21 deletions

diff --git a/extensions/markdown-language-features/src/languageFeatures/references.ts b/extensions/markdown-language-features/src/languageFeatures/references.ts
index 7695d5e4ba3..5f63973d701 100644
--- a/extensions/markdown-language-features/src/languageFeatures/references.ts
+++ b/extensions/markdown-language-features/src/languageFeatures/references.ts
@@ -97,7 +97,7 @@ export class MdReferencesProvider extends Disposable implements vscode.Reference
 	}
 
 	private async getReferencesToHeader(document: SkinnyTextDocument, header: TocEntry): Promise<MdReference[]> {
-		const links = (await this._linkCache.getAll()).flat();
+		const links = (await this._linkCache.values()).flat();
 
 		const references: MdReference[] = [];
 
@@ -150,7 +150,7 @@ export class MdReferencesProvider extends Disposable implements vscode.Reference
 	}
 
 	private async getReferencesToLink(sourceLink: MdLink, triggerPosition: vscode.Position, token: vscode.CancellationToken): Promise<MdReference[]> {
-		const allLinksInWorkspace = (await this._linkCache.getAll()).flat();
+		const allLinksInWorkspace = (await this._linkCache.values()).flat();
 		if (token.isCancellationRequested) {
 			return [];
 		}
@@ -227,7 +227,7 @@ export class MdReferencesProvider extends Disposable implements vscode.Reference
 	}
 
 	public async getAllReferencesToFile(resource: vscode.Uri, _token: vscode.CancellationToken): Promise<MdReference[]> {
-		const allLinksInWorkspace = (await this._linkCache.getAll()).flat();
+		const allLinksInWorkspace = (await this._linkCache.values()).flat();
 		return Array.from(this.findAllLinksToFile(resource, allLinksInWorkspace, undefined));
 	}
 
diff --git a/extensions/markdown-language-features/src/languageFeatures/workspaceCache.ts b/extensions/markdown-language-features/src/languageFeatures/workspaceCache.ts
index e4039810c30..295771bfa60 100644
--- a/extensions/markdown-language-features/src/languageFeatures/workspaceCache.ts
+++ b/extensions/markdown-language-features/src/languageFeatures/workspaceCache.ts
@@ -6,6 +6,7 @@
 import * as vscode from 'vscode';
 import { Disposable } from '../util/dispose';
 import { Lazy, lazy } from '../util/lazy';
+import { ResourceMap } from '../util/resourceMap';
 import { MdWorkspaceContents, SkinnyTextDocument } from '../workspaceContents';
 
 /**
@@ -13,8 +14,8 @@ import { MdWorkspaceContents, SkinnyTextDocument } from '../workspaceContents';
  */
 export class MdWorkspaceCache<T> extends Disposable {
 
-	private readonly _cache = new Map<string, Lazy<Promise<T>>>();
-	private _hasPopulatedCache = false;
+	private readonly _cache = new ResourceMap<Lazy<Promise<T>>>();
+	private _init?: Promise<void>;
 
 	public constructor(
 		private readonly workspaceContents: MdWorkspaceContents,
@@ -23,19 +24,29 @@ export class MdWorkspaceCache<T> extends Disposable {
 		super();
 	}
 
-	public async getAll(): Promise<T[]> {
-		if (!this._hasPopulatedCache) {
-			await this.populateCache();
-			this._hasPopulatedCache = true;
-
-			this.workspaceContents.onDidChangeMarkdownDocument(this.onDidChangeDocument, this, this._disposables);
-			this.workspaceContents.onDidCreateMarkdownDocument(this.onDidChangeDocument, this, this._disposables);
-			this.workspaceContents.onDidDeleteMarkdownDocument(this.onDidDeleteDocument, this, this._disposables);
-		}
+	public async entries(): Promise<Array<[vscode.Uri, T]>> {
+		await this.ensureInit();
+		return Promise.all(Array.from(this._cache.entries(), async ([key, entry]) => {
+			return [key, await entry.value];
+		}));
+	}
 
+	public async values(): Promise<Array<T>> {
+		await this.ensureInit();
 		return Promise.all(Array.from(this._cache.values(), x => x.value));
 	}
 
+	private async ensureInit(): Promise<void> {
+		if (!this._init) {
+			this._init = this.populateCache();
+
+			this._register(this.workspaceContents.onDidChangeMarkdownDocument(this.onDidChangeDocument, this));
+			this._register(this.workspaceContents.onDidCreateMarkdownDocument(this.onDidChangeDocument, this));
+			this._register(this.workspaceContents.onDidDeleteMarkdownDocument(this.onDidDeleteDocument, this));
+		}
+		await this._init;
+	}
+
 	private async populateCache(): Promise<void> {
 		const markdownDocumentUris = await this.workspaceContents.getAllMarkdownDocuments();
 		for (const document of markdownDocumentUris) {
@@ -43,12 +54,8 @@ export class MdWorkspaceCache<T> extends Disposable {
 		}
 	}
 
-	private key(resource: vscode.Uri): string {
-		return resource.toString();
-	}
-
 	private update(document: SkinnyTextDocument): void {
-		this._cache.set(this.key(document.uri), lazy(() => this.getValue(document)));
+		this._cache.set(document.uri, lazy(() => this.getValue(document)));
 	}
 
 	private onDidChangeDocument(document: SkinnyTextDocument) {
@@ -56,6 +63,6 @@ export class MdWorkspaceCache<T> extends Disposable {
 	}
 
 	private onDidDeleteDocument(resource: vscode.Uri) {
-		this._cache.delete(this.key(resource));
+		this._cache.delete(resource);
 	}
 }
diff --git a/extensions/markdown-language-features/src/languageFeatures/workspaceSymbolProvider.ts b/extensions/markdown-language-features/src/languageFeatures/workspaceSymbolProvider.ts
index f7b1dcbfd3d..5906cc1cc72 100644
--- a/extensions/markdown-language-features/src/languageFeatures/workspaceSymbolProvider.ts
+++ b/extensions/markdown-language-features/src/languageFeatures/workspaceSymbolProvider.ts
@@ -23,7 +23,7 @@ export class MdWorkspaceSymbolProvider extends Disposable implements vscode.Work
 	}
 
 	public async provideWorkspaceSymbols(query: string): Promise<vscode.SymbolInformation[]> {
-		const allSymbols = (await this._cache.getAll()).flat();
+		const allSymbols = (await this._cache.values()).flat();
 		return allSymbols.filter(symbolInformation => symbolInformation.name.toLowerCase().indexOf(query.toLowerCase()) !== -1);
 	}
 }
diff --git a/extensions/markdown-language-features/src/util/resourceMap.ts b/extensions/markdown-language-features/src/util/resourceMap.ts
new file mode 100644
index 00000000000..35935314e55
--- /dev/null
+++ b/extensions/markdown-language-features/src/util/resourceMap.ts
@@ -0,0 +1,62 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as vscode from 'vscode';
+
+export class ResourceMap<T> {
+
+	private readonly map = new Map<string, { readonly uri: vscode.Uri; readonly value: T }>();
+
+	public set(uri: vscode.Uri, value: T): this {
+		this.map.set(this.toKey(uri), { uri, value });
+		return this;
+	}
+
+	public get(resource: vscode.Uri): T | undefined {
+		return this.map.get(this.toKey(resource))?.value;
+	}
+
+	public has(resource: vscode.Uri): boolean {
+		return this.map.has(this.toKey(resource));
+	}
+
+	public get size(): number {
+		return this.map.size;
+	}
+
+	public clear(): void {
+		this.map.clear();
+	}
+
+	public delete(resource: vscode.Uri): boolean {
+		return this.map.delete(this.toKey(resource));
+	}
+
+	public *values(): IterableIterator<T> {
+		for (const entry of this.map.values()) {
+			yield entry.value;
+		}
+	}
+
+	public *keys(): IterableIterator<vscode.Uri> {
+		for (const entry of this.map.values()) {
+			yield entry.uri;
+		}
+	}
+
+	public *entries(): IterableIterator<[vscode.Uri, T]> {
+		for (const entry of this.map.values()) {
+			yield [entry.uri, entry.value];
+		}
+	}
+
+	public [Symbol.iterator](): IterableIterator<[vscode.Uri, T]> {
+		return this.entries();
+	}
+
+	private toKey(resource: vscode.Uri) {
+		return resource.toString();
+	}
+}