add codelens for running main func

docs/settings.md

@@ -210,12 +210,14 @@ Feature level setting to enable/disable code lens for references and run/debug t
 | --- | --- |
 | `references` | If true, enables the references code lens. Uses guru. Recalculates when there is change to the document followed by scrolling. Unnecessary when using the language server; use the call graph feature instead. <br/> Default: `false` |
 | `runtest` | If true, enables code lens for running and debugging tests <br/> Default: `true` |
+| `runmain` | If true, enables code lens for running main func <br/> Default: `true` |
 
 Default:
 ```
 {
 	"references" :	false,
 	"runtest" :	true,
+	"runmain" :	true,
 }
 ```
 ### `go.formatFlags`
@@ -497,7 +499,7 @@ Allowed Options: `package`, `workspace`, `off`
 Default: `"package"`
 ### `gopls`
 
-Customize `gopls` behavior by specifying the gopls' settings in this section. For example, 
+Customize `gopls` behavior by specifying the gopls' settings in this section. For example,
 ```
 "gopls" : {
 	"build.directoryFilters": ["-node_modules"]

package.json

@@ -489,6 +489,11 @@
         "title": "Go: Reset Global State",
         "description": "Reset keys in global state to undefined."
       },
+      {
+        "command": "go.runMain",
+        "title": "Go: Run main() func on file",
+        "description": "Run main() func on file"
+      },
       {
         "command": "go.explorer.refresh",
         "title": "Go Explorer: Refresh",
@@ -1606,12 +1611,18 @@
               "type": "boolean",
               "default": true,
               "description": "If true, enables code lens for running and debugging tests"
+            },
+            "runmain": {
+              "type": "boolean",
+              "default": true,
+              "description": "If true, enables code lens for running main func"
             }
           },
           "additionalProperties": false,
           "default": {
             "references": false,
-            "runtest": true
+            "runtest": true,
+            "runmain": true
           },
           "description": "Feature level setting to enable/disable code lens for references and run/debug tests",
           "scope": "resource"

src/goMain.ts

@@ -60,6 +60,7 @@ import { GO111MODULE, goModInit, isModSupported } from './goModules';
 import { playgroundCommand } from './goPlayground';
 import { GoReferencesCodeLensProvider } from './goReferencesCodelens';
 import { GoRunTestCodeLensProvider } from './goRunTestCodelens';
+import { GoMainCodeLensProvider, runMainFunc } from './goMainCodelens';
 import { disposeGoStatusBar, expandGoStatusBar, outputChannel, updateGoStatusBar } from './goStatus';
 import {
 	debugPrevious,
@@ -210,9 +211,11 @@ If you would like additional configuration for diagnostics from gopls, please se
 		})
 	);
 	const testCodeLensProvider = new GoRunTestCodeLensProvider();
+	const mainCodeLensProvider = new GoMainCodeLensProvider();
 	const referencesCodeLensProvider = new GoReferencesCodeLensProvider();
 
 	ctx.subscriptions.push(vscode.languages.registerCodeLensProvider(GO_MODE, testCodeLensProvider));
+	ctx.subscriptions.push(vscode.languages.registerCodeLensProvider(GO_MODE, mainCodeLensProvider));
 	ctx.subscriptions.push(vscode.languages.registerCodeLensProvider(GO_MODE, referencesCodeLensProvider));
 
 	// debug
@@ -481,6 +484,7 @@ If you would like additional configuration for diagnostics from gopls, please se
 
 			if (updatedGoConfig['enableCodeLens']) {
 				testCodeLensProvider.setEnabled(updatedGoConfig['enableCodeLens']['runtest']);
+				mainCodeLensProvider.setEnabled(updatedGoConfig['enableCodeLens']['runmain']);
 				referencesCodeLensProvider.setEnabled(updatedGoConfig['enableCodeLens']['references']);
 			}
 
@@ -564,6 +568,12 @@ If you would like additional configuration for diagnostics from gopls, please se
 		})
 	);
 
+	ctx.subscriptions.push(
+		vscode.commands.registerCommand('go.runMain', (args) => {
+			runMainFunc()
+		})
+	)
+
 	ctx.subscriptions.push(
 		vscode.commands.registerCommand('go.show.commands', () => {
 			const extCommands = getExtensionCommands();

src/goMainCodelens.ts

@@ -0,0 +1,121 @@
+/* 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 cp = require('child_process');
+
+import { CancellationToken, CodeLens, TextDocument } from 'vscode';
+import { getGoConfig } from './config';
+import { GoBaseCodeLensProvider } from './goBaseCodelens';
+import { GoDocumentSymbolProvider } from './goOutline';
+import { getBinPath } from './util';
+import { envPath, getCurrentGoRoot } from './utils/pathUtils';
+import { reject } from 'lodash';
+
+export class GoMainCodeLensProvider extends GoBaseCodeLensProvider {
+	private readonly mainRegex = /^main$/;
+
+	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['runmain'] : false;
+		if (!codelensEnabled || !document.fileName.match('main.go')) {
+			return [];
+		}
+
+		const codelenses = await Promise.all([
+			this.getCodeLensForMainFunc(document, token)
+		]);
+		return ([] as CodeLens[]).concat(...codelenses);
+	}
+
+	// Return the first main function
+	private async getMainFunc(
+		doc: vscode.TextDocument,
+		token: vscode.CancellationToken
+	): Promise<vscode.DocumentSymbol | undefined> {
+		const documentSymbolProvider = new GoDocumentSymbolProvider(true);
+		const symbols = await documentSymbolProvider.provideDocumentSymbols(doc, token);
+		if (!symbols || symbols.length === 0) {
+			return;
+		}
+		const symbol = symbols[0];
+		if (!symbol) {
+			return;
+		}
+		const children = symbol.children;
+
+		return children.find(sym => sym.kind === vscode.SymbolKind.Function && this.mainRegex.test(sym.name));
+	}
+
+	private async getCodeLensForMainFunc(document: TextDocument, token: CancellationToken): Promise<CodeLens[]> {
+		const mainPromise = async (): Promise<CodeLens[]> => {
+			const mainFunc = await this.getMainFunc(document, token);
+			if (!mainFunc) {
+				return [];
+			}
+
+			return [
+				new CodeLens(mainFunc.range, {
+					title: 'run main',
+					command: 'go.runMain',
+					arguments: [{ functionName: mainFunc.name }]
+				})
+			];
+		};
+
+		return await mainPromise();
+	}
+}
+
+const mainFuncOutputChannel = vscode.window.createOutputChannel('Go Main');
+
+export async function runMainFunc() {
+	let outputChannel = mainFuncOutputChannel
+	const goRuntimePath = getBinPath('go');
+	if (!goRuntimePath) {
+		vscode.window.showErrorMessage(
+			`Failed to run "go run ." as the "go" binary cannot be found in either GOROOT(${getCurrentGoRoot}) or PATH(${envPath})`
+		);
+		return Promise.resolve(false);
+	}
+
+	const editor = vscode.window.activeTextEditor;
+	const documentUri = editor ? editor.document.uri : null;
+	const args = ['run', documentUri.path];
+
+	outputChannel.clear()
+	outputChannel.show(true)
+	outputChannel.appendLine(["Running main func: ", goRuntimePath, ...args].join(' '))
+
+	cp.execFile(
+		goRuntimePath,
+		args,
+		{ },
+		(err, stdout, stderr) => {
+			try {
+				if (err) {
+					outputChannel.appendLine(err.message);
+					return;
+				}
+				if (stdout) {
+					outputChannel.append(stdout);
+				}
+				if (stderr) {
+					outputChannel.append(stderr);
+				}
+			} catch (e) {
+				reject(e);
+			}
+		}
+	)
+}