Move md workspace symbol search to language service (#154874)

* Move md workspace symbol search to language service

Also implements more of IWorkspace for the server

* Revert extra change

extensions/markdown-language-features/server/src/protocol.ts

@@ -0,0 +1,15 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { RequestType } from 'vscode-languageserver';
+import * as md from 'vscode-markdown-languageservice';
+
+declare const TextDecoder: any;
+
+export const parseRequestType: RequestType<{ uri: string }, md.Token[], any> = new RequestType('markdown/parse');
+
+export const readFileRequestType: RequestType<{ uri: string }, number[], any> = new RequestType('markdown/readFile');
+
+export const findFilesRequestTypes: RequestType<{}, string[], any> = new RequestType('markdown/findFiles');

extensions/markdown-language-features/server/src/server.ts

@@ -3,37 +3,15 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { Connection, Emitter, Event, InitializeParams, InitializeResult, RequestType, TextDocuments } from 'vscode-languageserver';
+import { Connection, InitializeParams, InitializeResult, TextDocuments } from 'vscode-languageserver';
 import { TextDocument } from 'vscode-languageserver-textdocument';
 import * as lsp from 'vscode-languageserver-types';
 import * as md from 'vscode-markdown-languageservice';
-import { URI } from 'vscode-uri';
 import { LogFunctionLogger } from './logging';
+import { parseRequestType } from './protocol';
+import { VsCodeClientWorkspace } from './workspace';
 
-
-const parseRequestType: RequestType<{ uri: string }, md.Token[], any> = new RequestType('markdown/parse');
-
-class TextDocumentToITextDocumentAdapter implements md.ITextDocument {
-	public readonly uri: md.IUri;
-
-	public get version(): number { return this._doc.version; }
-
-	public get lineCount(): number { return this._doc.lineCount; }
-
-	constructor(
-		private readonly _doc: TextDocument,
-	) {
-		this.uri = URI.parse(this._doc.uri);
-	}
-
-	getText(range?: md.IRange | undefined): string {
-		return this._doc.getText(range);
-	}
-
-	positionAt(offset: number): md.IPosition {
-		return this._doc.positionAt(offset);
-	}
-}
+declare const TextDecoder: any;
 
 export function startServer(connection: Connection) {
 	const documents = new TextDocuments(TextDocument);
@@ -45,11 +23,11 @@ export function startServer(connection: Connection) {
 				documentSymbolProvider: true,
 				foldingRangeProvider: true,
 				selectionRangeProvider: true,
+				workspaceSymbolProvider: true,
 			}
 		};
 	});
 
-
 	const parser = new class implements md.IMdParser {
 		slugifier = md.githubSlugifier;
 
@@ -58,42 +36,15 @@ export function startServer(connection: Connection) {
 		}
 	};
 
-	const workspace = new class implements md.IMdWorkspace {
-
-		private readonly _onDidChangeMarkdownDocument = new Emitter<md.ITextDocument>();
-		onDidChangeMarkdownDocument: Event<md.ITextDocument> = this._onDidChangeMarkdownDocument.event;
-
-		private readonly _onDidCreateMarkdownDocument = new Emitter<md.ITextDocument>();
-		onDidCreateMarkdownDocument: Event<md.ITextDocument> = this._onDidCreateMarkdownDocument.event;
-
-		private readonly _onDidDeleteMarkdownDocument = new Emitter<md.IUri>();
-		onDidDeleteMarkdownDocument: Event<md.IUri> = this._onDidDeleteMarkdownDocument.event;
-
-		async getAllMarkdownDocuments(): Promise<Iterable<md.ITextDocument>> {
-			return documents.all().map(doc => new TextDocumentToITextDocumentAdapter(doc));
-		}
-		hasMarkdownDocument(resource: md.IUri): boolean {
-			return !!documents.get(resource.toString());
-		}
-		async getOrLoadMarkdownDocument(_resource: md.IUri): Promise<md.ITextDocument | undefined> {
-			return undefined;
-		}
-		async pathExists(_resource: md.IUri): Promise<boolean> {
-			return false;
-		}
-		async readDirectory(_resource: md.IUri): Promise<[string, { isDir: boolean }][]> {
-			return [];
-		}
-	};
-
+	const workspace = new VsCodeClientWorkspace(connection, documents);
 	const logger = new LogFunctionLogger(connection.console.log.bind(connection.console));
-	const provider = md.createLanguageService(workspace, parser, logger);
+	const provider = md.createLanguageService({ workspace, parser, logger });
 
 	connection.onDocumentSymbol(async (params, token): Promise<lsp.DocumentSymbol[]> => {
 		try {
 			const document = documents.get(params.textDocument.uri);
 			if (document) {
-				return await provider.provideDocumentSymbols(new TextDocumentToITextDocumentAdapter(document), token);
+				return await provider.provideDocumentSymbols(document, token);
 			}
 		} catch (e) {
 			console.error(e.stack);
@@ -105,7 +56,7 @@ export function startServer(connection: Connection) {
 		try {
 			const document = documents.get(params.textDocument.uri);
 			if (document) {
-				return await provider.provideFoldingRanges(new TextDocumentToITextDocumentAdapter(document), token);
+				return await provider.provideFoldingRanges(document, token);
 			}
 		} catch (e) {
 			console.error(e.stack);
@@ -117,7 +68,7 @@ export function startServer(connection: Connection) {
 		try {
 			const document = documents.get(params.textDocument.uri);
 			if (document) {
-				return await provider.provideSelectionRanges(new TextDocumentToITextDocumentAdapter(document), params.positions, token);
+				return await provider.provideSelectionRanges(document, params.positions, token);
 			}
 		} catch (e) {
 			console.error(e.stack);
@@ -125,5 +76,15 @@ export function startServer(connection: Connection) {
 		return [];
 	});
 
+	connection.onWorkspaceSymbol(async (params, token): Promise<lsp.WorkspaceSymbol[]> => {
+		try {
+			return await provider.provideWorkspaceSymbols(params.query, token);
+		} catch (e) {
+			console.error(e.stack);
+		}
+		return [];
+	});
+
 	connection.listen();
 }
+

extensions/markdown-language-features/server/src/util/arrays.ts

@@ -0,0 +1,11 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+/**
+ * @returns New array with all falsy values removed. The original array IS NOT modified.
+ */
+export function coalesce<T>(array: ReadonlyArray<T | undefined | null>): T[] {
+	return <T[]>array.filter(e => !!e);
+}

extensions/markdown-language-features/server/src/util/file.ts

@@ -0,0 +1,27 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { TextDocument } from 'vscode-languageserver-textdocument';
+import * as URI from 'vscode-uri';
+
+const markdownFileExtensions = Object.freeze<string[]>([
+	'.md',
+	'.mkd',
+	'.mdwn',
+	'.mdown',
+	'.markdown',
+	'.markdn',
+	'.mdtxt',
+	'.mdtext',
+	'.workbook',
+]);
+
+export function looksLikeMarkdownPath(resolvedHrefPath: URI.URI) {
+	return markdownFileExtensions.includes(URI.Utils.extname(URI.URI.from(resolvedHrefPath)).toLowerCase());
+}
+
+export function isMarkdownDocument(document: TextDocument): boolean {
+	return document.languageId === 'markdown';
+}

extensions/markdown-language-features/server/src/util/limiter.ts

@@ -0,0 +1,67 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+interface ILimitedTaskFactory<T> {
+	factory: ITask<Promise<T>>;
+	c: (value: T | Promise<T>) => void;
+	e: (error?: unknown) => void;
+}
+
+interface ITask<T> {
+	(): T;
+}
+
+/**
+ * A helper to queue N promises and run them all with a max degree of parallelism. The helper
+ * ensures that at any time no more than M promises are running at the same time.
+ *
+ * Taken from 'src/vs/base/common/async.ts'
+ */
+export class Limiter<T> {
+
+	private _size = 0;
+	private runningPromises: number;
+	private readonly maxDegreeOfParalellism: number;
+	private readonly outstandingPromises: ILimitedTaskFactory<T>[];
+
+	constructor(maxDegreeOfParalellism: number) {
+		this.maxDegreeOfParalellism = maxDegreeOfParalellism;
+		this.outstandingPromises = [];
+		this.runningPromises = 0;
+	}
+
+	get size(): number {
+		return this._size;
+	}
+
+	queue(factory: ITask<Promise<T>>): Promise<T> {
+		this._size++;
+
+		return new Promise<T>((c, e) => {
+			this.outstandingPromises.push({ factory, c, e });
+			this.consume();
+		});
+	}
+
+	private consume(): void {
+		while (this.outstandingPromises.length && this.runningPromises < this.maxDegreeOfParalellism) {
+			const iLimitedTask = this.outstandingPromises.shift()!;
+			this.runningPromises++;
+
+			const promise = iLimitedTask.factory();
+			promise.then(iLimitedTask.c, iLimitedTask.e);
+			promise.then(() => this.consumed(), () => this.consumed());
+		}
+	}
+
+	private consumed(): void {
+		this._size--;
+		this.runningPromises--;
+
+		if (this.outstandingPromises.length > 0) {
+			this.consume();
+		}
+	}
+}

extensions/markdown-language-features/server/src/util/resourceMap.ts

@@ -0,0 +1,69 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { URI } from 'vscode-uri';
+
+
+type ResourceToKey = (uri: URI) => string;
+
+const defaultResourceToKey = (resource: URI): string => resource.toString();
+
+export class ResourceMap<T> {
+
+	private readonly map = new Map<string, { readonly uri: URI; readonly value: T }>();
+
+	private readonly toKey: ResourceToKey;
+
+	constructor(toKey: ResourceToKey = defaultResourceToKey) {
+		this.toKey = toKey;
+	}
+
+	public set(uri: URI, value: T): this {
+		this.map.set(this.toKey(uri), { uri, value });
+		return this;
+	}
+
+	public get(resource: URI): T | undefined {
+		return this.map.get(this.toKey(resource))?.value;
+	}
+
+	public has(resource: 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: 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<URI> {
+		for (const entry of this.map.values()) {
+			yield entry.uri;
+		}
+	}
+
+	public *entries(): IterableIterator<[URI, T]> {
+		for (const entry of this.map.values()) {
+			yield [entry.uri, entry.value];
+		}
+	}
+
+	public [Symbol.iterator](): IterableIterator<[URI, T]> {
+		return this.entries();
+	}
+}

extensions/markdown-language-features/server/src/workspace.ts

@@ -0,0 +1,155 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { Connection, Emitter, FileChangeType, TextDocuments } from 'vscode-languageserver';
+import { TextDocument } from 'vscode-languageserver-textdocument';
+import * as md from 'vscode-markdown-languageservice';
+import { URI } from 'vscode-uri';
+import * as protocol from './protocol';
+import { coalesce } from './util/arrays';
+import { isMarkdownDocument, looksLikeMarkdownPath } from './util/file';
+import { Limiter } from './util/limiter';
+import { ResourceMap } from './util/resourceMap';
+
+declare const TextDecoder: any;
+
+export class VsCodeClientWorkspace implements md.IWorkspace {
+
+	private readonly _onDidCreateMarkdownDocument = new Emitter<md.ITextDocument>();
+	public readonly onDidCreateMarkdownDocument = this._onDidCreateMarkdownDocument.event;
+
+	private readonly _onDidChangeMarkdownDocument = new Emitter<md.ITextDocument>();
+	public readonly onDidChangeMarkdownDocument = this._onDidChangeMarkdownDocument.event;
+
+	private readonly _onDidDeleteMarkdownDocument = new Emitter<URI>();
+	public readonly onDidDeleteMarkdownDocument = this._onDidDeleteMarkdownDocument.event;
+
+	private readonly _documentCache = new ResourceMap<md.ITextDocument>();
+
+	private readonly _utf8Decoder = new TextDecoder('utf-8');
+
+	constructor(
+		private readonly connection: Connection,
+		private readonly documents: TextDocuments<TextDocument>,
+	) {
+		documents.onDidOpen(e => {
+			this._documentCache.delete(URI.parse(e.document.uri));
+			if (this.isRelevantMarkdownDocument(e.document)) {
+				this._onDidCreateMarkdownDocument.fire(e.document);
+			}
+		});
+
+		documents.onDidChangeContent(e => {
+			if (this.isRelevantMarkdownDocument(e.document)) {
+				this._onDidChangeMarkdownDocument.fire(e.document);
+			}
+		});
+
+		documents.onDidClose(e => {
+			this._documentCache.delete(URI.parse(e.document.uri));
+		});
+
+		connection.onDidChangeWatchedFiles(async ({ changes }) => {
+			for (const change of changes) {
+				const resource = URI.parse(change.uri);
+				switch (change.type) {
+					case FileChangeType.Changed: {
+						this._documentCache.delete(resource);
+						const document = await this.getOrLoadMarkdownDocument(resource);
+						if (document) {
+							this._onDidChangeMarkdownDocument.fire(document);
+						}
+						break;
+					}
+					case FileChangeType.Created: {
+						const document = await this.getOrLoadMarkdownDocument(resource);
+						if (document) {
+							this._onDidCreateMarkdownDocument.fire(document);
+						}
+						break;
+					}
+					case FileChangeType.Deleted: {
+						this._documentCache.delete(resource);
+						this._onDidDeleteMarkdownDocument.fire(resource);
+						break;
+					}
+				}
+			}
+		});
+	}
+
+	async getAllMarkdownDocuments(): Promise<Iterable<md.ITextDocument>> {
+		const maxConcurrent = 20;
+
+		const foundFiles = new ResourceMap<void>();
+		const limiter = new Limiter<md.ITextDocument | undefined>(maxConcurrent);
+
+		// Add files on disk
+		const resources = await this.connection.sendRequest(protocol.findFilesRequestTypes, {});
+		const onDiskResults = await Promise.all(resources.map(strResource => {
+			return limiter.queue(async () => {
+				const resource = URI.parse(strResource);
+				const doc = await this.getOrLoadMarkdownDocument(resource);
+				if (doc) {
+					foundFiles.set(resource);
+				}
+				return doc;
+			});
+		}));
+
+		// Add opened files (such as untitled files)
+		const openTextDocumentResults = await Promise.all(this.documents.all()
+			.filter(doc => !foundFiles.has(URI.parse(doc.uri)) && this.isRelevantMarkdownDocument(doc)));
+
+		return coalesce([...onDiskResults, ...openTextDocumentResults]);
+	}
+
+	hasMarkdownDocument(resource: URI): boolean {
+		return !!this.documents.get(resource.toString());
+	}
+
+	async getOrLoadMarkdownDocument(resource: URI): Promise<md.ITextDocument | undefined> {
+		const existing = this._documentCache.get(resource);
+		if (existing) {
+			return existing;
+		}
+
+		const matchingDocument = this.documents.get(resource.toString());
+		if (matchingDocument) {
+			this._documentCache.set(resource, matchingDocument);
+			return matchingDocument;
+		}
+
+		if (!looksLikeMarkdownPath(resource)) {
+			return undefined;
+		}
+
+		try {
+			const response = await this.connection.sendRequest(protocol.readFileRequestType, { uri: resource.toString() });
+			// TODO: LSP doesn't seem to handle Array buffers well
+			const bytes = new Uint8Array(response);
+
+			// We assume that markdown is in UTF-8
+			const text = this._utf8Decoder.decode(bytes);
+			const doc = new md.InMemoryDocument(resource, text, 0);
+			this._documentCache.set(resource, doc);
+			return doc;
+		} catch (e) {
+			return undefined;
+		}
+	}
+
+	async pathExists(_resource: URI): Promise<boolean> {
+		return false;
+	}
+
+	async readDirectory(_resource: URI): Promise<[string, { isDir: boolean }][]> {
+		return [];
+	}
+
+	private isRelevantMarkdownDocument(doc: TextDocument) {
+		return isMarkdownDocument(doc) && URI.parse(doc.uri).scheme !== 'vscode-bulkeditpreview';
+	}
+}