Add references codelens support (#933)

Author: theSoenke

README.md

@@ -15,6 +15,7 @@ This extension adds rich language support for the Go language to VS Code, includ
 - Quick Info (using `gogetdoc` or `godef`+`godoc`)
 - Goto Definition (using `gogetdoc` or `godef`+`godoc`)
 - Find References (using `guru`)
+- References CodeLens
 - File outline (using `go-outline`)
 - Workspace symbol search (using `go-symbols`)
 - Rename (using `gorename`. Note: For Undo after rename to work in Windows you need to have `diff` tool in your path)
@@ -49,7 +50,7 @@ The Go extension is ready to use on the get go. If you want to customize the fea
 
 
 ### Go Language Server (Experimental)
-Set `go.useLanguageServer` to `true` to use the Go language server from [Sourcegraph](https://github.com/sourcegraph/go-langserver) for features like Hover, Definition, Find All References, Signature Help, Go to Symbol in File and Workspace. 
+Set `go.useLanguageServer` to `true` to use the Go language server from [Sourcegraph](https://github.com/sourcegraph/go-langserver) for features like Hover, Definition, Find All References, Signature Help, Go to Symbol in File and Workspace.
 * This is an experimental feature and is not available in Windows yet.
 * If set to true, you will be prompted to install the Go language server. Once installed, you will have to reload VS Code window. The language server will then be run by the Go extension in the background to provide services needed for the above mentioned features.
 * Everytime you change the value of the setting `go.useLanguageServer`, you need to reload the VS Code window for it to take effect.
@@ -112,7 +113,7 @@ For more read [Debugging Go Code Using VS Code](https://github.com/Microsoft/vsc
 
 #### Remote Debugging
 
-To remote debug using VS Code, read [Remote Debugging](https://github.com/Microsoft/vscode-go/wiki/Debugging-Go-code-using-VS-Code#remote-debugging) 
+To remote debug using VS Code, read [Remote Debugging](https://github.com/Microsoft/vscode-go/wiki/Debugging-Go-code-using-VS-Code#remote-debugging)
 
 ## Building and Debugging the Extension
 

package.json

@@ -510,16 +510,21 @@
           "default": false,
           "description": "If false, the import statements will be excluded while using the Go to Symbol in File feature"
         },
+        "go.referencesCodeLens.enabled": {
+          "type": "boolean",
+          "default": false,
+          "description": "Enable/disable references CodeLens"
+        },
         "go.addTags": {
           "type": "object",
           "properties": {
             "promptForTags": {
-              "type":"boolean",
+              "type": "boolean",
               "default": false,
               "description": "If true, Go: Add Tags command will prompt the user to provide tags and options instead of using the configured values"
             },
             "tags": {
-              "type":"string",
+              "type": "string",
               "default": "json",
               "description": "Comma separated tags to be used by Go: Add Tags command"
             },
@@ -559,7 +564,7 @@
               "default": 500,
               "description": "The number of milliseconds to delay before execution. Resets with each keystroke."
             }
-          },          
+          },
           "default": {
             "enabled": false,
             "delay": 500
@@ -570,12 +575,12 @@
           "type": "object",
           "properties": {
             "promptForTags": {
-              "type":"boolean",
+              "type": "boolean",
               "default": false,
               "description": "If true, Go: Remove Tags command will prompt the user to provide tags and options instead of using the configured values"
             },
             "tags": {
-              "type":"string",
+              "type": "string",
               "default": "json",
               "description": "Comma separated tags to be used by Go: Remove Tags command"
             },
@@ -624,7 +629,7 @@
               "type": "boolean",
               "default": true,
               "description": "If true, adds command to run all tests in the current package to the editor context menu"
-            },          
+            },
             "generateTestForFunction": {
               "type": "boolean",
               "default": true,
@@ -701,7 +706,7 @@
         {
           "when": "editorTextFocus && config.go.editorContextMenuCommands.testPackage && resourceLangId == go",
           "command": "go.test.package"
-        },       
+        },
         {
           "when": "editorTextFocus && config.go.editorContextMenuCommands.generateTestForFunction && resourceLangId == go",
           "command": "go.test.generate.function"

src/goCodelens.ts

@@ -0,0 +1,68 @@
+'use strict';
+
+import vscode = require('vscode');
+import { CodeLensProvider, SymbolInformation, SymbolKind, TextDocument, CancellationToken, CodeLens, Range, Command, Location, commands } from 'vscode';
+import { documentSymbols, GoDocumentSymbolProvider } from './goOutline';
+import { GoReferenceProvider } from './goReferences';
+
+export class GoCodeLensProvider implements CodeLensProvider {
+	public provideCodeLenses(document: TextDocument, token: CancellationToken): CodeLens[] | Thenable<CodeLens[]> {
+		let codelensEnabled = vscode.workspace.getConfiguration('go').get('referencesCodeLens.enabled');
+		if (!codelensEnabled) {
+			return Promise.resolve([]);
+		}
+
+		return this.provideDocumentSymbols(document, token).then(symbols => {
+			let symbolReferences = symbols.map(symbol => this.provideSymbolReferences(document, symbol, token));
+			return Promise.all(symbolReferences).then(values => {
+				let codelenses = [];
+				values.forEach(lens => {
+					if (lens) {
+						codelenses.push(lens);
+					}
+				});
+				return codelenses;
+			});
+		});
+	}
+
+	private provideDocumentSymbols(document: TextDocument, token: CancellationToken): Thenable<vscode.SymbolInformation[]> {
+		let symbolProvider = new GoDocumentSymbolProvider();
+		return symbolProvider.provideDocumentSymbols(document, token).then(symbols => {
+			return symbols.filter(symbol =>
+				symbol.kind === vscode.SymbolKind.Function ||
+				symbol.kind === vscode.SymbolKind.Interface);
+		});
+	}
+
+	private provideSymbolReferences(document: TextDocument, symbol: SymbolInformation, token: CancellationToken): Thenable<CodeLens> {
+		if (token.isCancellationRequested) {
+			return Promise.resolve(null);
+		}
+
+		let options = {
+			includeDeclaration: false
+		};
+		let position = symbol.location.range.start;
+
+		// Add offset for functions due to go parser returns always 1 as the start character in a line
+		if (symbol.kind === vscode.SymbolKind.Function) {
+			position = position.translate(0, 5);
+		}
+		let referenceProvider = new GoReferenceProvider();
+		return referenceProvider.provideReferences(document, position, options, token).then(references => {
+			if (!references) {
+				return Promise.resolve(null);
+			}
+
+			let showReferences: Command = {
+				title: references.length === 1
+					? '1 reference'
+					: references.length + ' references',
+				command: 'editor.action.showReferences',
+				arguments: [document.uri, position, references]
+			};
+			return new CodeLens(symbol.location.range, showReferences);
+		});
+	}
+}

src/goMain.ts

@@ -33,6 +33,7 @@ import { LanguageClient } from 'vscode-languageclient';
 import { clearCacheForTools } from './goPath';
 import { addTags, removeTags } from './goModifytags';
 import { parseLiveFile } from './goLiveErrors';
+import { GoCodeLensProvider } from './goCodelens';
 
 export let errorDiagnosticCollection: vscode.DiagnosticCollection;
 let warningDiagnosticCollection: vscode.DiagnosticCollection;
@@ -82,6 +83,7 @@ export function activate(ctx: vscode.ExtensionContext): void {
 	ctx.subscriptions.push(vscode.languages.registerDocumentFormattingEditProvider(GO_MODE, new GoDocumentFormattingEditProvider()));
 	ctx.subscriptions.push(vscode.languages.registerRenameProvider(GO_MODE, new GoRenameProvider()));
 	ctx.subscriptions.push(vscode.languages.registerCodeActionsProvider(GO_MODE, new GoCodeActionProvider()));
+	ctx.subscriptions.push(vscode.languages.registerCodeLensProvider(GO_MODE, new GoCodeLensProvider()));
 
 	errorDiagnosticCollection = vscode.languages.createDiagnosticCollection('go-error');
 	ctx.subscriptions.push(errorDiagnosticCollection);

src/goReferences.ts

@@ -14,9 +14,7 @@ import { promptForMissingTool } from './goInstallTools';
 export class GoReferenceProvider implements vscode.ReferenceProvider {
 
 	public provideReferences(document: vscode.TextDocument, position: vscode.Position, options: { includeDeclaration: boolean }, token: vscode.CancellationToken): Thenable<vscode.Location[]> {
-		return vscode.workspace.saveAll(false).then(() => {
-			return this.doFindReferences(document, position, options, token);
-		});
+		return this.doFindReferences(document, position, options, token);
 	}
 
 	private doFindReferences(document: vscode.TextDocument, position: vscode.Position, options: { includeDeclaration: boolean }, token: vscode.CancellationToken): Thenable<vscode.Location[]> {
@@ -41,10 +39,12 @@ export class GoReferenceProvider implements vscode.ReferenceProvider {
 						promptForMissingTool('guru');
 						return resolve(null);
 					}
-					if (err) {
+
+					if (err && (<any>err).killed !== true) {
 						console.log(err);
 						return resolve(null);
 					}
+
 					let lines = stdout.toString().split('\n');
 					let results: vscode.Location[] = [];
 					for (let i = 0; i < lines.length; i++) {
@@ -53,6 +53,14 @@ export class GoReferenceProvider implements vscode.ReferenceProvider {
 						if (!match) continue;
 						let [_, file, lineStartStr, colStartStr, lineEndStr, colEndStr] = match;
 						let referenceResource = vscode.Uri.file(path.resolve(cwd, file));
+
+						if (!options.includeDeclaration) {
+							if (document.uri.fsPath === referenceResource.fsPath &&
+								position.line === Number(lineStartStr) - 1) {
+								continue;
+							}
+						}
+
 						let range = new vscode.Range(
 							+lineStartStr - 1, +colStartStr - 1, +lineEndStr - 1, +colEndStr
 						);