Reduce recomputation of state in markdown extension (#152804)

* Reduce recomputation of state in markdown extension

- Use `getForDocument` more often to avoid refetching documents
- Debounce `MdTableOfContentsWatcher`. We don't want this to trigger on every keystroke :)

* Cache LinkDefinitionSet

* Add test file change

* Fix toc watcher for tests

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

@@ -44,7 +44,7 @@ export function activate(context: vscode.ExtensionContext) {
 	const cspArbiter = new ExtensionContentSecurityPolicyArbiter(context.globalState, context.workspaceState);
 	const commandManager = new CommandManager();
 
-	const engine = new MarkdownItEngine(contributions, githubSlugifier);
+	const engine = new MarkdownItEngine(contributions, githubSlugifier, logger);
 	const workspaceContents = new VsCodeMdWorkspaceContents();
 	const parser = new MdParsingProvider(engine, workspaceContents);
 	const tocProvider = new MdTableOfContentsProvider(parser, workspaceContents, logger);

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

@@ -9,13 +9,13 @@ import * as nls from 'vscode-nls';
 import { CommandManager } from '../commandManager';
 import { ILogger } from '../logging';
 import { MdTableOfContentsProvider } from '../tableOfContents';
-import { MdTableOfContentsWatcher } from '../test/tableOfContentsWatcher';
 import { Delayer } from '../util/async';
 import { noopToken } from '../util/cancellation';
 import { Disposable } from '../util/dispose';
 import { isMarkdownFile, looksLikeMarkdownPath } from '../util/file';
 import { Limiter } from '../util/limiter';
 import { ResourceMap } from '../util/resourceMap';
+import { MdTableOfContentsWatcher } from '../util/tableOfContentsWatcher';
 import { MdWorkspaceContents, SkinnyTextDocument } from '../workspaceContents';
 import { InternalHref, LinkDefinitionSet, MdLink, MdLinkProvider, MdLinkSource } from './documentLinks';
 import { MdReferencesProvider, tryResolveLinkPath } from './references';
@@ -347,7 +347,7 @@ export class DiagnosticManager extends Disposable {
 			}
 		}));
 
-		this.tableOfContentsWatcher = this._register(new MdTableOfContentsWatcher(workspaceContents, tocProvider));
+		this.tableOfContentsWatcher = this._register(new MdTableOfContentsWatcher(workspaceContents, tocProvider, delay));
 		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>();
@@ -491,7 +491,7 @@ export class DiagnosticComputer {
 			return [];
 		}
 
-		const toc = await this.tocProvider.get(doc.uri);
+		const toc = await this.tocProvider.getForDocument(doc);
 		if (token.isCancellationRequested) {
 			return [];
 		}

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

@@ -427,12 +427,17 @@ export class MdLinkComputer {
 	}
 }
 
+interface MdDocumentLinks {
+	readonly links: readonly MdLink[];
+	readonly definitions: LinkDefinitionSet;
+}
+
 /**
  * Stateful object which provides links for markdown files the workspace.
  */
 export class MdLinkProvider extends Disposable {
 
-	private readonly _linkCache: MdDocumentInfoCache<readonly MdLink[]>;
+	private readonly _linkCache: MdDocumentInfoCache<MdDocumentLinks>;
 
 	private readonly linkComputer: MdLinkComputer;
 
@@ -443,21 +448,19 @@ export class MdLinkProvider extends Disposable {
 	) {
 		super();
 		this.linkComputer = new MdLinkComputer(tokenizer);
-		this._linkCache = this._register(new MdDocumentInfoCache(workspaceContents, doc => {
+		this._linkCache = this._register(new MdDocumentInfoCache(workspaceContents, async doc => {
 			logger.verbose('LinkProvider', `compute - ${doc.uri}`);
-			return this.linkComputer.getAllLinks(doc, noopToken);
+
+			const links = await this.linkComputer.getAllLinks(doc, noopToken);
+			return {
+				links,
+				definitions: new LinkDefinitionSet(links),
+			};
 		}));
 	}
 
-	public async getLinks(document: SkinnyTextDocument): Promise<{
+	public async getLinks(document: SkinnyTextDocument): Promise<MdDocumentLinks> {
-		readonly links: readonly MdLink[];
+		return this._linkCache.getForDocument(document);
-		readonly definitions: LinkDefinitionSet;
-	}> {
-		const links = (await this._linkCache.get(document.uri)) ?? [];
-		return {
-			links,
-			definitions: new LinkDefinitionSet(links),
-		};
 	}
 }
 

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

@@ -56,7 +56,7 @@ export class MdFoldingProvider implements vscode.FoldingRangeProvider {
 	}
 
 	private async getHeaderFoldingRanges(document: SkinnyTextDocument): Promise<vscode.FoldingRange[]> {
-		const toc = await this.tocProvide.get(document.uri);
+		const toc = await this.tocProvide.getForDocument(document);
 		return toc.entries.map(entry => {
 			let endLine = entry.sectionLocation.range.end.line;
 			if (document.lineAt(endLine).isEmptyOrWhitespace && endLine >= entry.line + 1) {

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

@@ -83,7 +83,7 @@ export class MdReferencesProvider extends Disposable {
 	public async getReferencesAtPosition(document: SkinnyTextDocument, position: vscode.Position, token: vscode.CancellationToken): Promise<MdReference[]> {
 		this.logger.verbose('ReferencesProvider', `getReferencesAtPosition: ${document.uri}`);
 
-		const toc = await this.tocProvider.get(document.uri);
+		const toc = await this.tocProvider.getForDocument(document);
 		if (token.isCancellationRequested) {
 			return [];
 		}

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

@@ -53,7 +53,7 @@ export class MdSmartSelect implements vscode.SelectionRangeProvider {
 	}
 
 	private async getHeaderSelectionRange(document: SkinnyTextDocument, position: vscode.Position): Promise<vscode.SelectionRange | undefined> {
-		const toc = await this.tocProvider.get(document.uri);
+		const toc = await this.tocProvider.getForDocument(document);
 
 		const headerInfo = getHeadersForPosition(toc.entries, position);
 

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

@@ -6,13 +6,14 @@
 import type MarkdownIt = require('markdown-it');
 import type Token = require('markdown-it/lib/token');
 import * as vscode from 'vscode';
-import { MdDocumentInfoCache } from './util/workspaceCache';
+import { ILogger } from './logging';
 import { MarkdownContributionProvider } from './markdownExtensions';
 import { Slugifier } from './slugify';
 import { Disposable } from './util/dispose';
 import { stringHash } from './util/hash';
 import { WebviewResourceProvider } from './util/resources';
 import { isOfScheme, Schemes } from './util/schemes';
+import { MdDocumentInfoCache } from './util/workspaceCache';
 import { MdWorkspaceContents, SkinnyTextDocument } from './workspaceContents';
 
 const UNICODE_NEWLINE_REGEX = /\u2028|\u2029/g;
@@ -95,6 +96,7 @@ interface RenderEnv {
 
 export interface IMdParser {
 	readonly slugifier: Slugifier;
+
 	tokenize(document: SkinnyTextDocument): Promise<Token[]>;
 }
 
@@ -110,6 +112,7 @@ export class MarkdownItEngine implements IMdParser {
 	public constructor(
 		private readonly contributionProvider: MarkdownContributionProvider,
 		slugifier: Slugifier,
+		private readonly logger: ILogger,
 	) {
 		this.slugifier = slugifier;
 
@@ -180,6 +183,7 @@ export class MarkdownItEngine implements IMdParser {
 			return cached;
 		}
 
+		this.logger.verbose('MarkdownItEngine', `tokenizeDocument - ${document.uri}`);
 		const tokens = this.tokenizeString(document.getText(), engine);
 		this._tokenCache.update(document, config, tokens);
 		return tokens;

extensions/markdown-language-features/src/test/engine.ts

@@ -8,6 +8,7 @@ import { MarkdownItEngine } from '../markdownEngine';
 import { MarkdownContributionProvider, MarkdownContributions } from '../markdownExtensions';
 import { githubSlugifier } from '../slugify';
 import { Disposable } from '../util/dispose';
+import { nulLogger } from './nulLogging';
 
 const emptyContributions = new class extends Disposable implements MarkdownContributionProvider {
 	readonly extensionUri = vscode.Uri.file('/');
@@ -16,5 +17,5 @@ const emptyContributions = new class extends Disposable implements MarkdownContr
 };
 
 export function createNewMarkdownEngine(): MarkdownItEngine {
-	return new MarkdownItEngine(emptyContributions, githubSlugifier);
+	return new MarkdownItEngine(emptyContributions, githubSlugifier, nulLogger);
 }

extensions/markdown-language-features/src/util/tableOfContentsWatcher.ts

@@ -5,10 +5,11 @@
 
 import * as vscode from 'vscode';
 import { MdTableOfContentsProvider, TableOfContents } from '../tableOfContents';
-import { equals } from '../util/arrays';
+import { equals } from './arrays';
-import { Disposable } from '../util/dispose';
+import { Disposable } from './dispose';
-import { ResourceMap } from '../util/resourceMap';
+import { ResourceMap } from './resourceMap';
 import { MdWorkspaceContents, SkinnyTextDocument } from '../workspaceContents';
+import { Delayer } from './async';
 
 /**
  * Check if the items in a table of contents have changed.
@@ -27,15 +28,22 @@ export class MdTableOfContentsWatcher extends Disposable {
 		readonly toc: TableOfContents;
 	}>();
 
+	private readonly _pending = new ResourceMap<void>();
+
 	private readonly _onTocChanged = this._register(new vscode.EventEmitter<{ readonly uri: vscode.Uri }>);
 	public readonly onTocChanged = this._onTocChanged.event;
 
+	private readonly delayer: Delayer<void>;
+
 	public constructor(
 		private readonly workspaceContents: MdWorkspaceContents,
 		private readonly tocProvider: MdTableOfContentsProvider,
+		private readonly delay: number,
 	) {
 		super();
 
+		this.delayer = this._register(new Delayer<void>(delay));
+
 		this._register(this.workspaceContents.onDidChangeMarkdownDocument(this.onDidChangeDocument, this));
 		this._register(this.workspaceContents.onDidCreateMarkdownDocument(this.onDidCreateDocument, this));
 		this._register(this.workspaceContents.onDidDeleteMarkdownDocument(this.onDidDeleteDocument, this));
@@ -47,17 +55,34 @@ export class MdTableOfContentsWatcher extends Disposable {
 	}
 
 	private async onDidChangeDocument(document: SkinnyTextDocument) {
-		const existing = this._files.get(document.uri);
+		if (this.delay > 0) {
-		const newToc = await this.tocProvider.getForDocument(document);
+			this._pending.set(document.uri);
-
+			this.delayer.trigger(() => this.flushPending());
-		if (!existing || hasTableOfContentsChanged(existing.toc, newToc)) {
+		} else {
-			this._onTocChanged.fire({ uri: document.uri });
+			this.updateForResource(document.uri);
 		}
-
-		this._files.set(document.uri, { toc: newToc });
 	}
 
 	private onDidDeleteDocument(resource: vscode.Uri) {
 		this._files.delete(resource);
+		this._pending.delete(resource);
+	}
+
+	private async flushPending() {
+		const pending = [...this._pending.keys()];
+		this._pending.clear();
+
+		return Promise.all(pending.map(resource => this.updateForResource(resource)));
+	}
+
+	private async updateForResource(resource: vscode.Uri) {
+		const existing = this._files.get(resource);
+		const newToc = await this.tocProvider.get(resource);
+
+		if (!existing || hasTableOfContentsChanged(existing.toc, newToc)) {
+			this._onTocChanged.fire({ uri: resource });
+		}
+
+		this._files.set(resource, { toc: newToc });
 	}
 }