Revalidate linked files on header change (#152366)

Fixes #150945

With this change, when the headers in a file change, we should also try to revalidate all files that link to it

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

@@ -78,7 +78,7 @@ function registerMarkdownLanguageFeatures(
 		vscode.languages.registerDefinitionProvider(selector, new MdDefinitionProvider(referencesProvider)),
 		MdPathCompletionProvider.register(selector, engine, linkComputer),
 		registerDocumentLinkProvider(selector, linkComputer),
-		registerDiagnostics(selector, engine, workspaceContents, linkComputer, commandManager),
+		registerDiagnostics(selector, engine, workspaceContents, linkComputer, commandManager, referencesProvider),
 		registerDropIntoEditor(selector),
 		registerPasteProvider(selector),
 		registerFindFileReferences(commandManager, referencesProvider),

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

@@ -10,13 +10,15 @@ import { CommandManager } from '../commandManager';
 import { MarkdownEngine } from '../markdownEngine';
 import { TableOfContents } from '../tableOfContents';
 import { Delayer } from '../util/async';
+import { noopToken } from '../util/cancellation';
 import { Disposable } from '../util/dispose';
 import { isMarkdownFile } from '../util/file';
 import { Limiter } from '../util/limiter';
 import { ResourceMap } from '../util/resourceMap';
+import { MdTableOfContentsWatcher } from '../test/tableOfContentsWatcher';
 import { MdWorkspaceContents, SkinnyTextDocument } from '../workspaceContents';
 import { InternalHref, LinkDefinitionSet, MdLink, MdLinkComputer, MdLinkSource } from './documentLinkProvider';
-import { tryFindMdDocumentForLink } from './references';
+import { MdReferencesProvider, tryFindMdDocumentForLink } from './references';
 
 const localize = nls.loadMessageBundle();
 
@@ -93,19 +95,21 @@ class InflightDiagnosticRequests {
 
 	private readonly inFlightRequests = new ResourceMap<{ readonly cts: vscode.CancellationTokenSource }>();
 
-	public trigger(resource: vscode.Uri, compute: (token: vscode.CancellationToken) => Promise<void>) {
+	public async trigger(resource: vscode.Uri, compute: (token: vscode.CancellationToken) => Promise<void>): Promise<void> {
 		this.cancel(resource);
 
 		const cts = new vscode.CancellationTokenSource();
 		const entry = { cts };
 		this.inFlightRequests.set(resource, entry);
 
-		compute(cts.token).finally(() => {
+		try {
+			return await compute(cts.token);
+		} finally {
 			if (this.inFlightRequests.get(resource) === entry) {
 				this.inFlightRequests.delete(resource);
 			}
 			cts.dispose();
-		});
+		}
 	}
 
 	public cancel(resource: vscode.Uri) {
@@ -226,55 +230,136 @@ class LinkDoesNotExistDiagnostic extends vscode.Diagnostic {
 	}
 }
 
-export class DiagnosticManager extends Disposable {
+export abstract class DiagnosticReporter extends Disposable {
+	private readonly pending = new ResourceMap<Promise<any>>();
+
+	public clear(): void {
+		this.pending.clear();
+	}
+
+	public abstract set(uri: vscode.Uri, diagnostics: readonly vscode.Diagnostic[]): void;
+
+	public delete(uri: vscode.Uri): void {
+		this.pending.delete(uri);
+	}
+
+	public signalTriggered(uri: vscode.Uri, recompute: Promise<any>): void {
+		this.pending.set(uri, recompute);
+		recompute.finally(() => {
+			if (this.pending.get(uri) === recompute) {
+				this.pending.delete(uri);
+			}
+		});
+	}
+
+	public async waitAllPending(): Promise<void> {
+		await Promise.all([...this.pending.values()]);
+	}
+}
+
+export class DiagnosticCollectionReporter extends DiagnosticReporter {
 
 	private readonly collection: vscode.DiagnosticCollection;
 
+	constructor() {
+		super();
+		this.collection = this._register(vscode.languages.createDiagnosticCollection('markdown'));
+	}
+
+	public override clear(): void {
+		super.clear();
+		this.collection.clear();
+	}
+
+	public set(uri: vscode.Uri, diagnostics: readonly vscode.Diagnostic[]): void {
+		const tabs = this.getAllTabResources();
+		this.collection.set(uri, tabs.has(uri) ? diagnostics : []);
+	}
+
+	public override delete(uri: vscode.Uri): void {
+		super.delete(uri);
+		this.collection.delete(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);
+				}
+			}
+		}
+		return openedTabDocs;
+	}
+}
+
+export class DiagnosticManager extends Disposable {
+
 	private readonly diagnosticDelayer: Delayer<void>;
 	private readonly pendingDiagnostics = new Set<vscode.Uri>();
 	private readonly inFlightDiagnostics = this._register(new InflightDiagnosticRequests());
 
 	private readonly linkWatcher = this._register(new LinkWatcher());
+	private readonly tableOfContentsWatcher: MdTableOfContentsWatcher;
+
+	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,
+		delay = 300,
 	) {
 		super();
 
-		this.diagnosticDelayer = this._register(new Delayer(300));
-
-		this.collection = this._register(vscode.languages.createDiagnosticCollection('markdown'));
+		this.diagnosticDelayer = this._register(new Delayer(delay));
 
 		this._register(this.configuration.onDidChange(() => {
 			this.rebuild();
 		}));
 
-		this._register(vscode.workspace.onDidOpenTextDocument(doc => {
-			this.triggerDiagnostics(doc);
+		this._register(workspaceContents.onDidCreateMarkdownDocument(doc => {
+			this.triggerDiagnostics(doc.uri);
 		}));
 
-		this._register(vscode.workspace.onDidChangeTextDocument(e => {
-			this.triggerDiagnostics(e.document);
+		this._register(workspaceContents.onDidChangeMarkdownDocument(doc => {
+			this.triggerDiagnostics(doc.uri);
 		}));
 
 		this._register(vscode.workspace.onDidCloseTextDocument(({ uri }) => {
 			this.pendingDiagnostics.delete(uri);
 			this.inFlightDiagnostics.cancel(uri);
 			this.linkWatcher.deleteDocument(uri);
-			this.collection.delete(uri);
+			this.reporter.delete(uri);
 		}));
 
 		this._register(this.linkWatcher.onDidChangeLinkedToFile(changedDocuments => {
 			for (const resource of changedDocuments) {
 				const doc = vscode.workspace.textDocuments.find(doc => doc.uri.toString() === resource.toString());
-				if (doc) {
-					this.triggerDiagnostics(doc);
+				if (doc && isMarkdownFile(doc)) {
+					this.triggerDiagnostics(doc.uri);
 				}
 			}
 		}));
 
-		this.rebuild();
+		this.tableOfContentsWatcher = this._register(new MdTableOfContentsWatcher(engine, workspaceContents));
+		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>();
+			for (const ref of await this.referencesProvider.getAllReferencesToFile(e.uri, noopToken)) {
+				const file = ref.location.uri;
+				if (!triggered.has(file)) {
+					this.triggerDiagnostics(file);
+					triggered.set(file);
+				}
+			}
+		}));
+
+		this.ready = this.rebuild();
 	}
 
 	public override dispose() {
@@ -294,49 +379,33 @@ export class DiagnosticManager extends Disposable {
 		const pending = [...this.pendingDiagnostics];
 		this.pendingDiagnostics.clear();
 
-		for (const resource of pending) {
-			const doc = vscode.workspace.textDocuments.find(doc => doc.uri.fsPath === resource.fsPath);
+		await Promise.all(pending.map(async resource => {
+			const doc = await this.workspaceContents.getMarkdownDocument(resource);
 			if (doc) {
-				this.inFlightDiagnostics.trigger(doc.uri, async (token) => {
+				await this.inFlightDiagnostics.trigger(doc.uri, async (token) => {
 					const state = await this.recomputeDiagnosticState(doc, token);
 					this.linkWatcher.updateLinksForDocument(doc.uri, state.config.enabled && state.config.validateFileLinks ? state.links : []);
-					this.collection.set(doc.uri, state.diagnostics);
+					this.reporter.set(doc.uri, state.diagnostics);
 				});
 			}
-		}
+		}));
 	}
 
 	private async rebuild() {
-		this.collection.clear();
+		this.reporter.clear();
 		this.pendingDiagnostics.clear();
 		this.inFlightDiagnostics.clear();
 
-		const allOpenedTabResources = this.getAllTabResources();
-		await Promise.all(
-			vscode.workspace.textDocuments
-				.filter(doc => allOpenedTabResources.has(doc.uri) && isMarkdownFile(doc))
-				.map(doc => this.triggerDiagnostics(doc)));
+		for (const doc of await this.workspaceContents.getAllMarkdownDocuments()) {
+			this.triggerDiagnostics(doc.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);
-				}
-			}
-		}
-		return openedTabDocs;
-	}
+	private triggerDiagnostics(uri: vscode.Uri) {
+		this.inFlightDiagnostics.cancel(uri);
 
-	private triggerDiagnostics(doc: vscode.TextDocument) {
-		this.inFlightDiagnostics.cancel(doc.uri);
-
-		if (isMarkdownFile(doc)) {
-			this.pendingDiagnostics.add(doc.uri);
-			this.diagnosticDelayer.trigger(() => this.recomputePendingDiagnostics());
-		}
+		this.pendingDiagnostics.add(uri);
+		this.reporter.signalTriggered(uri, this.diagnosticDelayer.trigger(() => this.recomputePendingDiagnostics()));
 	}
 }
 
@@ -558,9 +627,16 @@ export function register(
 	workspaceContents: MdWorkspaceContents,
 	linkComputer: MdLinkComputer,
 	commandManager: CommandManager,
+	referenceProvider: MdReferencesProvider,
 ): vscode.Disposable {
 	const configuration = new VSCodeDiagnosticConfiguration();
-	const manager = new DiagnosticManager(new DiagnosticComputer(engine, workspaceContents, linkComputer), configuration);
+	const manager = new DiagnosticManager(
+		engine,
+		workspaceContents,
+		new DiagnosticComputer(engine, workspaceContents, linkComputer),
+		configuration,
+		new DiagnosticCollectionReporter(),
+		referenceProvider);
 	return vscode.Disposable.from(
 		configuration,
 		manager,

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

@@ -6,10 +6,13 @@
 import * as assert from 'assert';
 import 'mocha';
 import * as vscode from 'vscode';
-import { DiagnosticComputer, DiagnosticConfiguration, DiagnosticLevel, DiagnosticManager, DiagnosticOptions } from '../languageFeatures/diagnostics';
+import { DiagnosticCollectionReporter, DiagnosticComputer, DiagnosticConfiguration, DiagnosticLevel, DiagnosticManager, DiagnosticOptions, DiagnosticReporter } from '../languageFeatures/diagnostics';
 import { MdLinkComputer } from '../languageFeatures/documentLinkProvider';
+import { MdReferencesProvider } from '../languageFeatures/references';
+import { githubSlugifier } from '../slugify';
 import { noopToken } from '../util/cancellation';
 import { InMemoryDocument } from '../util/inMemoryDocument';
+import { ResourceMap } from '../util/resourceMap';
 import { MdWorkspaceContents } from '../workspaceContents';
 import { createNewMarkdownEngine } from './engine';
 import { InMemoryWorkspaceMarkdownDocuments } from './inMemoryWorkspace';
@@ -32,10 +35,22 @@ async function getComputedDiagnostics(doc: InMemoryDocument, workspaceContents:
 	).diagnostics;
 }
 
-function createDiagnosticsManager(workspaceContents: MdWorkspaceContents, configuration = new MemoryDiagnosticConfiguration({})) {
+function createDiagnosticsManager(
+	workspaceContents: MdWorkspaceContents,
+	configuration = new MemoryDiagnosticConfiguration({}),
+	reporter: DiagnosticReporter = new DiagnosticCollectionReporter(),
+) {
 	const engine = createNewMarkdownEngine();
 	const linkComputer = new MdLinkComputer(engine);
-	return new DiagnosticManager(new DiagnosticComputer(engine, workspaceContents, linkComputer), configuration);
+	const referencesProvider = new MdReferencesProvider(linkComputer, workspaceContents, engine, githubSlugifier);
+	return new DiagnosticManager(
+		engine,
+		workspaceContents,
+		new DiagnosticComputer(engine, workspaceContents, linkComputer),
+		configuration,
+		reporter,
+		referencesProvider,
+		0);
 }
 
 function assertDiagnosticsEqual(actual: readonly vscode.Diagnostic[], expectedRanges: readonly vscode.Range[]) {
@@ -72,6 +87,29 @@ class MemoryDiagnosticConfiguration implements DiagnosticConfiguration {
 	}
 }
 
+class MemoryDiagnosticReporter extends DiagnosticReporter {
+	public readonly diagnostics = new ResourceMap<readonly vscode.Diagnostic[]>();
+
+	override dispose(): void {
+		super.clear();
+		this.clear();
+	}
+
+	override clear(): void {
+		super.clear();
+		this.diagnostics.clear();
+	}
+
+	set(uri: vscode.Uri, diagnostics: readonly vscode.Diagnostic[]): void {
+		this.diagnostics.set(uri, diagnostics);
+	}
+
+	override delete(uri: vscode.Uri): void {
+		super.delete(uri);
+		this.diagnostics.delete(uri);
+	}
+}
+
 
 suite('markdown: Diagnostics', () => {
 	test('Should not return any diagnostics for empty document', async () => {
@@ -387,6 +425,64 @@ suite('markdown: Diagnostics', () => {
 			new vscode.Range(3, 14, 3, 22),
 		]);
 	});
+
+	test('Should revalidate linked files when header changes', async () => {
+		const doc1Uri = workspacePath('doc1.md');
+		const doc1 = new InMemoryDocument(doc1Uri, joinLines(
+			`[text](#no-such)`,
+			`[text](/doc2.md#header)`,
+		));
+		const doc2Uri = workspacePath('doc2.md');
+		const doc2 = new InMemoryDocument(doc2Uri, joinLines(
+			`# Header`,
+			`[text](#header)`,
+			`[text](#no-such-2)`,
+		));
+
+		const contents = new InMemoryWorkspaceMarkdownDocuments([doc1, doc2]);
+		const reporter = new MemoryDiagnosticReporter();
+
+		const manager = createDiagnosticsManager(contents, new MemoryDiagnosticConfiguration({}), reporter);
+		await manager.ready;
+
+		// Check initial state
+		await reporter.waitAllPending();
+		assertDiagnosticsEqual(reporter.diagnostics.get(doc1Uri)!, [
+			new vscode.Range(0, 7, 0, 15),
+		]);
+		assertDiagnosticsEqual(reporter.diagnostics.get(doc2Uri)!, [
+			new vscode.Range(2, 7, 2, 17),
+		]);
+
+		// Edit header
+		contents.updateDocument(new InMemoryDocument(doc2Uri, joinLines(
+			`# new header`,
+			`[text](#new-header)`,
+			`[text](#no-such-2)`,
+		)));
+		await reporter.waitAllPending();
+		assertDiagnosticsEqual(orderDiagnosticsByRange(reporter.diagnostics.get(doc1Uri)!), [
+			new vscode.Range(0, 7, 0, 15),
+			new vscode.Range(1, 15, 1, 22),
+		]);
+		assertDiagnosticsEqual(reporter.diagnostics.get(doc2Uri)!, [
+			new vscode.Range(2, 7, 2, 17),
+		]);
+
+		// Revert to original file
+		contents.updateDocument(new InMemoryDocument(doc2Uri, joinLines(
+			`# header`,
+			`[text](#header)`,
+			`[text](#no-such-2)`,
+		)));
+		await reporter.waitAllPending();
+		assertDiagnosticsEqual(orderDiagnosticsByRange(reporter.diagnostics.get(doc1Uri)!), [
+			new vscode.Range(0, 7, 0, 15)
+		]);
+		assertDiagnosticsEqual(reporter.diagnostics.get(doc2Uri)!, [
+			new vscode.Range(2, 7, 2, 17),
+		]);
+	});
 });
 
 function orderDiagnosticsByRange(diagnostics: Iterable<vscode.Diagnostic>): readonly vscode.Diagnostic[] {

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

@@ -0,0 +1,64 @@
+/*---------------------------------------------------------------------------------------------
+ *  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';
+import { MarkdownEngine } from '../markdownEngine';
+import { TableOfContents } from '../tableOfContents';
+import { MdWorkspaceContents, SkinnyTextDocument } from '../workspaceContents';
+import { equals } from '../util/arrays';
+import { Disposable } from '../util/dispose';
+import { ResourceMap } from '../util/resourceMap';
+
+/**
+ * Check if the items in a table of contents have changed.
+ *
+ * This only checks for changes in the entries themselves, not for any changes in their locations.
+ */
+function hasTableOfContentsChanged(a: TableOfContents, b: TableOfContents): boolean {
+	const aSlugs = a.entries.map(entry => entry.slug.value).sort();
+	const bSlugs = b.entries.map(entry => entry.slug.value).sort();
+	return !equals(aSlugs, bSlugs);
+}
+
+export class MdTableOfContentsWatcher extends Disposable {
+
+	private readonly _files = new ResourceMap<{
+		readonly toc: TableOfContents;
+	}>();
+
+	private readonly _onTocChanged = this._register(new vscode.EventEmitter<{ readonly uri: vscode.Uri }>);
+	public readonly onTocChanged = this._onTocChanged.event;
+
+	public constructor(
+		private readonly engine: MarkdownEngine,
+		private readonly workspaceContents: MdWorkspaceContents,
+	) {
+		super();
+
+		this._register(this.workspaceContents.onDidChangeMarkdownDocument(this.onDidChangeDocument, this));
+		this._register(this.workspaceContents.onDidCreateMarkdownDocument(this.onDidCreateDocument, this));
+		this._register(this.workspaceContents.onDidDeleteMarkdownDocument(this.onDidDeleteDocument, this));
+	}
+
+	private async onDidCreateDocument(document: SkinnyTextDocument) {
+		const toc = await TableOfContents.create(this.engine, document);
+		this._files.set(document.uri, { toc });
+	}
+
+	private async onDidChangeDocument(document: SkinnyTextDocument) {
+		const existing = this._files.get(document.uri);
+		const newToc = await TableOfContents.create(this.engine, document);
+
+		if (!existing || hasTableOfContentsChanged(existing.toc, newToc)) {
+			this._onTocChanged.fire({ uri: document.uri });
+		}
+
+		this._files.set(document.uri, { toc: newToc });
+	}
+
+	private onDidDeleteDocument(resource: vscode.Uri) {
+		this._files.delete(resource);
+	}
+}