Select the appropriate node in the AST viewer when the editor text selection changes

When a user clicks in an editor that whose source tree is currently being displayed in
the ast viewer, the viewer selection will stay in sync with the editor selection.

Author: aeisenberg

extensions/ql-vscode/package-lock.json

@@ -704,7 +704,6 @@
     "zip-a-folder": "~0.0.12"
   },
   "devDependencies": {
-    "@types/semver": "~7.2.0",
     "@types/chai": "^4.1.7",
     "@types/chai-as-promised": "~7.1.2",
     "@types/child-process-promise": "^2.2.1",
@@ -714,7 +713,7 @@
     "@types/google-protobuf": "^3.2.7",
     "@types/gulp": "^4.0.6",
     "@types/gulp-sourcemaps": "0.0.32",
-    "@types/js-yaml": "~3.12.2",
+    "@types/js-yaml": "^3.12.5",
     "@types/jszip": "~3.1.6",
     "@types/mocha": "~8.0.3",
     "@types/node": "^12.0.8",
@@ -723,6 +722,7 @@
     "@types/react": "^16.8.17",
     "@types/react-dom": "^16.8.4",
     "@types/sarif": "~2.1.2",
+    "@types/semver": "~7.2.0",
     "@types/sinon": "~7.5.2",
     "@types/sinon-chai": "~3.2.3",
     "@types/through2": "^2.0.36",

extensions/ql-vscode/package.json

@@ -1,68 +1,81 @@
-import * as vscode from 'vscode';
+import {
+  window,
+  ExtensionContext,
+  TreeDataProvider,
+  EventEmitter,
+  commands,
+  Event,
+  ProviderResult,
+  TreeItemCollapsibleState,
+  TreeItem,
+  TreeView,
+  TextEditorSelectionChangeEvent,
+  Location,
+  Range
+} from 'vscode';
+import * as path from 'path';
 
 import { DatabaseItem } from './databases';
 import { UrlValue, BqrsId } from './bqrs-cli-types';
-import fileRangeFromURI from './contextual/fileRangeFromURI';
 import { showLocation } from './interface-utils';
 import { isStringLoc, isWholeFileLoc, isLineColumnLoc } from './bqrs-utils';
 
 export interface AstItem {
   id: BqrsId;
   label?: string;
   location?: UrlValue;
+  fileLocation?: Location;
   parent: AstItem | RootAstItem;
   children: AstItem[];
   order: number;
 }
 
 export type RootAstItem = Omit<AstItem, 'parent'>;
 
-class AstViewerDataProvider implements vscode.TreeDataProvider<AstItem | RootAstItem> {
+class AstViewerDataProvider implements TreeDataProvider<AstItem | RootAstItem> {
 
   public roots: RootAstItem[] = [];
   public db: DatabaseItem | undefined;
 
   private _onDidChangeTreeData =
-    new vscode.EventEmitter<AstItem | undefined>();
-  readonly onDidChangeTreeData: vscode.Event<AstItem | undefined> =
+    new EventEmitter<AstItem | undefined>();
+  readonly onDidChangeTreeData: Event<AstItem | undefined> =
     this._onDidChangeTreeData.event;
 
   constructor() {
-    vscode.commands.registerCommand('codeQLAstViewer.gotoCode',
-      async (location: UrlValue, db: DatabaseItem) => {
-        if (location) {
-          await showLocation(fileRangeFromURI(location, db));
-        }
+    commands.registerCommand('codeQLAstViewer.gotoCode',
+      async (item: AstItem) => {
+        await showLocation(item.fileLocation);
       });
   }
 
   refresh(): void {
     this._onDidChangeTreeData.fire();
   }
-  getChildren(item?: AstItem): vscode.ProviderResult<(AstItem | RootAstItem)[]> {
+  getChildren(item?: AstItem): ProviderResult<(AstItem | RootAstItem)[]> {
     const children = item ? item.children : this.roots;
     return children.sort((c1, c2) => (c1.order - c2.order));
   }
 
-  getParent(item: AstItem): vscode.ProviderResult<AstItem> {
+  getParent(item: AstItem): ProviderResult<AstItem> {
     return item.parent as AstItem;
   }
 
-  getTreeItem(item: AstItem): vscode.TreeItem {
+  getTreeItem(item: AstItem): TreeItem {
     const line = this.extractLineInfo(item?.location);
 
     const state = item.children.length
-      ? vscode.TreeItemCollapsibleState.Collapsed
-      : vscode.TreeItemCollapsibleState.None;
-    const treeItem = new vscode.TreeItem(item.label || '', state);
+      ? TreeItemCollapsibleState.Collapsed
+      : TreeItemCollapsibleState.None;
+    const treeItem = new TreeItem(item.label || '', state);
     treeItem.description = line ? `Line ${line}` : '';
     treeItem.id = String(item.id);
     treeItem.tooltip = `${treeItem.description} ${treeItem.label}`;
     treeItem.command = {
       command: 'codeQLAstViewer.gotoCode',
       title: 'Go To Code',
       tooltip: `Go To ${item.location}`,
-      arguments: [item.location, this.db]
+      arguments: [item]
     };
     return treeItem;
   }
@@ -83,33 +96,83 @@ class AstViewerDataProvider implements vscode.TreeDataProvider<AstItem | RootAst
 }
 
 export class AstViewer {
-  private treeView: vscode.TreeView<AstItem | RootAstItem>;
+  private treeView: TreeView<AstItem | RootAstItem>;
   private treeDataProvider: AstViewerDataProvider;
+  private currentFile: string | undefined;
 
-  constructor() {
+  constructor(ctx: ExtensionContext) {
     this.treeDataProvider = new AstViewerDataProvider();
-    this.treeView = vscode.window.createTreeView('codeQLAstViewer', {
+    this.treeView = window.createTreeView('codeQLAstViewer', {
       treeDataProvider: this.treeDataProvider,
       showCollapseAll: true
     });
 
-    vscode.commands.registerCommand('codeQLAstViewer.clear', () => {
+    commands.registerCommand('codeQLAstViewer.clear', () => {
       this.clear();
     });
+
+    ctx.subscriptions.push(window.onDidChangeTextEditorSelection(this.updateTreeSelection, this));
   }
 
   updateRoots(roots: RootAstItem[], db: DatabaseItem, fileName: string) {
     this.treeDataProvider.roots = roots;
     this.treeDataProvider.db = db;
     this.treeDataProvider.refresh();
-    this.treeView.message = `AST for ${fileName}`;
-    this.treeView.reveal(roots[0], { focus: true });
+    this.treeView.message = `AST for ${path.basename(fileName)}`;
+    this.treeView.reveal(roots[0], { focus: false });
+    this.currentFile = fileName;
+  }
+
+  private updateTreeSelection(e: TextEditorSelectionChangeEvent) {
+    function isInside(selectedRange: Range, astRange?: Range): boolean {
+      return !!astRange?.contains(selectedRange);
+    }
+
+    // Recursively iterate all children until we find the node with the smallest
+    // range that contains the selection.
+    // Some nodes do not have a location, but their children might, so must
+    // recurse though location-less AST nodes to see if children are correct.
+    function findBest(selectedRange: Range, items?: RootAstItem[]): RootAstItem | undefined {
+      if (!items || !items.length) {
+        return;
+      }
+      for (const item of items) {
+        let candidate: RootAstItem | undefined = undefined;
+        if (isInside(selectedRange, item.fileLocation?.range)) {
+          candidate = item;
+        }
+        // always iterate through children since the location of an AST node in code QL does not
+        // always cover the complete text of the node.
+        candidate = findBest(selectedRange, item.children) || candidate;
+        if (candidate) {
+          return candidate;
+        }
+      }
+      return;
+    }
+
+    if (
+      this.treeView.visible &&
+      e.textEditor.document.uri.fsPath === this.currentFile &&
+      e.selections.length === 1
+    ) {
+      const selection = e.selections[0];
+      const range = selection.anchor.isBefore(selection.active)
+        ? new Range(selection.anchor, selection.active)
+        : new Range(selection.active, selection.anchor);
+
+      const targetItem = findBest(range, this.treeDataProvider.roots);
+      if (targetItem) {
+        this.treeView.reveal(targetItem);
+      }
+    }
   }
 
   private clear() {
     this.treeDataProvider.roots = [];
     this.treeDataProvider.db = undefined;
     this.treeDataProvider.refresh();
     this.treeView.message = undefined;
+    this.currentFile = undefined;
   }
 }

extensions/ql-vscode/src/astViewer.ts

@@ -3,6 +3,7 @@ import { CodeQLCliServer } from '../cli';
 import { DecodedBqrsChunk, BqrsId, EntityValue } from '../bqrs-cli-types';
 import { DatabaseItem } from '../databases';
 import { AstItem, RootAstItem } from '../astViewer';
+import fileRangeFromURI from './fileRangeFromURI';
 
 /**
  * A class that wraps a tree of QL results from a query that
@@ -87,6 +88,7 @@ export default class AstBuilder {
             id,
             label: entity.label,
             location: entity.url,
+            fileLocation: fileRangeFromURI(entity.url, this.db),
             children: [] as AstItem[],
             order: Number.MAX_SAFE_INTEGER
           };
@@ -95,8 +97,8 @@ export default class AstBuilder {
           const parent = idToItem.get(childToParent.has(id) ? childToParent.get(id)! : -1);
 
           if (parent) {
-            const astItem = item as AstItem;
-            astItem.parent = parent;
+            const astItem = item as unknown as AstItem;
+            (astItem).parent = parent;
             parent.children.push(astItem);
           }
           const children = parentToChildren.has(id) ? parentToChildren.get(id)! : [];

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

@@ -4,8 +4,8 @@ import { UrlValue, LineColumnLocation } from '../bqrs-cli-types';
 import { DatabaseItem } from '../databases';
 
 
-export default function fileRangeFromURI(uri: UrlValue, db: DatabaseItem): vscode.Location | undefined {
-  if (typeof uri === 'string') {
+export default function fileRangeFromURI(uri: UrlValue | undefined, db: DatabaseItem): vscode.Location | undefined {
+  if (!uri || typeof uri === 'string') {
     return undefined;
   } else if ('startOffset' in uri) {
     return undefined;

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

@@ -1,5 +1,4 @@
 import * as vscode from 'vscode';
-import * as path from 'path';
 
 import { decodeSourceArchiveUri, zipArchiveScheme } from '../archive-filesystem-provider';
 import { CodeQLCliServer } from '../cli';
@@ -118,7 +117,7 @@ export class TemplatePrintAstProvider {
     return new AstBuilder(
       queryResults, this.cli,
       this.dbm.findDatabaseItem(vscode.Uri.parse(queryResults.database.databaseUri!))!,
-      path.basename(document.fileName)
+      document.fileName
     );
   }
 

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

@@ -549,7 +549,7 @@ async function activateWithInstalledDistribution(
     new TemplateQueryReferenceProvider(cliServer, qs, dbm)
   );
 
-  const astViewer = new AstViewer();
+  const astViewer = new AstViewer(ctx);
   ctx.subscriptions.push(commands.registerCommand('codeQL.viewAst', async () => {
     const ast = await new TemplatePrintAstProvider(cliServer, qs, dbm)
       .provideAst(window.activeTextEditor?.document);

extensions/ql-vscode/src/extension.ts

@@ -262,7 +262,7 @@ export async function getQlPackForDbscheme(cliServer: CodeQLCliServer, dbschemeP
     });
   for (const { packDir, packName } of packs) {
     if (packDir !== undefined) {
-      const qlpack = yaml.safeLoad(await fs.readFile(path.join(packDir, 'qlpack.yml'), 'utf8'));
+      const qlpack = yaml.safeLoad(await fs.readFile(path.join(packDir, 'qlpack.yml'), 'utf8')) as { dbscheme: string };
       if (qlpack.dbscheme !== undefined && path.basename(qlpack.dbscheme) === path.basename(dbschemePath)) {
         return packName;
       }