Fix error messages for ast viewers and update caching

This commit does two things:

1. Add more appropriate error messages when asts can't be viewed.
2. Make better use of cached operations for asts. In the past, we were
not actually using cached operations. Each time an ast view request
occurred, we created a new TemplatePrintAstProvider instance. With this
change, we reuse the TemplatePrintAstProvider between calls and ensure
that an AST that is called once is reused on subsequent calls.

Author: aeisenberg

extensions/ql-vscode/CHANGELOG.md

@@ -3,6 +3,8 @@
 ## [UNRELEASED]
 
 - Avoid displaying an error when removing orphaned databases and the storage folder does not exist. [#748](https://github.com/github/vscode-codeql/pull/748)
+- Add better error messages when AST Viewer is unable to create an AST. [#753](https://github.com/github/vscode-codeql/pull/753)
+- Cache AST viewing operations so that subsequent calls to view the AST of a single file will be extremely fast. [#753](https://github.com/github/vscode-codeql/pull/753)
 
 ## 1.4.2 - 2 February 2021
 

extensions/ql-vscode/src/contextual/templateProvider.ts

@@ -1,4 +1,15 @@
-import * as vscode from 'vscode';
+import {
+  CancellationToken,
+  DefinitionProvider,
+  Location,
+  LocationLink,
+  Position,
+  ProgressLocation,
+  ReferenceContext,
+  ReferenceProvider,
+  TextDocument,
+  Uri
+} from 'vscode';
 
 import { decodeSourceArchiveUri, encodeArchiveBasePath, zipArchiveScheme } from '../archive-filesystem-provider';
 import { CodeQLCliServer } from '../cli';
@@ -22,20 +33,20 @@ import { qlpackOfDatabase, resolveQueries } from './queryResolver';
  * or from a selected identifier.
  */
 
-export class TemplateQueryDefinitionProvider implements vscode.DefinitionProvider {
-  private cache: CachedOperation<vscode.LocationLink[]>;
+export class TemplateQueryDefinitionProvider implements DefinitionProvider {
+  private cache: CachedOperation<LocationLink[]>;
 
   constructor(
     private cli: CodeQLCliServer,
     private qs: QueryServerClient,
     private dbm: DatabaseManager,
   ) {
-    this.cache = new CachedOperation<vscode.LocationLink[]>(this.getDefinitions.bind(this));
+    this.cache = new CachedOperation<LocationLink[]>(this.getDefinitions.bind(this));
   }
 
-  async provideDefinition(document: vscode.TextDocument, position: vscode.Position, _token: vscode.CancellationToken): Promise<vscode.LocationLink[]> {
+  async provideDefinition(document: TextDocument, position: Position, _token: CancellationToken): Promise<LocationLink[]> {
     const fileLinks = await this.cache.get(document.uri.toString());
-    const locLinks: vscode.LocationLink[] = [];
+    const locLinks: LocationLink[] = [];
     for (const link of fileLinks) {
       if (link.originSelectionRange!.contains(position)) {
         locLinks.push(link);
@@ -44,9 +55,9 @@ export class TemplateQueryDefinitionProvider implements vscode.DefinitionProvide
     return locLinks;
   }
 
-  private async getDefinitions(uriString: string): Promise<vscode.LocationLink[]> {
+  private async getDefinitions(uriString: string): Promise<LocationLink[]> {
     return withProgress({
-      location: vscode.ProgressLocation.Notification,
+      location: ProgressLocation.Notification,
       cancellable: true,
       title: 'Finding definitions'
     }, async (progress, token) => {
@@ -64,7 +75,7 @@ export class TemplateQueryDefinitionProvider implements vscode.DefinitionProvide
   }
 }
 
-export class TemplateQueryReferenceProvider implements vscode.ReferenceProvider {
+export class TemplateQueryReferenceProvider implements ReferenceProvider {
   private cache: CachedOperation<FullLocationLink[]>;
 
   constructor(
@@ -76,13 +87,13 @@ export class TemplateQueryReferenceProvider implements vscode.ReferenceProvider
   }
 
   async provideReferences(
-    document: vscode.TextDocument,
-    position: vscode.Position,
-    _context: vscode.ReferenceContext,
-    _token: vscode.CancellationToken
-  ): Promise<vscode.Location[]> {
+    document: TextDocument,
+    position: Position,
+    _context: ReferenceContext,
+    _token: CancellationToken
+  ): Promise<Location[]> {
     const fileLinks = await this.cache.get(document.uri.toString());
-    const locLinks: vscode.Location[] = [];
+    const locLinks: Location[] = [];
     for (const link of fileLinks) {
       if (link.targetRange!.contains(position)) {
         locLinks.push({ range: link.originSelectionRange!, uri: link.originUri });
@@ -93,7 +104,7 @@ export class TemplateQueryReferenceProvider implements vscode.ReferenceProvider
 
   private async getReferences(uriString: string): Promise<FullLocationLink[]> {
     return withProgress({
-      location: vscode.ProgressLocation.Notification,
+      location: ProgressLocation.Notification,
       cancellable: true,
       title: 'Finding references'
     }, async (progress, token) => {
@@ -112,40 +123,41 @@ export class TemplateQueryReferenceProvider implements vscode.ReferenceProvider
 }
 
 export class TemplatePrintAstProvider {
-  private cache: CachedOperation<QueryWithResults | undefined>;
+  private cache: CachedOperation<QueryWithResults>;
 
   constructor(
     private cli: CodeQLCliServer,
     private qs: QueryServerClient,
     private dbm: DatabaseManager,
-
-    // Note: progress and token are only used if a cached value is not available
-    private progress: ProgressCallback,
-    private token: vscode.CancellationToken
   ) {
-    this.cache = new CachedOperation<QueryWithResults | undefined>(this.getAst.bind(this));
+    this.cache = new CachedOperation<QueryWithResults>(this.getAst.bind(this));
   }
 
-  async provideAst(document?: vscode.TextDocument): Promise<AstBuilder | undefined> {
+  async provideAst(
+    progress: ProgressCallback,
+    token: CancellationToken,
+    document?: TextDocument
+  ): Promise<AstBuilder | undefined> {
     if (!document) {
-      return;
-    }
-    const queryResults = await this.cache.get(document.uri.toString());
-    if (!queryResults) {
-      return;
+      throw new Error('Cannot view the AST. Please select a valid source file inside a CodeQL database.');
     }
+    const queryResults = await this.cache.get(document.uri.toString(), progress, token);
 
     return new AstBuilder(
       queryResults, this.cli,
-      this.dbm.findDatabaseItem(vscode.Uri.parse(queryResults.database.databaseUri!, true))!,
+      this.dbm.findDatabaseItem(Uri.parse(queryResults.database.databaseUri!, true))!,
       document.fileName
     );
   }
 
-  private async getAst(uriString: string): Promise<QueryWithResults> {
-    const uri = vscode.Uri.parse(uriString, true);
+  private async getAst(
+    uriString: string,
+    progress: ProgressCallback,
+    token: CancellationToken
+  ): Promise<QueryWithResults> {
+    const uri = Uri.parse(uriString, true);
     if (uri.scheme !== zipArchiveScheme) {
-      throw new Error('AST Viewing is only available for databases with zipped source archives.');
+      throw new Error('Cannot view the AST. Please select a valid source file inside a CodeQL database.');
     }
 
     const zippedArchive = decodeSourceArchiveUri(uri);
@@ -181,9 +193,9 @@ export class TemplatePrintAstProvider {
       this.qs,
       db,
       false,
-      vscode.Uri.file(query),
-      this.progress,
-      this.token,
+      Uri.file(query),
+      progress,
+      token,
       templates
     );
   }

extensions/ql-vscode/src/extension.ts

@@ -692,13 +692,18 @@ async function activateWithInstalledDistribution(
   );
 
   const astViewer = new AstViewer();
+  const templateProvider = new TemplatePrintAstProvider(cliServer, qs, dbm);
+
   ctx.subscriptions.push(astViewer);
   ctx.subscriptions.push(commandRunnerWithProgress('codeQL.viewAst', async (
     progress: ProgressCallback,
     token: CancellationToken
   ) => {
-    const ast = await new TemplatePrintAstProvider(cliServer, qs, dbm, progress, token)
-      .provideAst(window.activeTextEditor?.document);
+    const ast = await templateProvider.provideAst(
+      progress,
+      token,
+      window.activeTextEditor?.document,
+    );
     if (ast) {
       astViewer.updateRoots(await ast.getRoots(), ast.db, ast.fileName);
     }

extensions/ql-vscode/src/helpers.ts

@@ -296,19 +296,19 @@ export async function getPrimaryDbscheme(datasetFolder: string): Promise<string>
  * A cached mapping from strings to value of type U.
  */
 export class CachedOperation<U> {
-  private readonly operation: (t: string) => Promise<U>;
+  private readonly operation: (t: string, ...args: any[]) => Promise<U>;
   private readonly cached: Map<string, U>;
   private readonly lru: string[];
   private readonly inProgressCallbacks: Map<string, [(u: U) => void, (reason?: any) => void][]>;
 
-  constructor(operation: (t: string) => Promise<U>, private cacheSize = 100) {
+  constructor(operation: (t: string, ...args: any[]) => Promise<U>, private cacheSize = 100) {
     this.operation = operation;
     this.lru = [];
     this.inProgressCallbacks = new Map<string, [(u: U) => void, (reason?: any) => void][]>();
     this.cached = new Map<string, U>();
   }
 
-  async get(t: string): Promise<U> {
+  async get(t: string, ...args: any[]): Promise<U> {
     // Try and retrieve from the cache
     const fromCache = this.cached.get(t);
     if (fromCache !== undefined) {
@@ -329,7 +329,7 @@ export class CachedOperation<U> {
     const callbacks: [(u: U) => void, (reason?: any) => void][] = [];
     this.inProgressCallbacks.set(t, callbacks);
     try {
-      const result = await this.operation(t);
+      const result = await this.operation(t, ...args);
       callbacks.forEach(f => f[0](result));
       this.inProgressCallbacks.delete(t);
       if (this.lru.length > this.cacheSize) {