Improve markdown references

For #146277

- Find all references on links
- Better support for links without a file extension

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

@@ -22,17 +22,30 @@ export interface InternalLinkTarget {
 	readonly kind: 'internal';
 
 	readonly fromResource: vscode.Uri;
+
 	readonly path: vscode.Uri;
 	readonly fragment: string;
 }
 
-export type LinkTarget = ExternalLinkTarget | InternalLinkTarget;
+export interface ReferenceLinkTarget {
+	readonly kind: 'reference';
+
+	readonly position: vscode.Position;
+}
+
+export interface DefinitionLinkTarget {
+	readonly kind: 'definition';
+
+	readonly target: ExternalLinkTarget | InternalLinkTarget;
+}
+
+export type LinkTarget = ExternalLinkTarget | InternalLinkTarget | ReferenceLinkTarget | DefinitionLinkTarget;
 
 
 function parseLink(
 	document: SkinnyTextDocument,
 	link: string,
-): LinkTarget | undefined {
+): ExternalLinkTarget | InternalLinkTarget | undefined {
 	const cleanLink = stripAngleBrackets(link);
 	const externalSchemeUri = getUriForLinkWithKnownExternalScheme(cleanLink);
 	if (externalSchemeUri) {
@@ -172,15 +185,24 @@ function isLinkInsideCode(code: CodeInDocument, link: LinkData) {
 		code.inline.some(position => position.intersection(link.sourceRange));
 }
 
-function createDocumentLink(sourceRange: vscode.Range, target: LinkTarget) {
-	if (target.kind === 'external') {
-		return new vscode.DocumentLink(sourceRange, target.uri);
-	} else {
-
-		const uri = OpenDocumentLinkCommand.createCommandUri(target.fromResource, target.path, target.fragment);
-		const documentLink = new vscode.DocumentLink(sourceRange, uri);
-		documentLink.tooltip = localize('documentLink.tooltip', 'Follow link');
-		return documentLink;
+function createDocumentLink(sourceRange: vscode.Range, target: LinkTarget): vscode.DocumentLink {
+	switch (target.kind) {
+		case 'external': {
+			return new vscode.DocumentLink(sourceRange, target.uri);
+		}
+		case 'internal': {
+			const uri = OpenDocumentLinkCommand.createCommandUri(target.fromResource, target.path, target.fragment);
+			const documentLink = new vscode.DocumentLink(sourceRange, uri);
+			documentLink.tooltip = localize('documentLink.tooltip', 'Follow link');
+			return documentLink;
+		}
+		case 'reference': {
+			return new vscode.DocumentLink(
+				sourceRange,
+				vscode.Uri.parse(`command:_markdown.moveCursorToPosition?${encodeURIComponent(JSON.stringify([target.position.line, target.position.character]))}`));
+		}
+		case 'definition':
+			return createDocumentLink(sourceRange, target.target);
 	}
 }
 
@@ -193,15 +215,20 @@ export class MdLinkProvider implements vscode.DocumentLinkProvider {
 		document: SkinnyTextDocument,
 		_token: vscode.CancellationToken
 	): Promise<vscode.DocumentLink[]> {
-		const text = document.getText();
-		const inlineLinks = await this.getInlineLinks(text, document);
-		return [
-			...inlineLinks.map(data => createDocumentLink(data.sourceRange, data.target)),
-			...this.getReferenceLinks(text, document)
-		];
+		return (await this.getAllLinks(document)).map(data => createDocumentLink(data.sourceRange, data.target));
 	}
 
-	public async getInlineLinks(text: string, document: SkinnyTextDocument): Promise<LinkData[]> {
+	public async getAllLinks(document: SkinnyTextDocument): Promise<LinkData[]> {
+		return Array.from([
+			...(await this.getInlineLinks(document)),
+			...this.getReferenceLinks(document),
+			...this.getDefinitionLinks(document),
+		]);
+	}
+
+	public async getInlineLinks(document: SkinnyTextDocument): Promise<LinkData[]> {
+		const text = document.getText();
+
 		const results: LinkData[] = [];
 		const codeInDocument = await findCode(document, this.engine);
 		for (const match of text.matchAll(linkPattern)) {
@@ -217,8 +244,9 @@ export class MdLinkProvider implements vscode.DocumentLinkProvider {
 		return results;
 	}
 
-	public *getReferenceLinks(text: string, document: SkinnyTextDocument): Iterable<vscode.DocumentLink> {
-		const definitions = this.getDefinitions(text, document);
+	public *getReferenceLinks(document: SkinnyTextDocument): Iterable<LinkData> {
+		const text = document.getText();
+		const definitions = this.getDefinitions(document);
 		for (const match of text.matchAll(referenceLinkPattern)) {
 			let linkStart: vscode.Position;
 			let linkEnd: vscode.Position;
@@ -237,23 +265,33 @@ export class MdLinkProvider implements vscode.DocumentLinkProvider {
 				continue;
 			}
 
-			try {
-				const link = definitions.get(reference);
-				if (link) {
-					yield new vscode.DocumentLink(
-						new vscode.Range(linkStart, linkEnd),
-						vscode.Uri.parse(`command:_markdown.moveCursorToPosition?${encodeURIComponent(JSON.stringify([link.linkRange.start.line, link.linkRange.start.character]))}`));
-				}
-			} catch (e) {
-				// noop
+			const link = definitions.get(reference);
+			if (link) {
+				yield {
+					sourceRange: new vscode.Range(linkStart, linkEnd),
+					target: {
+						kind: 'reference',
+						position: link.linkRange.start
+					}
+				};
+
 			}
 		}
+	}
 
+	public *getDefinitionLinks(document: SkinnyTextDocument): Iterable<LinkData> {
+		const definitions = this.getDefinitions(document);
 		for (const definition of definitions.values()) {
 			try {
 				const target = parseLink(document, definition.link);
 				if (target) {
-					yield createDocumentLink(definition.linkRange, target);
+					yield {
+						sourceRange: definition.linkRange,
+						target: {
+							kind: 'definition',
+							target
+						}
+					};
 				}
 			} catch (e) {
 				// noop
@@ -261,7 +299,8 @@ export class MdLinkProvider implements vscode.DocumentLinkProvider {
 		}
 	}
 
-	public getDefinitions(text: string, document: SkinnyTextDocument): Map<string, { readonly link: string; readonly linkRange: vscode.Range }> {
+	public getDefinitions(document: SkinnyTextDocument): Map<string, { readonly link: string; readonly linkRange: vscode.Range }> {
+		const text = document.getText();
 		const out = new Map<string, { link: string; linkRange: vscode.Range }>();
 		for (const match of text.matchAll(definitionPattern)) {
 			const pre = match[1];

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

@@ -236,7 +236,7 @@ export class MdPathCompletionProvider implements vscode.CompletionItemProvider {
 		const insertionRange = new vscode.Range(context.linkTextStartPosition, position);
 		const replacementRange = new vscode.Range(insertionRange.start, position.translate({ characterDelta: context.linkSuffix.length }));
 
-		const definitions = this.linkProvider.getDefinitions(document.getText(), document);
+		const definitions = this.linkProvider.getDefinitions(document);
 		for (const def of definitions) {
 			yield {
 				kind: vscode.CompletionItemKind.Reference,

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

@@ -2,28 +2,34 @@
  *  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 * as uri from 'vscode-uri';
 import { MarkdownEngine } from '../markdownEngine';
-import { TableOfContents } from '../tableOfContents';
+import { TableOfContents, TocEntry } from '../tableOfContents';
 import { Disposable } from '../util/dispose';
 import { MdWorkspaceContents, SkinnyTextDocument } from '../workspaceContents';
-import { InternalLinkTarget, LinkData, MdLinkProvider } from './documentLinkProvider';
+import { InternalLinkTarget, LinkData, LinkTarget, MdLinkProvider } from './documentLinkProvider';
 import { MdWorkspaceCache } from './workspaceCache';
 
 
+function isLinkToHeader(target: LinkTarget, header: TocEntry, headerDocument: vscode.Uri): target is InternalLinkTarget {
+	return target.kind === 'internal'
+		&& target.path.fsPath === headerDocument.fsPath
+		&& target.fragment === header.slug.value;
+}
+
 export class MdReferencesProvider extends Disposable implements vscode.ReferenceProvider {
 
 	private readonly _linkCache: MdWorkspaceCache<Promise<LinkData[]>>;
 
 	public constructor(
-		linkProvider: MdLinkProvider,
-		workspaceContents: MdWorkspaceContents,
+		private readonly linkProvider: MdLinkProvider,
+		private readonly workspaceContents: MdWorkspaceContents,
 		private readonly engine: MarkdownEngine,
 	) {
 		super();
 
-		this._linkCache = this._register(new MdWorkspaceCache(workspaceContents, doc => linkProvider.getInlineLinks(doc.getText(), doc)));
+		this._linkCache = this._register(new MdWorkspaceCache(workspaceContents, doc => linkProvider.getAllLinks(doc)));
 	}
 
 	async provideReferences(document: SkinnyTextDocument, position: vscode.Position, context: vscode.ReferenceContext, token: vscode.CancellationToken): Promise<vscode.Location[] | undefined> {
@@ -33,29 +39,80 @@ export class MdReferencesProvider extends Disposable implements vscode.Reference
 		}
 
 		const header = toc.entries.find(entry => entry.line === position.line);
-		if (!header) {
-			return undefined;
+		if (header) {
+			return this.getReferencesToHeader(document, header, context);
+		} else {
+			return this.getReferencesToLink(document, position, context);
 		}
+	}
 
-		const locations: vscode.Location[] = [];
+	private async getReferencesToHeader(document: SkinnyTextDocument, header: TocEntry, context: vscode.ReferenceContext,): Promise<vscode.Location[] | undefined> {
+		const links = (await Promise.all(await this._linkCache.getAll())).flat();
+
+		const references: vscode.Location[] = [];
 
 		if (context.includeDeclaration) {
 			const line = document.lineAt(header.line);
-			locations.push(new vscode.Location(document.uri, new vscode.Range(header.line, 0, header.line, line.text.length)));
+			references.push(new vscode.Location(document.uri, new vscode.Range(header.line, 0, header.line, line.text.length)));
 		}
 
-		(await Promise.all(await this._linkCache.getAll()))
-			.flat()
-			.filter(link => {
-				return link.target.kind === 'internal'
-					&& link.target.path.fsPath === document.uri.fsPath
-					&& link.target.fragment === header.slug.value;
-			})
-			.forEach(link => {
-				const target = link.target as InternalLinkTarget;
-				locations.push(new vscode.Location(target.fromResource, link.sourceRange));
-			});
+		for (const link of links) {
+			if (isLinkToHeader(link.target, header, document.uri)) {
+				references.push(new vscode.Location(link.target.fromResource, link.sourceRange));
+			} else if (link.target.kind === 'definition' && isLinkToHeader(link.target.target, header, document.uri)) {
+				references.push(new vscode.Location(link.target.target.fromResource, link.sourceRange));
+			}
+		}
 
-		return locations;
+		return references;
+	}
+
+	private async getReferencesToLink(document: SkinnyTextDocument, position: vscode.Position, context: vscode.ReferenceContext): Promise<vscode.Location[] | undefined> {
+		const links = (await Promise.all(await this._linkCache.getAll())).flat();
+
+		const docLinks = await this.linkProvider.getInlineLinks(document);
+		const sourceLink = docLinks.find(link => link.sourceRange.contains(position));
+
+		if (sourceLink?.target.kind !== 'internal') {
+			return undefined;
+		}
+
+
+		let targetDoc = await this.workspaceContents.getMarkdownDocument(sourceLink.target.path);
+		if (!targetDoc) {
+			// We don't think the file exists. If it doesn't already have an extension, try tacking on a `.md` and using that instead
+			if (uri.Utils.extname(sourceLink.target.path) === '') {
+				const dotMdResource = sourceLink.target.path.with({ path: sourceLink.target.path.path + '.md' });
+				targetDoc = await this.workspaceContents.getMarkdownDocument(dotMdResource);
+			}
+		}
+
+		if (!targetDoc) {
+			return undefined;
+		}
+
+		const references: vscode.Location[] = [];
+
+		if (context.includeDeclaration) {
+			const toc = await TableOfContents.create(this.engine, targetDoc);
+			const entry = toc.lookup(sourceLink.target.fragment);
+			if (entry) {
+				references.push(entry.location);
+			}
+		}
+
+		for (const link of links) {
+			if (link.target.kind === 'internal'
+				&& link.target.fragment === sourceLink.target.fragment
+				&& (
+					link.target.path.fsPath === targetDoc.uri.fsPath
+					|| uri.Utils.extname(link.target.path) === '' && link.target.path.with({ path: link.target.path.path + '.md' }).fsPath === targetDoc.uri.fsPath
+				)
+			) {
+				references.push(new vscode.Location(link.target.fromResource, link.sourceRange));
+			}
+		}
+
+		return references;
 	}
 }