src/goCodeLens: show codelens for implementations of symbols

docs/settings.md

@@ -211,11 +211,15 @@ Feature level setting to enable/disable code lens for references and run/debug t
 | Properties | Description |
 | --- | --- |
 | `runtest` | If true, enables code lens for running and debugging tests <br/> Default: `true` |
+| `implementation` | If true, enables code lens for showing implementations <br/> Default: `true` |
+| `prefetchImpls` | If true, enables code lens for showing implementations <br/> Default: `false` |
 
 Default:
 ```
 {
 	"runtest" :	true,
+	"implementation" :	true,
+	"prefetchImpls" :	false,
 }
 ```
 ### `go.formatFlags`

extension/package.json

@@ -1551,11 +1551,23 @@
               "type": "boolean",
               "default": true,
               "description": "If true, enables code lens for running and debugging tests"
+            },
+            "implementation": {
+              "type": "boolean",
+              "default": true,
+              "description": "If true, enables code lens for showing implementations"
+            },
+            "prefetchImpls": {
+              "type": "boolean",
+              "default": false,
+              "description": "If true, shows implementation count in the code lens text"
             }
           },
           "additionalProperties": false,
           "default": {
-            "runtest": true
+            "runtest": true,
+            "implementation": true,
+            "prefetchImpls": false
           },
           "description": "Feature level setting to enable/disable code lens for references and run/debug tests",
           "scope": "resource"

extension/src/goCodeLens.ts

@@ -0,0 +1,192 @@
+/* eslint-disable @typescript-eslint/no-unused-vars */
+/* eslint-disable @typescript-eslint/no-explicit-any */
+/*---------------------------------------------------------
+ * Copyright (C) Microsoft Corporation. All rights reserved.
+ * Licensed under the MIT License. See LICENSE in the project root for license information.
+ *--------------------------------------------------------*/
+
+'use strict';
+
+import vscode = require('vscode');
+import { CancellationToken, CodeLens, TextDocument } from 'vscode';
+import { getGoConfig } from './config';
+import { GoBaseCodeLensProvider } from './goBaseCodelens';
+import { GoDocumentSymbolProvider } from './goDocumentSymbols';
+import { GoExtensionContext } from './context';
+import { GO_MODE } from './goMode';
+import { getSymbolImplementations } from './language/goLanguageServer';
+
+export class GoCodeLensProvider extends GoBaseCodeLensProvider {
+	static activate(ctx: vscode.ExtensionContext, goCtx: GoExtensionContext) {
+		const codeLensProvider = new this(goCtx);
+		ctx.subscriptions.push(vscode.languages.registerCodeLensProvider(GO_MODE, codeLensProvider));
+		ctx.subscriptions.push(
+			vscode.workspace.onDidChangeConfiguration(async (e: vscode.ConfigurationChangeEvent) => {
+				if (!e.affectsConfiguration('go')) {
+					return;
+				}
+				const updatedGoConfig = getGoConfig();
+				if (updatedGoConfig['enableCodeLens']) {
+					codeLensProvider.setEnabled(updatedGoConfig['enableCodeLens']['implementation']);
+				}
+			})
+		);
+
+		codeLensProvider.goToImplementations = codeLensProvider.goToImplementations.bind(codeLensProvider);
+
+		vscode.commands.registerCommand('go.codeLens.goToImplementations', codeLensProvider.goToImplementations);
+	}
+
+	constructor(private readonly goCtx: GoExtensionContext) {
+		super();
+	}
+
+	public async provideCodeLenses(document: TextDocument, token: CancellationToken): Promise<CodeLens[]> {
+		if (!this.enabled) {
+			return [];
+		}
+		const config = getGoConfig(document.uri);
+		const codeLensConfig = config.get<{ [key: string]: any }>('enableCodeLens');
+		const codelensEnabled = codeLensConfig ? codeLensConfig['implementation'] : false;
+		if (!codelensEnabled || !document.fileName.endsWith('.go')) {
+			return [];
+		}
+
+		const prefetchImpls = codeLensConfig ? codeLensConfig['prefetchImpls'] : false;
+
+		const abstractCodelenses = this.getCodeLensForAbstractSymbols(document, token, prefetchImpls);
+		const concreteCodelenses = this.getCodeLensForConcreteSymbols(document, token, prefetchImpls);
+
+		const codeLenses = await Promise.all([abstractCodelenses, concreteCodelenses]);
+		return codeLenses.flat();
+	}
+
+	private async getCodeLensForConcreteSymbols(
+		document: TextDocument,
+		token: CancellationToken,
+		prefetchImpls: boolean
+	): Promise<CodeLens[]> {
+		const concreteTypes = await this.getConcreteTypes(document);
+		if (concreteTypes && concreteTypes.length) {
+			const concreteTypesCodeLens = await this.mapSymbolsToCodeLenses(document, concreteTypes, prefetchImpls);
+			return concreteTypesCodeLens;
+		}
+
+		return [];
+	}
+
+	private async getCodeLensForAbstractSymbols(
+		document: TextDocument,
+		token: CancellationToken,
+		prefetchImpls: boolean
+	): Promise<CodeLens[]> {
+		const interfaces = await this.getInterfaces(document);
+		if (interfaces && interfaces.length) {
+			const interfacesCodeLens = this.mapSymbolsToCodeLenses(document, interfaces, prefetchImpls);
+
+			const methodsCodeLens = this.mapSymbolsToCodeLenses(
+				document,
+				interfaces.flatMap((i) => i.children),
+				prefetchImpls
+			);
+
+			const codeLenses = await Promise.all([interfacesCodeLens, methodsCodeLens]);
+
+			return codeLenses.flat();
+		}
+		return [];
+	}
+
+	private async getInterfaces(document: TextDocument): Promise<vscode.DocumentSymbol[]> {
+		const documentSymbolProvider = GoDocumentSymbolProvider(this.goCtx);
+		const symbols = await documentSymbolProvider.provideDocumentSymbols(document);
+		if (!symbols || symbols.length === 0) {
+			return [];
+		}
+		const pkg = symbols[0];
+		if (!pkg) {
+			return [];
+		}
+		const children = pkg.children;
+		const interfaces = children.filter((s) => s.kind === vscode.SymbolKind.Interface);
+		if (!interfaces) {
+			return [];
+		}
+
+		return interfaces;
+	}
+
+	private async getConcreteTypes(document: TextDocument): Promise<vscode.DocumentSymbol[]> {
+		const documentSymbolProvider = GoDocumentSymbolProvider(this.goCtx);
+		const symbols = await documentSymbolProvider.provideDocumentSymbols(document);
+		if (!symbols || symbols.length === 0) {
+			return [];
+		}
+		const pkg = symbols[0];
+		if (!pkg) {
+			return [];
+		}
+		const children = pkg.children;
+		const concreteTypes = children.filter((s) =>
+			[vscode.SymbolKind.Struct, vscode.SymbolKind.Method].includes(s.kind)
+		);
+		if (!concreteTypes) {
+			return [];
+		}
+
+		return concreteTypes;
+	}
+
+	private async mapSymbolsToCodeLenses(
+		document: vscode.TextDocument,
+		symbols: vscode.DocumentSymbol[],
+		prefetchImpls: boolean
+	): Promise<vscode.CodeLens[]> {
+		if (prefetchImpls) {
+			return Promise.all(
+				symbols.map(async (s) => {
+					const implementations = await this.getImplementations(document, s);
+					if (implementations.length) {
+						return new CodeLens(s.range, {
+							title: `${implementations.length} implementation${implementations.length > 1 ? 's' : ''}`,
+							command: 'editor.action.goToLocations',
+							arguments: [document.uri, s.range.start, implementations, 'peek']
+						});
+					}
+
+					return new CodeLens(s.range, {
+						title: 'no implementation found',
+						command: ''
+					});
+				})
+			);
+		}
+
+		return symbols.map((s) => {
+			return new CodeLens(s.range, {
+				title: 'implementations',
+				command: 'go.codeLens.goToImplementations',
+				arguments: [document, s]
+			});
+		});
+	}
+
+	private async goToImplementations(document: vscode.TextDocument, symbol: vscode.DocumentSymbol) {
+		const implementations = await this.getImplementations(document, symbol);
+		await vscode.commands.executeCommand(
+			'editor.action.goToLocations',
+			document.uri,
+			symbol.range.start,
+			implementations,
+			'peek',
+			'No implementation found'
+		);
+	}
+
+	private async getImplementations(
+		document: vscode.TextDocument,
+		symbol: vscode.DocumentSymbol
+	): Promise<vscode.Location[]> {
+		return getSymbolImplementations(this.goCtx, document, symbol);
+	}
+}

extension/src/goMain.ts

@@ -73,6 +73,7 @@ import * as commands from './commands';
 import { toggleVulncheckCommandFactory } from './goVulncheck';
 import { GoTaskProvider } from './goTaskProvider';
 import { setTelemetryEnvVars, telemetryReporter } from './goTelemetry';
+import { GoCodeLensProvider } from './goCodeLens';
 
 const goCtx: GoExtensionContext = {};
 
@@ -144,6 +145,7 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<ExtensionA
 	registerCommand('go.builds.run', commands.runBuilds);
 	registerCommand('go.environment.status', expandGoStatusBar);
 
+	GoCodeLensProvider.activate(ctx, goCtx);
 	GoRunTestCodeLensProvider.activate(ctx, goCtx);
 	GoDebugConfigurationProvider.activate(ctx, goCtx);
 	GoDebugFactory.activate(ctx, goCtx);

extension/src/language/goLanguageServer.ts

@@ -33,7 +33,8 @@ import {
 	ProvideCompletionItemsSignature,
 	ProvideDocumentFormattingEditsSignature,
 	ResponseError,
-	RevealOutputChannelOn
+	RevealOutputChannelOn,
+	Location
 } from 'vscode-languageclient';
 import { LanguageClient, ServerOptions } from 'vscode-languageclient/node';
 import { getGoConfig, getGoplsConfig, extensionInfo } from '../config';
@@ -1665,3 +1666,39 @@ async function getGoplsStats(binpath?: string) {
 		return `gopls stats -anon failed after ${duration} ms. Please check if gopls is killed by OS.`;
 	}
 }
+
+export async function getSymbolImplementations(
+	goCtx: GoExtensionContext,
+	document: vscode.TextDocument,
+	symbol: vscode.DocumentSymbol
+): Promise<vscode.Location[]> {
+	const languageClient = goCtx.languageClient;
+	if (languageClient) {
+		const params = {
+			textDocument: { uri: document.uri.toString() },
+			position: {
+				line: symbol.selectionRange.start.line,
+				character: symbol.selectionRange.start.character
+			}
+		};
+		try {
+			const implementations = await languageClient.sendRequest<Location[]>('textDocument/implementation', params);
+			return implementations
+				? implementations.map(
+						(i) =>
+							new vscode.Location(
+								vscode.Uri.parse(i.uri),
+								new vscode.Range(
+									new vscode.Position(i.range.start.line, i.range.start.character),
+									new vscode.Position(i.range.end.line, i.range.end.character)
+								)
+							)
+				  )
+				: [];
+		} catch (error) {
+			console.error(`unable to get implementations for ${symbol.name}:`, error);
+			return [];
+		}
+	}
+	return [];
+}