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;
 	}
 }

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

@@ -13,14 +13,18 @@ export class InMemoryWorkspaceMarkdownDocuments implements MdWorkspaceContents {
 
 	constructor(documents: SkinnyTextDocument[]) {
 		for (const doc of documents) {
-			this._documents.set(doc.uri.toString(), doc);
+			this._documents.set(this.getKey(doc.uri), doc);
 		}
 	}
 
-	async getAllMarkdownDocuments() {
+	public async getAllMarkdownDocuments() {
 		return Array.from(this._documents.values());
 	}
 
+	public async getMarkdownDocument(resource: vscode.Uri): Promise<SkinnyTextDocument | undefined> {
+		return this._documents.get(this.getKey(resource));
+	}
+
 	private readonly _onDidChangeMarkdownDocumentEmitter = new vscode.EventEmitter<SkinnyTextDocument>();
 	public onDidChangeMarkdownDocument = this._onDidChangeMarkdownDocumentEmitter.event;
 
@@ -31,19 +35,23 @@ export class InMemoryWorkspaceMarkdownDocuments implements MdWorkspaceContents {
 	public onDidDeleteMarkdownDocument = this._onDidDeleteMarkdownDocumentEmitter.event;
 
 	public updateDocument(document: SkinnyTextDocument) {
-		this._documents.set(document.uri.toString(), document);
+		this._documents.set(this.getKey(document.uri), document);
 		this._onDidChangeMarkdownDocumentEmitter.fire(document);
 	}
 
 	public createDocument(document: SkinnyTextDocument) {
-		assert.ok(!this._documents.has(document.uri.toString()));
+		assert.ok(!this._documents.has(this.getKey(document.uri)));
 
-		this._documents.set(document.uri.toString(), document);
+		this._documents.set(this.getKey(document.uri), document);
 		this._onDidCreateMarkdownDocumentEmitter.fire(document);
 	}
 
 	public deleteDocument(resource: vscode.Uri) {
-		this._documents.delete(resource.toString());
+		this._documents.delete(this.getKey(resource));
 		this._onDidDeleteMarkdownDocumentEmitter.fire(resource);
 	}
+
+	private getKey(resource: vscode.Uri): string {
+		return resource.fsPath;
+	}
 }

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

@@ -22,7 +22,7 @@ function getReferences(doc: InMemoryDocument, pos: vscode.Position, workspaceCon
 	return provider.provideReferences(doc, pos, { includeDeclaration: true }, noopToken);
 }
 
-suite('markdown header references', () => {
+suite('markdown references', () => {
 	test('Should not return references when not on header', async () => {
 		const doc = new InMemoryDocument(workspaceFile('doc.md'), joinLines(
 			`# abc`,
@@ -41,7 +41,7 @@ suite('markdown header references', () => {
 		}
 	});
 
-	test('Should find simple references within same file', async () => {
+	test('Should find references from header within same file', async () => {
 		const doc = new InMemoryDocument(workspaceFile('doc.md'), joinLines(
 			`# abc`,
 			``,
@@ -54,7 +54,7 @@ suite('markdown header references', () => {
 		assert.deepStrictEqual(refs!.length, 3);
 
 		{
-			const ref = refs![0]; // Header own ref
+			const ref = refs![0]; // Header definition
 			assert.deepStrictEqual(ref.range.start.line, 0);
 		}
 		{
@@ -67,7 +67,7 @@ suite('markdown header references', () => {
 		}
 	});
 
-	test('Should find simple references across files', async () => {
+	test('Should find references from header across files', async () => {
 		const docUri = workspaceFile('doc.md');
 		const other1Uri = workspaceFile('sub', 'other.md');
 		const other2Uri = workspaceFile('other2.md');
@@ -94,7 +94,7 @@ suite('markdown header references', () => {
 		assert.deepStrictEqual(refs!.length, 4);
 
 		{
-			const ref = refs![0]; // Header own ref
+			const ref = refs![0]; // Header definition
 			assert.deepStrictEqual(ref.uri.toString(), docUri.toString());
 			assert.deepStrictEqual(ref.range.start.line, 0);
 		}
@@ -114,4 +114,140 @@ suite('markdown header references', () => {
 			assert.deepStrictEqual(ref.range.start.line, 2);
 		}
 	});
+
+	test('Should find references from header to link definitions ', async () => {
+		const doc = new InMemoryDocument(workspaceFile('doc.md'), joinLines(
+			`# abc`,
+			``,
+			`[bla]: #abc`
+		));
+		const refs = await getReferences(doc, new vscode.Position(0, 3), new InMemoryWorkspaceMarkdownDocuments([doc]));
+
+		assert.deepStrictEqual(refs!.length, 2);
+
+		{
+			const ref = refs![0]; // Header definition
+			assert.deepStrictEqual(ref.range.start.line, 0);
+		}
+		{
+			const ref = refs![1];
+			assert.deepStrictEqual(ref.range.start.line, 2);
+		}
+	});
+
+	test('Should find references from link within same file', async () => {
+		const doc = new InMemoryDocument(workspaceFile('doc.md'), joinLines(
+			`# abc`,
+			``,
+			`[link 1](#abc)`,
+			`[not link](#noabc)`,
+			`[link 2](#abc)`,
+		));
+		const refs = await getReferences(doc, new vscode.Position(2, 10), new InMemoryWorkspaceMarkdownDocuments([doc]));
+
+		assert.deepStrictEqual(refs!.length, 3);
+
+		{
+			const ref = refs![0]; // Header definition
+			assert.deepStrictEqual(ref.range.start.line, 0);
+		}
+		{
+			const ref = refs![1];
+			assert.deepStrictEqual(ref.range.start.line, 2);
+		}
+		{
+			const ref = refs![2];
+			assert.deepStrictEqual(ref.range.start.line, 4);
+		}
+	});
+
+	test('Should find references from link across files', async () => {
+		const docUri = workspaceFile('doc.md');
+		const other1Uri = workspaceFile('sub', 'other.md');
+		const other2Uri = workspaceFile('other2.md');
+
+		const doc = new InMemoryDocument(docUri, joinLines(
+			`# abc`,
+			``,
+			`[link 1](#abc)`,
+		));
+		const refs = await getReferences(doc, new vscode.Position(2, 10), new InMemoryWorkspaceMarkdownDocuments([
+			doc,
+			new InMemoryDocument(other1Uri, joinLines(
+				`[not link](#abc)`,
+				`[not link](/doc.md#abz)`,
+				`[with ext](/doc.md#abc)`,
+				`[without ext](/doc#abc)`,
+			)),
+			new InMemoryDocument(other2Uri, joinLines(
+				`[not link](#abc)`,
+				`[not link](./doc.md#abz)`,
+				`[link](./doc.md#abc)`,
+			))
+		]));
+
+		assert.deepStrictEqual(refs!.length, 5);
+
+		{
+			const ref = refs![0]; // Header definition
+			assert.deepStrictEqual(ref.uri.toString(), docUri.toString());
+			assert.deepStrictEqual(ref.range.start.line, 0);
+		}
+		{
+			const ref = refs![1]; // Within file
+			assert.deepStrictEqual(ref.uri.toString(), docUri.toString());
+			assert.deepStrictEqual(ref.range.start.line, 2);
+		}
+		{
+			const ref = refs![2]; // Other with ext
+			assert.deepStrictEqual(ref.uri.toString(), other1Uri.toString());
+			assert.deepStrictEqual(ref.range.start.line, 2);
+		}
+		{
+			const ref = refs![3]; // Other without ext
+			assert.deepStrictEqual(ref.uri.toString(), other1Uri.toString());
+			assert.deepStrictEqual(ref.range.start.line, 3);
+		}
+		{
+			const ref = refs![4]; // Other2
+			assert.deepStrictEqual(ref.uri.toString(), other2Uri.toString());
+			assert.deepStrictEqual(ref.range.start.line, 2);
+		}
+	});
+
+	test('Should find references from link across files when triggered on link without file extension ', async () => {
+		const docUri = workspaceFile('doc.md');
+		const other1Uri = workspaceFile('sub', 'other.md');
+
+		const doc = new InMemoryDocument(docUri, joinLines(
+			`[with ext](./sub/other#header)`,
+			`[without ext](./sub/other.md#header)`,
+		));
+		const refs = await getReferences(doc, new vscode.Position(0, 15), new InMemoryWorkspaceMarkdownDocuments([
+			doc,
+			new InMemoryDocument(other1Uri, joinLines(
+				`pre`,
+				`# header`,
+				`post`,
+			)),
+		]));
+
+		assert.deepStrictEqual(refs!.length, 3);
+
+		{
+			const ref = refs![0]; // Header definition
+			assert.deepStrictEqual(ref.uri.toString(), other1Uri.toString());
+			assert.deepStrictEqual(ref.range.start.line, 1);
+		}
+		{
+			const ref = refs![1];
+			assert.deepStrictEqual(ref.uri.toString(), docUri.toString());
+			assert.deepStrictEqual(ref.range.start.line, 0);
+		}
+		{
+			const ref = refs![2];
+			assert.deepStrictEqual(ref.uri.toString(), docUri.toString());
+			assert.deepStrictEqual(ref.range.start.line, 1);
+		}
+	});
 });

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

@@ -38,6 +38,8 @@ export interface MdWorkspaceContents {
 	 */
 	getAllMarkdownDocuments(): Promise<Iterable<SkinnyTextDocument>>;
 
+	getMarkdownDocument(resource: vscode.Uri): Promise<SkinnyTextDocument | undefined>;
+
 	readonly onDidChangeMarkdownDocument: vscode.Event<SkinnyTextDocument>;
 	readonly onDidCreateMarkdownDocument: vscode.Event<SkinnyTextDocument>;
 	readonly onDidDeleteMarkdownDocument: vscode.Event<vscode.Uri>;
@@ -124,7 +126,7 @@ export class VsCodeMdWorkspaceContents extends Disposable implements MdWorkspace
 		}));
 	}
 
-	private async getMarkdownDocument(resource: vscode.Uri): Promise<SkinnyTextDocument | undefined> {
+	public async getMarkdownDocument(resource: vscode.Uri): Promise<SkinnyTextDocument | undefined> {
 		const matchingDocument = vscode.workspace.textDocuments.find((doc) => doc.uri.toString() === resource.toString());
 		if (matchingDocument) {
 			return matchingDocument;