Introduce separate tree item types in the methods usage panel

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();
     });