Introduce separate tree item types in the methods usage panel

This creates new tree item types for methods and usages such that these
can contain references to their parent and children. This allows us to
easily find the parent of a usage and to find the children of a method.
This removes an expensive `find` call in `getParent`.

Author: koesie10

extensions/ql-vscode/src/model-editor/methods-usage/methods-usage-data-provider.ts

@@ -27,7 +27,7 @@ export class MethodsUsageDataProvider
   private methods: readonly Method[] = [];
   // sortedMethods is a separate field so we can check if the methods have changed
   // by reference, which is faster than checking if the methods have changed by value.
-  private sortedMethods: readonly Method[] = [];
+  private sortedTreeItems: readonly MethodTreeViewItem[] = [];
   private databaseItem: DatabaseItem | undefined = undefined;
   private sourceLocationPrefix: string | undefined = undefined;
   private hideModeledMethods: boolean = INITIAL_HIDE_MODELED_METHODS_VALUE;
@@ -72,7 +72,9 @@ export class MethodsUsageDataProvider
       this.modifiedMethodSignatures !== modifiedMethodSignatures
     ) {
       this.methods = methods;
-      this.sortedMethods = sortMethodsInGroups(methods, mode);
+      this.sortedTreeItems = createTreeItems(
+        sortMethodsInGroups(methods, mode),
+      );
       this.databaseItem = databaseItem;
       this.sourceLocationPrefix =
         await this.databaseItem.getSourceLocationPrefix(this.cliServer);
@@ -86,15 +88,15 @@ export class MethodsUsageDataProvider
   }
 
   getTreeItem(item: MethodsUsageTreeViewItem): TreeItem {
-    if (isExternalApiUsage(item)) {
+    if (isMethodTreeViewItem(item)) {
       return {
         label: `${item.packageName}.${item.typeName}.${item.methodName}${item.methodParameters}`,
         collapsibleState: TreeItemCollapsibleState.Collapsed,
         iconPath: this.getModelingStatusIcon(item),
       };
     } else {
       const method = this.getParent(item);
-      if (!method || !isExternalApiUsage(method)) {
+      if (!method || !isMethodTreeViewItem(method)) {
         throw new Error("Parent not found for tree item");
       }
       return {
@@ -144,12 +146,12 @@ export class MethodsUsageDataProvider
   getChildren(item?: MethodsUsageTreeViewItem): MethodsUsageTreeViewItem[] {
     if (item === undefined) {
       if (this.hideModeledMethods) {
-        return this.sortedMethods.filter((api) => !api.supported);
+        return this.sortedTreeItems.filter((api) => !api.supported);
       } else {
-        return [...this.sortedMethods];
+        return [...this.sortedTreeItems];
       }
-    } else if (isExternalApiUsage(item)) {
-      return [...item.usages];
+    } else if (isMethodTreeViewItem(item)) {
+      return item.children;
     } else {
       return [];
     }
@@ -158,29 +160,42 @@ export class MethodsUsageDataProvider
   getParent(
     item: MethodsUsageTreeViewItem,
   ): MethodsUsageTreeViewItem | undefined {
-    if (isExternalApiUsage(item)) {
+    if (isMethodTreeViewItem(item)) {
       return undefined;
     } else {
-      return this.methods.find((e) => e.usages.includes(item));
+      return item.parent;
     }
   }
 
-  public resolveCanonicalUsage(usage: Usage): Usage | undefined {
-    for (const method of this.methods) {
-      for (const u of method.usages) {
-        if (usagesAreEqual(u, usage)) {
-          return u;
-        }
-      }
+  public resolveUsageTreeViewItem(
+    methodSignature: string,
+    usage: Usage,
+  ): UsageTreeViewItem | undefined {
+    const method = this.sortedTreeItems.find(
+      (m) => m.signature === methodSignature,
+    );
+    if (!method) {
+      return undefined;
     }
-    return undefined;
+
+    return method.children.find((u) => usagesAreEqual(u, usage));
   }
 }
 
-export type MethodsUsageTreeViewItem = Method | Usage;
+type MethodTreeViewItem = Method & {
+  children: UsageTreeViewItem[];
+};
+
+type UsageTreeViewItem = Usage & {
+  parent: MethodTreeViewItem;
+};
+
+export type MethodsUsageTreeViewItem = MethodTreeViewItem | UsageTreeViewItem;
 
-function isExternalApiUsage(item: MethodsUsageTreeViewItem): item is Method {
-  return (item as any).usages !== undefined;
+function isMethodTreeViewItem(
+  item: MethodsUsageTreeViewItem,
+): item is MethodTreeViewItem {
+  return "children" in item && "usages" in item;
 }
 
 function usagesAreEqual(u1: Usage, u2: Usage): boolean {
@@ -206,3 +221,20 @@ function sortMethodsInGroups(methods: readonly Method[], mode: Mode): Method[] {
     return sortMethods(group);
   });
 }
+
+function createTreeItems(methods: readonly Method[]): MethodTreeViewItem[] {
+  return methods.map((method) => {
+    const newMethod: MethodTreeViewItem = {
+      ...method,
+      children: [],
+    };
+
+    newMethod.children = method.usages.map((usage) => ({
+      ...usage,
+      // This needs to be a reference to the parent method, not a copy of it.
+      parent: newMethod,
+    }));
+
+    return newMethod;
+  });
+}

extensions/ql-vscode/src/model-editor/methods-usage/methods-usage-panel.ts

@@ -56,10 +56,16 @@ export class MethodsUsagePanel extends DisposableObject {
     };
   }
 
-  public async revealItem(usage: Usage): Promise<void> {
-    const canonicalUsage = this.dataProvider.resolveCanonicalUsage(usage);
-    if (canonicalUsage !== undefined) {
-      await this.treeView.reveal(canonicalUsage);
+  public async revealItem(
+    methodSignature: string,
+    usage: Usage,
+  ): Promise<void> {
+    const usageTreeViewItem = this.dataProvider.resolveUsageTreeViewItem(
+      methodSignature,
+      usage,
+    );
+    if (usageTreeViewItem !== undefined) {
+      await this.treeView.reveal(usageTreeViewItem);
     }
   }
 

extensions/ql-vscode/src/model-editor/model-editor-module.ts

@@ -102,7 +102,7 @@ export class ModelEditorModule extends DisposableObject {
     method: Method,
     usage: Usage,
   ): Promise<void> {
-    await this.methodsUsagePanel.revealItem(usage);
+    await this.methodsUsagePanel.revealItem(method.signature, usage);
     await this.methodModelingPanel.setMethod(databaseItem, method);
     await showResolvableLocation(usage.url, databaseItem, this.app.logger);
   }

extensions/ql-vscode/test/vscode-tests/no-workspace/model-editor/methods-usage/methods-usage-data-provider.test.ts

@@ -3,7 +3,10 @@ import {
   CallClassification,
   Method,
 } from "../../../../../src/model-editor/method";
-import { MethodsUsageDataProvider } from "../../../../../src/model-editor/methods-usage/methods-usage-data-provider";
+import {
+  MethodsUsageDataProvider,
+  MethodsUsageTreeViewItem,
+} from "../../../../../src/model-editor/methods-usage/methods-usage-data-provider";
 import { DatabaseItem } from "../../../../../src/databases/local-databases";
 import {
   createMethod,
@@ -242,13 +245,23 @@ describe("MethodsUsageDataProvider", () => {
 
     const usage = createUsage({});
 
+    const methodTreeItem: MethodsUsageTreeViewItem = {
+      ...supportedMethod,
+      children: [],
+    };
+
+    const usageTreeItem: MethodsUsageTreeViewItem = {
+      ...usage,
+      parent: methodTreeItem,
+    };
+    methodTreeItem.children = [usageTreeItem];
+
     it("should return [] if item is a usage", async () => {
-      expect(dataProvider.getChildren(usage)).toEqual([]);
+      expect(dataProvider.getChildren(usageTreeItem)).toEqual([]);
     });
 
-    it("should return usages if item is external api usage", async () => {
-      const method = createMethod({ usages: [usage] });
-      expect(dataProvider.getChildren(method)).toEqual([usage]);
+    it("should return usages if item is method", async () => {
+      expect(dataProvider.getChildren(methodTreeItem)).toEqual([usageTreeItem]);
     });
 
     it("should show all methods if hideModeledMethods is false and looking at the root", async () => {

extensions/ql-vscode/test/vscode-tests/no-workspace/model-editor/methods-usage/methods-usage-panel.test.ts

@@ -68,11 +68,10 @@ describe("MethodsUsagePanel", () => {
     });
 
     it("should reveal the correct item in the tree view", async () => {
-      const methods = [
-        createMethod({
-          usages: [usage],
-        }),
-      ];
+      const method = createMethod({
+        usages: [usage],
+      });
+      const methods = [method];
 
       const panel = new MethodsUsagePanel(modelingStore, mockCliServer);
       await panel.setState(
@@ -84,13 +83,16 @@ describe("MethodsUsagePanel", () => {
         modifiedMethodSignatures,
       );
 
-      await panel.revealItem(usage);
+      await panel.revealItem(method.signature, usage);
 
-      expect(mockTreeView.reveal).toHaveBeenCalledWith(usage);
+      expect(mockTreeView.reveal).toHaveBeenCalledWith(
+        expect.objectContaining(usage),
+      );
     });
 
     it("should do nothing if usage cannot be found", async () => {
-      const methods = [createMethod({})];
+      const method = createMethod({});
+      const methods = [method];
       const panel = new MethodsUsagePanel(modelingStore, mockCliServer);
       await panel.setState(
         methods,
@@ -101,7 +103,7 @@ describe("MethodsUsagePanel", () => {
         modifiedMethodSignatures,
       );
 
-      await panel.revealItem(usage);
+      await panel.revealItem(method.signature, usage);
 
       expect(mockTreeView.reveal).not.toHaveBeenCalled();
     });