Reduce number of times MD docs are re-tokenized (#152674)

This change reduces the number of times we retokenize a markdown file by doing the following:

- Use `MdTableOfContentsProvider` in more places
- Introduce a `IMarkdownParser` interface that lets us drop in a caching version of the tokenizer

extensions/markdown-language-features/src/languageFeatures/diagnostics.ts

@@ -7,7 +7,6 @@ import * as picomatch from 'picomatch';
 import * as vscode from 'vscode';
 import * as nls from 'vscode-nls';
 import { CommandManager } from '../commandManager';
-import { MarkdownEngine } from '../markdownEngine';
 import { MdTableOfContentsProvider } from '../tableOfContents';
 import { MdTableOfContentsWatcher } from '../test/tableOfContentsWatcher';
 import { Delayer } from '../util/async';
@@ -305,12 +304,12 @@ export class DiagnosticManager extends Disposable {
 	public readonly ready: Promise<void>;
 
 	constructor(
-		engine: MarkdownEngine,
 		private readonly workspaceContents: MdWorkspaceContents,
 		private readonly computer: DiagnosticComputer,
 		private readonly configuration: DiagnosticConfiguration,
 		private readonly reporter: DiagnosticReporter,
 		private readonly referencesProvider: MdReferencesProvider,
+		tocProvider: MdTableOfContentsProvider,
 		delay = 300,
 	) {
 		super();
@@ -346,7 +345,7 @@ export class DiagnosticManager extends Disposable {
 			}
 		}));
 
-		this.tableOfContentsWatcher = this._register(new MdTableOfContentsWatcher(engine, workspaceContents));
+		this.tableOfContentsWatcher = this._register(new MdTableOfContentsWatcher(workspaceContents, tocProvider));
 		this._register(this.tableOfContentsWatcher.onTocChanged(async e => {
 			// When the toc of a document changes, revalidate every file that linked to it too
 			const triggered = new ResourceMap<void>();
@@ -638,7 +637,6 @@ class AddToIgnoreLinksQuickFixProvider implements vscode.CodeActionProvider {
 
 export function registerDiagnosticSupport(
 	selector: vscode.DocumentSelector,
-	engine: MarkdownEngine,
 	workspaceContents: MdWorkspaceContents,
 	linkProvider: MdLinkProvider,
 	commandManager: CommandManager,
@@ -647,12 +645,12 @@ export function registerDiagnosticSupport(
 ): vscode.Disposable {
 	const configuration = new VSCodeDiagnosticConfiguration();
 	const manager = new DiagnosticManager(
-		engine,
 		workspaceContents,
 		new DiagnosticComputer(workspaceContents, linkProvider, tocProvider),
 		configuration,
 		new DiagnosticCollectionReporter(),
-		referenceProvider);
+		referenceProvider,
+		tocProvider);
 	return vscode.Disposable.from(
 		configuration,
 		manager,

extensions/markdown-language-features/src/languageFeatures/documentLinkProvider.ts

@@ -7,7 +7,7 @@ import * as vscode from 'vscode';
 import * as nls from 'vscode-nls';
 import * as uri from 'vscode-uri';
 import { OpenDocumentLinkCommand } from '../commands/openDocumentLink';
-import { MarkdownEngine } from '../markdownEngine';
+import { IMdParser } from '../markdownEngine';
 import { coalesce } from '../util/arrays';
 import { noopToken } from '../util/cancellation';
 import { Disposable } from '../util/dispose';
@@ -233,8 +233,8 @@ const definitionPattern = /^([\t ]*\[(?!\^)((?:\\\]|[^\]])+)\]:\s*)([^<]\S*|<[^>
 const inlineCodePattern = /(?:^|[^`])(`+)(?:.+?|.*?(?:(?:\r?\n).+?)*?)(?:\r?\n)?\1(?:$|[^`])/gm;
 
 class NoLinkRanges {
-	public static async compute(document: SkinnyTextDocument, engine: MarkdownEngine): Promise<NoLinkRanges> {
-		const tokens = await engine.parse(document);
+	public static async compute(tokenizer: IMdParser, document: SkinnyTextDocument): Promise<NoLinkRanges> {
+		const tokens = await tokenizer.tokenize(document);
 		const multiline = tokens.filter(t => (t.type === 'code_block' || t.type === 'fence' || t.type === 'html_block') && !!t.map).map(t => t.map) as [number, number][];
 
 		const text = document.getText();
@@ -270,11 +270,11 @@ class NoLinkRanges {
 export class MdLinkComputer {
 
 	constructor(
-		private readonly engine: MarkdownEngine
+		private readonly tokenizer: IMdParser,
 	) { }
 
 	public async getAllLinks(document: SkinnyTextDocument, token: vscode.CancellationToken): Promise<MdLink[]> {
-		const noLinkRanges = await NoLinkRanges.compute(document, this.engine);
+		const noLinkRanges = await NoLinkRanges.compute(this.tokenizer, document);
 		if (token.isCancellationRequested) {
 			return [];
 		}
@@ -436,11 +436,11 @@ export class MdLinkProvider extends Disposable {
 	private readonly linkComputer: MdLinkComputer;
 
 	constructor(
-		engine: MarkdownEngine,
+		tokenizer: IMdParser,
 		workspaceContents: MdWorkspaceContents,
 	) {
 		super();
-		this.linkComputer = new MdLinkComputer(engine);
+		this.linkComputer = new MdLinkComputer(tokenizer);
 		this._linkCache = this._register(new MdDocumentInfoCache(workspaceContents, doc => this.linkComputer.getAllLinks(doc, noopToken)));
 	}
 

extensions/markdown-language-features/src/languageFeatures/foldingProvider.ts

@@ -3,9 +3,9 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import Token = require('markdown-it/lib/token');
+import type Token = require('markdown-it/lib/token');
 import * as vscode from 'vscode';
-import { MarkdownEngine } from '../markdownEngine';
+import { IMdParser } from '../markdownEngine';
 import { MdTableOfContentsProvider } from '../tableOfContents';
 import { SkinnyTextDocument } from '../workspaceContents';
 
@@ -18,7 +18,7 @@ interface MarkdownItTokenWithMap extends Token {
 export class MdFoldingProvider implements vscode.FoldingRangeProvider {
 
 	constructor(
-		private readonly engine: MarkdownEngine,
+		private readonly parser: IMdParser,
 		private readonly tocProvide: MdTableOfContentsProvider,
 	) { }
 
@@ -36,7 +36,7 @@ export class MdFoldingProvider implements vscode.FoldingRangeProvider {
 	}
 
 	private async getRegions(document: SkinnyTextDocument): Promise<vscode.FoldingRange[]> {
-		const tokens = await this.engine.parse(document);
+		const tokens = await this.parser.tokenize(document);
 		const regionMarkers = tokens.filter(isRegionMarker)
 			.map(token => ({ line: token.map[0], isStart: isStartRegion(token.content) }));
 
@@ -67,7 +67,7 @@ export class MdFoldingProvider implements vscode.FoldingRangeProvider {
 	}
 
 	private async getBlockFoldingRanges(document: SkinnyTextDocument): Promise<vscode.FoldingRange[]> {
-		const tokens = await this.engine.parse(document);
+		const tokens = await this.parser.tokenize(document);
 		const multiLineListItems = tokens.filter(isFoldableToken);
 		return multiLineListItems.map(listItem => {
 			const start = listItem.map[0];
@@ -115,8 +115,8 @@ const isFoldableToken = (token: Token): token is MarkdownItTokenWithMap => {
 
 export function registerFoldingSupport(
 	selector: vscode.DocumentSelector,
-	engine: MarkdownEngine,
+	parser: IMdParser,
 	tocProvider: MdTableOfContentsProvider,
 ): vscode.Disposable {
-	return vscode.languages.registerFoldingRangeProvider(selector, new MdFoldingProvider(engine, tocProvider));
+	return vscode.languages.registerFoldingRangeProvider(selector, new MdFoldingProvider(parser, tocProvider));
 }

extensions/markdown-language-features/src/languageFeatures/pathCompletions.ts

@@ -5,7 +5,7 @@
 
 import { dirname, resolve } from 'path';
 import * as vscode from 'vscode';
-import { MarkdownEngine } from '../markdownEngine';
+import { IMdParser } from '../markdownEngine';
 import { TableOfContents } from '../tableOfContents';
 import { resolveUriToMarkdownFile } from '../util/openDocumentLink';
 import { SkinnyTextDocument } from '../workspaceContents';
@@ -82,7 +82,7 @@ function tryDecodeUriComponent(str: string): string {
 export class MdVsCodePathCompletionProvider implements vscode.CompletionItemProvider {
 
 	constructor(
-		private readonly engine: MarkdownEngine,
+		private readonly parser: IMdParser,
 		private readonly linkProvider: MdLinkProvider,
 	) { }
 
@@ -249,7 +249,7 @@ export class MdVsCodePathCompletionProvider implements vscode.CompletionItemProv
 	}
 
 	private async *provideHeaderSuggestions(document: SkinnyTextDocument, position: vscode.Position, context: CompletionContext, insertionRange: vscode.Range): AsyncIterable<vscode.CompletionItem> {
-		const toc = await TableOfContents.createForDocumentOrNotebook(this.engine, document);
+		const toc = await TableOfContents.createForDocumentOrNotebook(this.parser, document);
 		for (const entry of toc.entries) {
 			const replacementRange = new vscode.Range(insertionRange.start, position.translate({ characterDelta: context.linkSuffix.length }));
 			yield {
@@ -349,8 +349,8 @@ export class MdVsCodePathCompletionProvider implements vscode.CompletionItemProv
 
 export function registerPathCompletionSupport(
 	selector: vscode.DocumentSelector,
-	engine: MarkdownEngine,
+	parser: IMdParser,
 	linkProvider: MdLinkProvider,
 ): vscode.Disposable {
-	return vscode.languages.registerCompletionItemProvider(selector, new MdVsCodePathCompletionProvider(engine, linkProvider), '.', '/', '#');
+	return vscode.languages.registerCompletionItemProvider(selector, new MdVsCodePathCompletionProvider(parser, linkProvider), '.', '/', '#');
 }

extensions/markdown-language-features/src/languageFeatures/references.ts

@@ -4,7 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 import * as vscode from 'vscode';
 import * as uri from 'vscode-uri';
-import { MarkdownEngine } from '../markdownEngine';
+import { IMdParser } from '../markdownEngine';
 import { MdTableOfContentsProvider, TocEntry } from '../tableOfContents';
 import { noopToken } from '../util/cancellation';
 import { Disposable } from '../util/dispose';
@@ -68,13 +68,13 @@ export class MdReferencesProvider extends Disposable {
 	private readonly _linkComputer: MdLinkComputer;
 
 	public constructor(
-		private readonly engine: MarkdownEngine,
+		private readonly parser: IMdParser,
 		private readonly workspaceContents: MdWorkspaceContents,
 		private readonly tocProvider: MdTableOfContentsProvider,
 	) {
 		super();
 
-		this._linkComputer = new MdLinkComputer(engine);
+		this._linkComputer = new MdLinkComputer(parser);
 		this._linkCache = this._register(new MdWorkspaceInfoCache(workspaceContents, doc => this._linkComputer.getAllLinks(doc, noopToken)));
 	}
 
@@ -114,7 +114,7 @@ export class MdReferencesProvider extends Disposable {
 		for (const link of links) {
 			if (link.href.kind === 'internal'
 				&& this.looksLikeLinkToDoc(link.href, document.uri)
-				&& this.engine.slugifier.fromHeading(link.href.fragment).value === header.slug.value
+				&& this.parser.slugifier.fromHeading(link.href.fragment).value === header.slug.value
 			) {
 				references.push({
 					kind: 'link',
@@ -204,7 +204,7 @@ export class MdReferencesProvider extends Disposable {
 					continue;
 				}
 
-				if (this.engine.slugifier.fromHeading(link.href.fragment).equals(this.engine.slugifier.fromHeading(sourceLink.href.fragment))) {
+				if (this.parser.slugifier.fromHeading(link.href.fragment).equals(this.parser.slugifier.fromHeading(sourceLink.href.fragment))) {
 					const isTriggerLocation = sourceLink.source.resource.fsPath === link.source.resource.fsPath && sourceLink.source.hrefRange.isEqual(link.source.hrefRange);
 					references.push({
 						kind: 'link',

extensions/markdown-language-features/src/languageFeatures/smartSelect.ts

@@ -4,7 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 import Token = require('markdown-it/lib/token');
 import * as vscode from 'vscode';
-import { MarkdownEngine } from '../markdownEngine';
+import { IMdParser } from '../markdownEngine';
 import { MdTableOfContentsProvider, TocEntry } from '../tableOfContents';
 import { SkinnyTextDocument } from '../workspaceContents';
 
@@ -15,7 +15,7 @@ interface MarkdownItTokenWithMap extends Token {
 export class MdSmartSelect implements vscode.SelectionRangeProvider {
 
 	constructor(
-		private readonly engine: MarkdownEngine,
+		private readonly parser: IMdParser,
 		private readonly tocProvider: MdTableOfContentsProvider,
 	) { }
 
@@ -37,9 +37,7 @@ export class MdSmartSelect implements vscode.SelectionRangeProvider {
 	}
 
 	private async getBlockSelectionRange(document: SkinnyTextDocument, position: vscode.Position, headerRange?: vscode.SelectionRange): Promise<vscode.SelectionRange | undefined> {
-
-		const tokens = await this.engine.parse(document);
-
+		const tokens = await this.parser.tokenize(document);
 		const blockTokens = getBlockTokensForPosition(tokens, position, headerRange);
 
 		if (blockTokens.length === 0) {
@@ -253,8 +251,8 @@ function getFirstChildHeader(document: SkinnyTextDocument, header?: TocEntry, to
 
 export function registerSmartSelectSupport(
 	selector: vscode.DocumentSelector,
-	engine: MarkdownEngine,
+	parser: IMdParser,
 	tocProvider: MdTableOfContentsProvider,
 ): vscode.Disposable {
-	return vscode.languages.registerSelectionRangeProvider(selector, new MdSmartSelect(engine, tocProvider));
+	return vscode.languages.registerSelectionRangeProvider(selector, new MdSmartSelect(parser, tocProvider));
 }

extensions/markdown-language-features/src/languageFeatures/workspaceCache.ts

@@ -65,6 +65,15 @@ export class MdDocumentInfoCache<T> extends Disposable {
 		return doc && this.onDidChangeDocument(doc, true)?.value;
 	}
 
+	public async getForDocument(document: SkinnyTextDocument): Promise<T> {
+		const existing = this._cache.get(document.uri);
+		if (existing) {
+			return existing;
+		}
+
+		return this.onDidChangeDocument(document, true)!.value;
+	}
+
 	public async entries(): Promise<Array<[vscode.Uri, T]>> {
 		return this._cache.entries();
 	}