Add code lens for quick eval command

Author: marcnjaramillo

extensions/ql-vscode/CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## [UNRELEASED]
 
+- Add a CodeLens to make the Quick Evaluation command more accessible. Click the `Quick Evaluation` prompt above a predicate definition in the editor to evaluate that predicate on its own. [#1035](https://github.com/github/vscode-codeql/pull/1035)
+  
 ## 1.5.8 - 2 December 2021
 
 - Emit a more explicit error message when a user tries to add a database with an unzipped source folder to the workspace. [#1021](https://github.com/github/vscode-codeql/pull/1021)

extensions/ql-vscode/src/extension.ts

@@ -11,7 +11,8 @@ import {
   window as Window,
   env,
   window,
-  QuickPickItem
+  QuickPickItem,
+  Range
 } from 'vscode';
 import { LanguageClient } from 'vscode-languageclient';
 import * as os from 'os';
@@ -21,6 +22,7 @@ import { testExplorerExtensionId, TestHub } from 'vscode-test-adapter-api';
 
 import { AstViewer } from './astViewer';
 import * as archiveFilesystemProvider from './archive-filesystem-provider';
+import QuickEvalCodeLensProvider from './quickEvalCodeLensProvider';
 import { CodeQLCliServer, CliVersionConstraint } from './cli';
 import {
   CliConfigListener,
@@ -156,6 +158,7 @@ export interface CodeQLExtensionInterface {
  * @returns CodeQLExtensionInterface
  */
 export async function activate(ctx: ExtensionContext): Promise<CodeQLExtensionInterface | Record<string, never>> {
+
   void logger.log(`Starting ${extensionId} extension`);
   if (extension === undefined) {
     throw new Error(`Can't find extension ${extensionId}`);
@@ -166,6 +169,9 @@ export async function activate(ctx: ExtensionContext): Promise<CodeQLExtensionIn
   await initializeTelemetry(extension, ctx);
   languageSupport.install();
 
+  const codelensProvider = new QuickEvalCodeLensProvider();
+  languages.registerCodeLensProvider({ scheme: 'file', language: 'ql' }, codelensProvider);
+
   ctx.subscriptions.push(distributionConfigListener);
   const codeQlVersionRange = DEFAULT_DISTRIBUTION_VERSION_RANGE;
   const distributionManager = new DistributionManager(distributionConfigListener, codeQlVersionRange, ctx);
@@ -471,6 +477,7 @@ async function activateWithInstalledDistribution(
     progress: ProgressCallback,
     token: CancellationToken,
     databaseItem: DatabaseItem | undefined,
+    range?: Range
   ): Promise<void> {
     if (qs !== undefined) {
       // If no databaseItem is specified, use the database currently selected in the Databases UI
@@ -485,7 +492,9 @@ async function activateWithInstalledDistribution(
         quickEval,
         selectedQuery,
         progress,
-        token
+        token,
+        undefined,
+        range
       );
       const item = qhm.buildCompletedQuery(info);
       await showResultsForCompletedQuery(item, WebviewReveal.NotForced);
@@ -733,6 +742,22 @@ async function activateWithInstalledDistribution(
         cancellable: true
       })
   );
+
+  ctx.subscriptions.push(
+    commandRunnerWithProgress(
+      'codeQL.codeLensQuickEval',
+      async (
+        progress: ProgressCallback,
+        token: CancellationToken,
+        uri: Uri,
+        range: Range
+      ) => await compileAndRunQuery(true, uri, progress, token, undefined, range),
+      {
+        title: 'Running query',
+        cancellable: true
+      })
+  );
+
   ctx.subscriptions.push(
     commandRunnerWithProgress('codeQL.quickQuery', async (
       progress: ProgressCallback,

extensions/ql-vscode/src/quickEvalCodeLensProvider.ts

@@ -0,0 +1,39 @@
+import {
+  CodeLensProvider,
+  TextDocument,
+  CodeLens, 
+  Command,
+  Range
+} from 'vscode';
+  
+class QuickEvalCodeLensProvider implements CodeLensProvider {
+  async provideCodeLenses(document: TextDocument): Promise<CodeLens[]> {
+
+    const codeLenses: CodeLens[] = [];
+
+    for (let index = 0; index < document.lineCount; index++) {
+      const textLine = document.lineAt(index);
+      // Match a predicate signature, including predicate name, parameter list, and opening brace.
+      const regex = new RegExp(/(\w+)\s*\(\s*.*(?:,\s*)*\)\s*\{/);
+      const matches = textLine.text.match(regex);
+
+      if (matches) {
+        const range: Range = new Range(
+          textLine.range.start.line, matches.index!,
+          textLine.range.end.line, matches.index! + 1
+        );
+
+        const command: Command = {
+          command: 'codeQL.codeLensQuickEval',
+          title: `Quick Evaluation: ${matches[1]}`,
+          arguments: [document.uri, range]
+        };
+        const codeLens = new CodeLens(range, command);
+        codeLenses.push(codeLens);
+      }
+    }
+    return codeLenses;
+  }
+}
+  
+export default QuickEvalCodeLensProvider;

extensions/ql-vscode/src/run-queries.ts

@@ -5,6 +5,7 @@ import * as tmp from 'tmp-promise';
 import {
   CancellationToken,
   ConfigurationTarget,
+  Range,
   TextDocument,
   TextEditor,
   Uri,
@@ -332,17 +333,18 @@ async function convertToQlPath(filePath: string): Promise<string> {
 
 
 /** Gets the selected position within the given editor. */
-async function getSelectedPosition(editor: TextEditor): Promise<messages.Position> {
-  const pos = editor.selection.start;
-  const posEnd = editor.selection.end;
-  // Convert from 0-based to 1-based line and column numbers.
-  return {
-    fileName: await convertToQlPath(editor.document.fileName),
-    line: pos.line + 1,
-    column: pos.character + 1,
-    endLine: posEnd.line + 1,
-    endColumn: posEnd.character + 1
-  };
+async function getSelectedPosition(editor: TextEditor, range?: Range): Promise<messages.Position> {
+    const selectedRange = range || editor.selection;
+    const pos = selectedRange.start;
+    const posEnd = selectedRange.end;
+    // Convert from 0-based to 1-based line and column numbers.
+    return {
+      fileName: await convertToQlPath(editor.document.fileName),
+      line: pos.line + 1,
+      column: pos.character + 1,
+      endLine: posEnd.line + 1,
+      endColumn: posEnd.character + 1
+    };
 }
 
 /**
@@ -490,7 +492,7 @@ type SelectedQuery = {
  * @param selectedResourceUri The selected resource when the command was run.
  * @param quickEval Whether the command being run is `Quick Evaluation`.
 */
-export async function determineSelectedQuery(selectedResourceUri: Uri | undefined, quickEval: boolean): Promise<SelectedQuery> {
+export async function determineSelectedQuery(selectedResourceUri: Uri | undefined, quickEval: boolean, range?: Range): Promise<SelectedQuery> {
   const editor = window.activeTextEditor;
 
   // Choose which QL file to use.
@@ -544,7 +546,7 @@ export async function determineSelectedQuery(selectedResourceUri: Uri | undefine
       // Report an error if we end up in this (hopefully unlikely) situation.
       throw new Error('The selected resource for quick evaluation should match the active editor.');
     }
-    quickEvalPosition = await getSelectedPosition(editor);
+    quickEvalPosition = await getSelectedPosition(editor, range);
     quickEvalText = editor.document.getText(editor.selection);
   }
 
@@ -560,13 +562,14 @@ export async function compileAndRunQueryAgainstDatabase(
   progress: ProgressCallback,
   token: CancellationToken,
   templates?: messages.TemplateDefinitions,
+  range?: Range
 ): Promise<QueryWithResults> {
   if (!db.contents || !db.contents.dbSchemeUri) {
     throw new Error(`Database ${db.databaseUri} does not have a CodeQL database scheme.`);
   }
 
   // Determine which query to run, based on the selection and the active editor.
-  const { queryPath, quickEvalPosition, quickEvalText } = await determineSelectedQuery(selectedQueryUri, quickEval);
+  const { queryPath, quickEvalPosition, quickEvalText } = await determineSelectedQuery(selectedQueryUri, quickEval, range);
 
   const historyItemOptions: QueryHistoryItemOptions = {};
   historyItemOptions.isQuickQuery === isQuickQueryPath(queryPath);