Add experimental support for Jump-to-def and Find-references

extensions/ql-vscode/src/bqrs-cli-types.ts

@@ -1,14 +1,34 @@
 
 export const PAGE_SIZE = 1000;
 
-export type ColumnKind = "f" | "i" | "s" | "b" | "d" | "e";
+/**
+ * The single-character codes used in the bqrs format for the the kind
+ * of a result column. This namespace is intentionally not an enum, see
+ * the "for the sake of extensibility" comment in messages.ts.
+ */
+// eslint-disable-next-line @typescript-eslint/no-namespace
+export namespace ColumnKindCode {
+  export const FLOAT = "f";
+  export const INTEGER = "i";
+  export const STRING = "s";
+  export const BOOLEAN = "b";
+  export const DATE = "d";
+  export const ENTITY = "e";
+}
+
+export type ColumnKind =
+  | typeof ColumnKindCode.FLOAT
+  | typeof ColumnKindCode.INTEGER
+  | typeof ColumnKindCode.STRING
+  | typeof ColumnKindCode.BOOLEAN
+  | typeof ColumnKindCode.DATE
+  | typeof ColumnKindCode.ENTITY;
 
 export interface Column {
   name?: string;
   kind: ColumnKind;
 }
 
-
 export interface ResultSetSchema {
   name: string;
   rows: number;

extensions/ql-vscode/src/cli.ts

@@ -10,10 +10,10 @@ import * as tk from 'tree-kill';
 import * as util from 'util';
 import { CancellationToken, Disposable } from 'vscode';
 import { BQRSInfo, DecodedBqrsChunk } from "./bqrs-cli-types";
-import { DistributionProvider } from './distribution';
-import { assertNever } from './helpers-pure';
-import { QueryMetadata, SortDirection } from './interface-types';
-import { Logger, ProgressReporter } from './logging';
+import { DistributionProvider } from "./distribution";
+import { assertNever } from "./helpers-pure";
+import { QueryMetadata, SortDirection } from "./interface-types";
+import { Logger, ProgressReporter } from "./logging";
 
 /**
  * The version of the SARIF format that we are using.
@@ -614,6 +614,26 @@ export class CodeQLCliServer implements Disposable {
       "Resolving qlpack information",
     );
   }
+
+  /**
+   * Gets information about queries in a query suite.
+   * @param additionalPacks A list of directories to search for qlpacks before searching in `searchPath`.
+   * @param searchPath A list of directories to search for packs not found in `additionalPacks`. If undefined,
+   *   the default CLI search path is used.
+   * @returns A list of query files found
+   */
+  resolveQueriesInSuite(suite: string, additionalPacks: string[], searchPath?: string[]): Promise<string[]> {
+    const args = ['--additional-packs', additionalPacks.join(path.delimiter)];
+    if (searchPath !== undefined) {
+      args.push('--search-path', path.join(...searchPath));
+    }
+    args.push(suite);
+    return this.runJsonCodeQlCliCommand<string[]>(
+      ['resolve', 'queries'],
+      args,
+      "Resolving queries",
+    );
+  }
 }
 
 /**

extensions/ql-vscode/src/config.ts

@@ -39,6 +39,17 @@ class Setting {
 
 const ROOT_SETTING = new Setting('codeQL');
 
+// Enable experimental features
+
+/**
+ * This setting is deliberately not in package.json so that it does
+ * not appear in the settings ui in vscode itself. If users want to
+ * enable experimental features, they can add
+ * "codeQl.experimentalFeatures" directly in their vscode settings
+ * json file.
+ */
+export const EXPERIMENTAL_FEATURES_SETTING = new Setting('experimentalFeatures', ROOT_SETTING);
+
 // Distribution configuration
 
 const DISTRIBUTION_SETTING = new Setting('cli', ROOT_SETTING);

extensions/ql-vscode/src/databases.ts

@@ -611,6 +611,11 @@ export class DatabaseManager extends DisposableObject {
     return this._databaseItems.find(item => item.databaseUri.toString(true) === uriString);
   }
 
+  public findDatabaseItemBySourceArchive(uri: vscode.Uri): DatabaseItem | undefined {
+    const uriString = uri.toString(true);
+    return this._databaseItems.find(item => item.sourceArchive && item.sourceArchive.toString(true) === uriString);
+  }
+
   private async addDatabaseItem(item: DatabaseItemImpl) {
     this._databaseItems.push(item);
     this.updatePersistedDatabaseList();

extensions/ql-vscode/src/definitions.ts

@@ -0,0 +1,202 @@
+import * as fs from 'fs-extra';
+import * as yaml from 'js-yaml';
+import * as tmp from 'tmp';
+import * as vscode from "vscode";
+import { decodeSourceArchiveUri, zipArchiveScheme } from "./archive-filesystem-provider";
+import { EntityValue, getResultSetSchema, LineColumnLocation, UrlValue, ColumnKindCode } from "./bqrs-cli-types";
+import { CodeQLCliServer } from "./cli";
+import { DatabaseItem, DatabaseManager } from "./databases";
+import * as helpers from './helpers';
+import { CachedOperation } from './helpers';
+import * as messages from "./messages";
+import { QueryServerClient } from "./queryserver-client";
+import { compileAndRunQueryAgainstDatabase, QueryWithResults } from "./run-queries";
+
+/**
+ * Run templated CodeQL queries to find definitions and references in
+ * source-language files. We may eventually want to find a way to
+ * generalize this to other custom queries, e.g. showing dataflow to
+ * or from a selected identifier.
+ */
+
+const TEMPLATE_NAME = "selectedSourceFile";
+const SELECT_QUERY_NAME = "#select";
+
+enum KeyType {
+  DefinitionQuery, ReferenceQuery
+}
+
+function tagOfKeyType(keyType: KeyType): string {
+  switch (keyType) {
+    case KeyType.DefinitionQuery: return "ide-contextual-queries/local-definitions";
+    case KeyType.ReferenceQuery: return "ide-contextual-queries/local-references";
+  }
+}
+
+async function resolveQueries(cli: CodeQLCliServer, qlpack: string, keyType: KeyType): Promise<string[]> {
+  const suiteFile = tmp.fileSync({ postfix: '.qls' }).name;
+  const suiteYaml = { qlpack, include: { kind: 'definitions', 'tags contain': tagOfKeyType(keyType) } };
+  await fs.writeFile(suiteFile, yaml.safeDump(suiteYaml), 'utf8');
+
+  const queries = await cli.resolveQueriesInSuite(suiteFile, helpers.getOnDiskWorkspaceFolders());
+  if (queries.length === 0) {
+    throw new Error("Couldn't find any queries for qlpack");
+  }
+  return queries;
+}
+
+async function qlpackOfDatabase(cli: CodeQLCliServer, db: DatabaseItem): Promise<string | undefined> {
+  if (db.contents === undefined)
+    return undefined;
+  const datasetPath = db.contents.datasetUri.fsPath;
+  const { qlpack } = await helpers.resolveDatasetFolder(cli, datasetPath);
+  return qlpack;
+}
+
+interface FullLocationLink extends vscode.LocationLink {
+  originUri: vscode.Uri;
+}
+
+export class TemplateQueryDefinitionProvider implements vscode.DefinitionProvider {
+  private cache: CachedOperation<vscode.LocationLink[]>;
+
+  constructor(
+    private cli: CodeQLCliServer,
+    private qs: QueryServerClient,
+    private dbm: DatabaseManager,
+  ) {
+    this.cache = new CachedOperation<vscode.LocationLink[]>(this.getDefinitions.bind(this));
+  }
+
+  async getDefinitions(uriString: string): Promise<vscode.LocationLink[]> {
+    return getLinksForUriString(this.cli, this.qs, this.dbm, uriString, (src, _dest) => src === uriString);
+  }
+
+  async provideDefinition(document: vscode.TextDocument, position: vscode.Position, _token: vscode.CancellationToken): Promise<vscode.LocationLink[]> {
+    const fileLinks = await this.cache.get(document.uri.toString());
+    const locLinks: vscode.LocationLink[] = [];
+    for (const link of fileLinks) {
+      if (link.originSelectionRange!.contains(position)) {
+        locLinks.push(link);
+      }
+    }
+    return locLinks;
+  }
+}
+
+export class TemplateQueryReferenceProvider implements vscode.ReferenceProvider {
+  private cache: CachedOperation<FullLocationLink[]>;
+
+  constructor(
+    private cli: CodeQLCliServer,
+    private qs: QueryServerClient,
+    private dbm: DatabaseManager,
+  ) {
+    this.cache = new CachedOperation<FullLocationLink[]>(this.getReferences.bind(this));
+  }
+
+  async getReferences(uriString: string): Promise<FullLocationLink[]> {
+    return getLinksForUriString(this.cli, this.qs, this.dbm, uriString, (_src, dest) => dest === uriString);
+  }
+
+  async provideReferences(document: vscode.TextDocument, position: vscode.Position, _context: vscode.ReferenceContext, _token: vscode.CancellationToken): Promise<vscode.Location[]> {
+    const fileLinks = await this.cache.get(document.uri.toString());
+    const locLinks: vscode.Location[] = [];
+    for (const link of fileLinks) {
+      if (link.targetRange!.contains(position)) {
+        locLinks.push({ range: link.originSelectionRange!, uri: link.originUri });
+      }
+    }
+    return locLinks;
+  }
+}
+
+interface FileRange {
+  file: vscode.Uri;
+  range: vscode.Range;
+}
+
+async function getLinksFromResults(results: QueryWithResults, cli: CodeQLCliServer, db: DatabaseItem, filter: (srcFile: string, destFile: string) => boolean): Promise<FullLocationLink[]> {
+  const localLinks: FullLocationLink[] = [];
+  const bqrsPath = results.query.resultsPaths.resultsPath;
+  const info = await cli.bqrsInfo(bqrsPath);
+  const selectInfo = getResultSetSchema(SELECT_QUERY_NAME, info);
+  if (selectInfo && selectInfo.columns.length == 3
+    && selectInfo.columns[0].kind == ColumnKindCode.ENTITY
+    && selectInfo.columns[1].kind == ColumnKindCode.ENTITY
+    && selectInfo.columns[2].kind == ColumnKindCode.STRING) {
+    // TODO: Page this
+    const allTuples = await cli.bqrsDecode(bqrsPath, SELECT_QUERY_NAME);
+    for (const tuple of allTuples.tuples) {
+      const src = tuple[0] as EntityValue;
+      const dest = tuple[1] as EntityValue;
+      const srcFile = src.url && fileRangeFromURI(src.url, db);
+      const destFile = dest.url && fileRangeFromURI(dest.url, db);
+      if (srcFile && destFile && filter(srcFile.file.toString(), destFile.file.toString())) {
+        localLinks.push({ targetRange: destFile.range, targetUri: destFile.file, originSelectionRange: srcFile.range, originUri: srcFile.file });
+      }
+    }
+  }
+  return localLinks;
+}
+
+async function getLinksForUriString(
+  cli: CodeQLCliServer,
+  qs: QueryServerClient,
+  dbm: DatabaseManager,
+  uriString: string,
+  filter: (src: string, dest: string) => boolean
+) {
+  const uri = decodeSourceArchiveUri(vscode.Uri.parse(uriString));
+  const sourceArchiveUri = vscode.Uri.file(uri.sourceArchiveZipPath).with({ scheme: zipArchiveScheme });
+
+  const db = dbm.findDatabaseItemBySourceArchive(sourceArchiveUri);
+  if (db) {
+    const qlpack = await qlpackOfDatabase(cli, db);
+    if (qlpack === undefined) {
+      throw new Error("Can't infer qlpack from database source archive");
+    }
+    const links: FullLocationLink[] = []
+    for (const query of await resolveQueries(cli, qlpack, KeyType.ReferenceQuery)) {
+      const templates: messages.TemplateDefinitions = {
+        [TEMPLATE_NAME]: {
+          values: {
+            tuples: [[{
+              stringValue: uri.pathWithinSourceArchive
+            }]]
+          }
+        }
+      };
+      const results = await compileAndRunQueryAgainstDatabase(cli, qs, db, false, vscode.Uri.file(query), templates);
+      if (results.result.resultType == messages.QueryResultType.SUCCESS) {
+        links.push(...await getLinksFromResults(results, cli, db, filter));
+      }
+    }
+    return links;
+  } else {
+    return [];
+  }
+}
+
+function fileRangeFromURI(uri: UrlValue, db: DatabaseItem): FileRange | undefined {
+  if (typeof uri === "string") {
+    return undefined;
+  } else if ('startOffset' in uri) {
+    return undefined;
+  } else {
+    const loc = uri as LineColumnLocation;
+    const range = new vscode.Range(Math.max(0, loc.startLine - 1),
+      Math.max(0, loc.startColumn - 1),
+      Math.max(0, loc.endLine - 1),
+      Math.max(0, loc.endColumn));
+    try {
+      const parsed = vscode.Uri.parse(uri.uri, true);
+      if (parsed.scheme === "file") {
+        return { file: db.resolveSourceFile(parsed.fsPath), range };
+      }
+      return undefined;
+    } catch (e) {
+      return undefined;
+    }
+  }
+}

extensions/ql-vscode/src/discovery.ts

@@ -84,4 +84,4 @@ export abstract class Discovery<T> extends DisposableObject {
    * @param results The discovery results returned by the `discover` function.
    */
   protected abstract update(results: T): void;
-}
+}

extensions/ql-vscode/src/extension.ts

@@ -1,25 +1,23 @@
-import { commands, Disposable, ExtensionContext, extensions, ProgressLocation, ProgressOptions, window as Window, Uri } from 'vscode';
+import { commands, Disposable, ExtensionContext, extensions, languages, ProgressLocation, ProgressOptions, Uri, window as Window } from 'vscode';
 import { LanguageClient } from 'vscode-languageclient';
+import { testExplorerExtensionId, TestHub } from 'vscode-test-adapter-api';
 import * as archiveFilesystemProvider from './archive-filesystem-provider';
-import { DistributionConfigListener, QueryServerConfigListener, QueryHistoryConfigListener } from './config';
+import { CodeQLCliServer } from './cli';
+import { DistributionConfigListener, QueryHistoryConfigListener, QueryServerConfigListener, EXPERIMENTAL_FEATURES_SETTING } from './config';
 import { DatabaseManager } from './databases';
 import { DatabaseUI } from './databases-ui';
-import {
-  DistributionUpdateCheckResultKind, DistributionManager, FindDistributionResult, FindDistributionResultKind, GithubApiError,
-  DEFAULT_DISTRIBUTION_VERSION_CONSTRAINT, GithubRateLimitedError
-} from './distribution';
+import { TemplateQueryDefinitionProvider, TemplateQueryReferenceProvider } from './definitions';
+import { DEFAULT_DISTRIBUTION_VERSION_CONSTRAINT, DistributionManager, DistributionUpdateCheckResultKind, FindDistributionResult, FindDistributionResultKind, GithubApiError, GithubRateLimitedError } from './distribution';
 import * as helpers from './helpers';
+import { assertNever } from './helpers-pure';
 import { spawnIdeServer } from './ide-server';
 import { InterfaceManager, WebviewReveal } from './interface';
 import { ideServerLogger, logger, queryServerLogger } from './logging';
-import { compileAndRunQueryAgainstDatabase, tmpDirDisposal, UserCancellationException } from './run-queries';
-import { CompletedQuery } from './query-results';
 import { QueryHistoryManager } from './query-history';
+import { CompletedQuery } from './query-results';
 import * as qsClient from './queryserver-client';
-import { CodeQLCliServer } from './cli';
-import { assertNever } from './helpers-pure';
 import { displayQuickQuery } from './quick-query';
-import { TestHub, testExplorerExtensionId } from 'vscode-test-adapter-api';
+import { compileAndRunQueryAgainstDatabase, tmpDirDisposal, UserCancellationException } from './run-queries';
 import { QLTestAdapterFactory } from './test-adapter';
 import { TestUIService } from './test-ui';
 
@@ -337,6 +335,17 @@ async function activateWithInstalledDistribution(ctx: ExtensionContext, distribu
   }));
 
   ctx.subscriptions.push(client.start());
+
+  if (EXPERIMENTAL_FEATURES_SETTING.getValue()) {
+    languages.registerDefinitionProvider(
+      { scheme: archiveFilesystemProvider.zipArchiveScheme },
+      new TemplateQueryDefinitionProvider(cliServer, qs, dbm)
+    );
+    languages.registerReferenceProvider(
+      { scheme: archiveFilesystemProvider.zipArchiveScheme },
+      new TemplateQueryReferenceProvider(cliServer, qs, dbm)
+    );
+  }
 }
 
 function initializeLogging(ctx: ExtensionContext): void {