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.

extensions/ql-vscode/package-lock.json

@@ -278,9 +278,9 @@
       }
     },
     "@types/js-yaml": {
-      "version": "3.12.2",
-      "resolved": "https://registry.npmjs.org/@types/js-yaml/-/js-yaml-3.12.2.tgz",
-      "integrity": "sha512-0CFu/g4mDSNkodVwWijdlr8jH7RoplRWNgovjFLEZeT+QEbbZXjBmCe3HwaWheAlCbHwomTwzZoSedeOycABug==",
+      "version": "3.12.5",
+      "resolved": "https://registry.npmjs.org/@types/js-yaml/-/js-yaml-3.12.5.tgz",
+      "integrity": "sha512-JCcp6J0GV66Y4ZMDAQCXot4xprYB+Zfd3meK9+INSJeVZwJmHAW30BBEEkPzXswMXuiyReUGOP3GxrADc9wPww==",
       "dev": true
     },
     "@types/json-schema": {

extensions/ql-vscode/package.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/src/astViewer.ts

@@ -1,8 +1,22 @@
-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';
 
@@ -10,6 +24,7 @@ export interface AstItem {
   id: BqrsId;
   label?: string;
   location?: UrlValue;
+  fileLocation?: Location;
   parent: AstItem | RootAstItem;
   children: AstItem[];
   order: number;
@@ -17,44 +32,42 @@ export interface AstItem {
 
 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}`;
@@ -62,7 +75,7 @@ class AstViewerDataProvider implements vscode.TreeDataProvider<AstItem | RootAst
       command: 'codeQLAstViewer.gotoCode',
       title: 'Go To Code',
       tooltip: `Go To ${item.location}`,
-      arguments: [item.location, this.db]
+      arguments: [item]
     };
     return treeItem;
   }
@@ -83,27 +96,76 @@ 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() {
@@ -111,5 +173,6 @@ export class AstViewer {
     this.treeDataProvider.db = undefined;
     this.treeDataProvider.refresh();
     this.treeView.message = undefined;
+    this.currentFile = undefined;
   }
 }

extensions/ql-vscode/src/contextual/astBuilder.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/fileRangeFromURI.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/templateProvider.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/extension.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/helpers.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;
       }

extensions/ql-vscode/src/vscode-tests/no-workspace/astViewer.test.ts

@@ -0,0 +1,150 @@
+import * as fs from 'fs-extra';
+import * as chai from 'chai';
+import * as chaiAsPromised from 'chai-as-promised';
+import * as sinon from 'sinon';
+import * as yaml from 'js-yaml';
+
+import { AstViewer, RootAstItem } from '../../astViewer';
+import { ExtensionContext, commands, Range } from 'vscode';
+import { DatabaseItem } from '../../databases';
+
+chai.use(chaiAsPromised);
+const expect = chai.expect;
+
+
+
+describe('AstViewer', () => {
+  let astRoots: RootAstItem[];
+  let viewer: AstViewer;
+  beforeEach(async () => {
+    // the ast is stored in yaml because there are back pointers
+    // making a json representation impossible.
+    // The complication here is that yaml files are not copied into the 'out' directory by tsc.
+    astRoots = await buildAst();
+
+    sinon.stub(commands, 'registerCommand');
+    sinon.stub(commands, 'executeCommand');
+  });
+
+  afterEach(() => {
+    sinon.restore();
+  });
+
+  it('should update the viewer roots', () => {
+    const item = {} as DatabaseItem;
+    viewer = new AstViewer({ subscriptions: [] as any[] } as ExtensionContext);
+    viewer.updateRoots(astRoots, item, 'def/abc');
+
+    expect((viewer as any).treeDataProvider.roots).to.eq(astRoots);
+    expect((viewer as any).treeDataProvider.db).to.eq(item);
+    expect((viewer as any).treeView.message).to.eq('AST for abc');
+  });
+
+  it('should update the tree selection based on a change in the editor selection', () => {
+    // Should select the namespace
+    doSelectionTest(astRoots[0], astRoots[0].fileLocation?.range);
+  });
+
+  it('should select an AssignExpr', () => {
+    // this one is interesting because it spans a couple of other nodes
+    const expr = findNodeById(300, astRoots);
+    expect(expr.label).to.eq('[AssignExpr] ... = ...');
+    doSelectionTest(expr, expr.fileLocation?.range);
+  });
+
+  it('should select nothing', () => {
+    doSelectionTest(undefined, new Range(2, 3, 4, 5));
+  });
+
+
+  function doSelectionTest(
+    expectedSelection: any,
+    selectionRange: Range | undefined,
+    fsPath = 'def/abc',
+  ) {
+    const item = {} as DatabaseItem;
+    viewer = new AstViewer({ subscriptions: [] as any[] } as ExtensionContext);
+    viewer.updateRoots(astRoots, item, fsPath);
+    const spy = sinon.spy();
+    (viewer as any).treeView.reveal = spy;
+    Object.defineProperty((viewer as any).treeView, 'visible', {
+      value: true
+    });
+
+    const mockEvent = createMockEvent(selectionRange, fsPath);
+    (viewer as any).updateTreeSelection(mockEvent);
+    if (expectedSelection) {
+      expect(spy).to.have.been.calledWith(expectedSelection);
+    } else {
+      expect(spy).not.to.have.been.called;
+    }
+  }
+
+  function createMockEvent(
+    selectionRange: Range | undefined,
+    fsPath: string,
+  ) {
+    return {
+      selections: [{
+        anchor: selectionRange?.start,
+        active: selectionRange?.end
+      }],
+      textEditor: {
+        document: {
+          uri: {
+            fsPath
+          }
+        }
+      }
+    };
+  }
+
+  function findNodeById(id: number, ast: any): any {
+    if (Array.isArray(ast)) {
+      for (const elt of ast) {
+        const candidate = findNodeById(id, elt);
+        if (candidate) {
+          return candidate;
+        }
+      }
+
+    } else if (typeof ast === 'object' && ast) {
+      if (ast.id === id) {
+        return ast;
+      } else {
+        for (const [name, prop] of Object.entries(ast)) {
+          if (name !== 'parent') {
+            const candidate = findNodeById(id, prop);
+            if (candidate) {
+              return candidate;
+            }
+          }
+        }
+      }
+    }
+  }
+
+  async function buildAst() {
+    const astRoots = yaml.safeLoad(await fs.readFile(`${__dirname}/../../../src/vscode-tests/no-workspace/data/astViewer.yml`, 'utf8')) as RootAstItem[];
+
+    // convert range properties into vscode.Range instances
+    function convertToRangeInstances(obj: any) {
+      if (Array.isArray(obj)) {
+        obj.forEach(elt => convertToRangeInstances(elt));
+      } else if (typeof obj === 'object' && obj) {
+        if ('range' in obj && '_start' in obj.range && '_end' in obj.range) {
+          obj.range = new Range(
+            obj.range._start._line,
+            obj.range._start._character,
+            obj.range._end._line,
+            obj.range._end._character,
+          );
+        } else {
+          Object.entries(obj).forEach(([name, prop]) => name !== 'parent' && convertToRangeInstances(prop));
+        }
+      }
+    }
+    convertToRangeInstances(astRoots);
+    return astRoots;
+  }
+});

extensions/ql-vscode/src/vscode-tests/no-workspace/contextual/astBuilder.test.ts

@@ -112,6 +112,7 @@ const expectedRoots = [
   {
     id: 0,
     label: '[TopLevelFunction] int disable_interrupts()',
+    fileLocation: undefined,
     location: {
       uri: 'file:/opt/src/arch/sandbox/lib/interrupts.c',
       startLine: 19,
@@ -125,6 +126,7 @@ const expectedRoots = [
   {
     id: 26363,
     label: '[TopLevelFunction] void enable_interrupts()',
+    fileLocation: undefined,
     location: {
       uri: 'file:/opt/src/arch/sandbox/lib/interrupts.c',
       startLine: 15,
@@ -138,6 +140,7 @@ const expectedRoots = [
   {
     id: 26364,
     label: '[TopLevelFunction] int interrupt_init()',
+    fileLocation: undefined,
     location: {
       uri: 'file:/opt/src/arch/sandbox/lib/interrupts.c',
       startLine: 10,