run benchmark at cursor (#972)

package.json

@@ -90,6 +90,11 @@
         "title": "Go: Test Function At Cursor",
         "description": "Runs a unit test at the cursor."
       },
+      {
+        "command": "go.benchmark.cursor",
+        "title": "Go: Benchmark Function At Cursor",
+        "description": "Runs a benchmark at the cursor."
+      },
       {
         "command": "go.test.file",
         "title": "Go: Test File",
@@ -825,6 +830,11 @@
           "command": "go.test.cursor",
           "group": "Go group 1"
         },
+        {
+          "when": "editorTextFocus && config.go.editorContextMenuCommands.benchmarkAtCursor && resourceLangId == go && !config.editor.codeLens",
+          "command": "go.benchmark.cursor",
+          "group": "Go group 1"
+        },
         {
           "when": "editorTextFocus && config.go.editorContextMenuCommands.testFile && resourceLangId == go",
           "command": "go.test.file",

src/goMain.ts

@@ -172,7 +172,14 @@ export function activate(ctx: vscode.ExtensionContext): void {
 
 	ctx.subscriptions.push(vscode.commands.registerCommand('go.test.cursor', (args) => {
 		let goConfig = vscode.workspace.getConfiguration('go', vscode.window.activeTextEditor ? vscode.window.activeTextEditor.document.uri : null);
-		testAtCursor(goConfig, args);
+		let isBenchmark = false;
+		testAtCursor(goConfig, isBenchmark, args);
+	}));
+
+	ctx.subscriptions.push(vscode.commands.registerCommand('go.benchmark.cursor', (args) => {
+		let goConfig = vscode.workspace.getConfiguration('go', vscode.window.activeTextEditor ? vscode.window.activeTextEditor.document.uri : null);
+		let isBenchmark = true;
+		testAtCursor(goConfig, isBenchmark, args);
 	}));
 
 	ctx.subscriptions.push(vscode.commands.registerCommand('go.test.package', (args) => {

src/goRunTestCodelens.ts

@@ -8,7 +8,7 @@
 import vscode = require('vscode');
 import path = require('path');
 import { CodeLensProvider, TextDocument, CancellationToken, CodeLens, Command } from 'vscode';
-import { getTestFunctions, getTestEnvVars, getTestFlags } from './testUtils';
+import { getTestFunctions, getBenchmarkFunctions, getTestEnvVars, getTestFlags } from './testUtils';
 import { GoDocumentSymbolProvider } from './goOutline';
 import { getCurrentGoPath } from './util';
 import { GoBaseCodeLensProvider } from './goBaseCodelens';
@@ -65,14 +65,14 @@ export class GoRunTestCodeLensProvider extends GoBaseCodeLensProvider {
 	}
 
 	private getCodeLensForFunctions(vsConfig: vscode.WorkspaceConfiguration, document: TextDocument): Thenable<CodeLens[]> {
-		return getTestFunctions(document).then(testFunctions => {
-			let codelens = [];
+		const codelens: CodeLens[] = [];
 
+		const testPromise = getTestFunctions(document).then(testFunctions => {
 			testFunctions.forEach(func => {
 				let runTestCmd: Command = {
 					title: 'run test',
 					command: 'go.test.cursor',
-					arguments: [ { functionName: func.name} ]
+					arguments: [ { functionName: func.name } ]
 				};
 
 				const args = ['-test.run', func.name];
@@ -95,7 +95,21 @@ export class GoRunTestCodeLensProvider extends GoBaseCodeLensProvider {
 				codelens.push(new CodeLens(func.location.range, runTestCmd));
 				codelens.push(new CodeLens(func.location.range, debugTestCmd));
 			});
-			return codelens;
 		});
+
+		const benchmarkPromise = getBenchmarkFunctions(document).then(benchmarkFunctions => {
+			benchmarkFunctions.forEach(func => {
+				let runBenchmarkCmd: Command = {
+					title: 'run benchmark',
+					command: 'go.benchmark.cursor',
+					arguments: [ { functionName: func.name, isBenchmark: true } ]
+				};
+
+				codelens.push(new CodeLens(func.location.range, runBenchmarkCmd));
+			});
+
+		});
+
+		return Promise.all([testPromise, benchmarkPromise]).then(() => codelens);
 	}
 }

src/goTest.ts

@@ -8,7 +8,7 @@
 import path = require('path');
 import vscode = require('vscode');
 import os = require('os');
-import { goTest, TestConfig, getTestEnvVars, getTestFlags, getTestFunctions } from './testUtils';
+import { goTest, TestConfig, getTestEnvVars, getTestFlags, getTestFunctions, getBenchmarkFunctions } from './testUtils';
 import { getCoverage } from './goCover';
 
 // lastTestConfig holds a reference to the last executed TestConfig which allows
@@ -21,7 +21,7 @@ let lastTestConfig: TestConfig;
 *
 * @param goConfig Configuration for the Go extension.
 */
-export function testAtCursor(goConfig: vscode.WorkspaceConfiguration, args: any) {
+export function testAtCursor(goConfig: vscode.WorkspaceConfiguration, isBenchmark: boolean, args: any) {
 	let editor = vscode.window.activeTextEditor;
 	if (!editor) {
 		vscode.window.showInformationMessage('No editor is active.');
@@ -31,39 +31,44 @@ export function testAtCursor(goConfig: vscode.WorkspaceConfiguration, args: any)
 		vscode.window.showInformationMessage('No tests found. Current file is not a test file.');
 		return;
 	}
+
+	const getFunctions = isBenchmark ? getBenchmarkFunctions : getTestFunctions;
+
 	editor.document.save().then(() => {
-		return getTestFunctions(editor.document).then(testFunctions => {
-		let testFunctionName: string;
-
-		// We use functionName if it was provided as argument
-		// Otherwise find any test function containing the cursor.
-		if (args && args.functionName) {
-			testFunctionName = args.functionName;
-		} else {
-			for (let func of testFunctions) {
-				let selection = editor.selection;
-				if (selection && func.location.range.contains(selection.start)) {
-					testFunctionName = func.name;
-					break;
-				}
-			};
-		}
+		return getFunctions(editor.document).then(testFunctions => {
+			let testFunctionName: string;
+
+			// We use functionName if it was provided as argument
+			// Otherwise find any test function containing the cursor.
+			if (args && args.functionName) {
+				testFunctionName = args.functionName;
+			} else {
+				for (let func of testFunctions) {
+					let selection = editor.selection;
+					if (selection && func.location.range.contains(selection.start)) {
+						testFunctionName = func.name;
+						break;
+					}
+				};
+			}
+
+			if (!testFunctionName) {
+				vscode.window.showInformationMessage('No test function found at cursor.');
+				return;
+			}
 
-		if (!testFunctionName) {
-			vscode.window.showInformationMessage('No test function found at cursor.');
-			return;
-		}
+			const testConfig = {
+				goConfig: goConfig,
+				dir: path.dirname(editor.document.fileName),
+				flags: getTestFlags(goConfig, args),
+				functions: [testFunctionName],
+				isBenchmark: isBenchmark
+			};
 
-		const testConfig = {
-			goConfig: goConfig,
-			dir: path.dirname(editor.document.fileName),
-			flags: getTestFlags(goConfig, args),
-			functions: [testFunctionName]
-		};
-		// Remember this config as the last executed test.
-		lastTestConfig = testConfig;
+			// Remember this config as the last executed test.
+			lastTestConfig = testConfig;
 
-		return goTest(testConfig);
+			return goTest(testConfig);
 		});
 	}).then(null, err => {
 		console.error(err);

src/testUtils.ts

@@ -43,6 +43,10 @@ export interface TestConfig {
 	 * Run all tests from all sub directories under `dir`
 	 */
 	includeSubDirectories?: boolean;
+	/**
+	 * Whether this is a benchmark.
+	 */
+	isBenchmark?: boolean;
 }
 
 export function getTestEnvVars(config: vscode.WorkspaceConfiguration): any {
@@ -85,19 +89,25 @@ export function getTestFunctions(doc: vscode.TextDocument): Thenable<vscode.Symb
 		.then(symbols =>
 			symbols.filter(sym =>
 				sym.kind === vscode.SymbolKind.Function
-				&& hasTestFunctionPrefix(sym.name))
+				&& (sym.name.startsWith('Test') || sym.name.startsWith('Example')))
 		);
 }
 
 /**
- * Returns whether a given function name has a test prefix.
- * Test functions have "Test" or "Example" as a prefix.
+ * Returns all Benchmark functions in the given source file.
  *
- * @param the function name.
- * @return whether the name has a test function prefix.
+ * @param the URI of a Go source file.
+ * @return benchmark function symbols for the source file.
  */
-function hasTestFunctionPrefix(name: string): boolean {
-	return name.startsWith('Test') || name.startsWith('Example');
+export function getBenchmarkFunctions(doc: vscode.TextDocument): Thenable<vscode.SymbolInformation[]> {
+	let documentSymbolProvider = new GoDocumentSymbolProvider();
+	return documentSymbolProvider
+		.provideDocumentSymbols(doc, null)
+		.then(symbols =>
+			symbols.filter(sym =>
+				sym.kind === vscode.SymbolKind.Function
+				&& sym.name.startsWith('Benchmark'))
+		);
 }
 
 /**
@@ -114,11 +124,18 @@ export function goTest(testconfig: TestConfig): Thenable<boolean> {
 		}
 
 		let buildTags: string = testconfig.goConfig['buildTags'];
-		let args = ['test', ...testconfig.flags, '-timeout', testconfig.goConfig['testTimeout']];
+		let args: Array<string> = ['test', ...testconfig.flags];
+		let testType: string = testconfig.isBenchmark ? 'Benchmarks' : 'Tests';
+
+		if (testconfig.isBenchmark) {
+			args.push('-benchmem', '-run=^$');
+		} else {
+			args.push('-timeout', testconfig.goConfig['testTimeout']);
+		}
 		if (buildTags && testconfig.flags.indexOf('-tags') === -1) {
-			args.push('-tags');
-			args.push(buildTags);
+			args.push('-tags', buildTags);
 		}
+
 		let testEnvVars = getTestEnvVars(testconfig.goConfig);
 		let goRuntimePath = getGoRuntimePath();
 
@@ -163,14 +180,14 @@ export function goTest(testconfig: TestConfig): Thenable<boolean> {
 				errBuf.done();
 
 				if (code) {
-					outputChannel.appendLine('Error: Tests failed.');
+					outputChannel.appendLine(`Error: ${testType} failed.`);
 				} else {
-					outputChannel.appendLine('Success: Tests passed.');
+					outputChannel.appendLine(`Success: ${testType} passed.`);
 				}
 				resolve(code === 0);
 			});
 		}, err => {
-			outputChannel.appendLine('Error: Tests failed.');
+			outputChannel.appendLine(`Error: ${testType} failed.`);
 			outputChannel.appendLine(err);
 			resolve(false);
 		});
@@ -195,15 +212,22 @@ function expandFilePathInOutput(output: string, cwd: string): string {
 	return lines.join('\n');
 }
 
+/**
+ * Function type for getting the target arguments.
+ *
+ * @param testconfig Configuration for the Go extension.
+ */
+type argsHandleFunc = (testconfig: TestConfig) => Thenable<Array<string>>;
+
 /**
  * Get the test target arguments.
  *
  * @param testconfig Configuration for the Go extension.
  */
 function targetArgs(testconfig: TestConfig): Thenable<Array<string>> {
 	if (testconfig.functions) {
-		return Promise.resolve(['-run', util.format('^%s$', testconfig.functions.join('|'))]);
-	} else if (testconfig.includeSubDirectories) {
+		return Promise.resolve([testconfig.isBenchmark ? '-bench' : '-run', util.format('^%s$', testconfig.functions.join('|'))]);
+	} else if (testconfig.includeSubDirectories && !testconfig.isBenchmark) {
 		return getGoVersion().then((ver: SemVersion) => {
 			if (ver && (ver.major > 1 || (ver.major === 1 && ver.minor >= 9))) {
 				return ['./...'];
@@ -213,3 +237,4 @@ function targetArgs(testconfig: TestConfig): Thenable<Array<string>> {
 	}
 	return Promise.resolve([]);
 }
+