Sort methods in usages panel according to model editor sort order

This sorts the methods in the methods usages panel the same as in the
model editor. Since this is dependent on the mode, we need to keep track
of the mode in the modeling store, so this also adds a mode field to the
db state.

Author: koesie10

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

@@ -17,15 +17,21 @@ import { INITIAL_HIDE_MODELED_METHODS_VALUE } from "../shared/hide-modeled-metho
 import { getModelingStatus } from "../shared/modeling-status";
 import { assertNever } from "../../common/helpers-pure";
 import { ModeledMethod } from "../modeled-method";
+import { groupMethods, sortGroupNames, sortMethods } from "../shared/sorting";
+import { INITIAL_MODE, Mode } from "../shared/mode";
 
 export class MethodsUsageDataProvider
   extends DisposableObject
   implements TreeDataProvider<MethodsUsageTreeViewItem>
 {
   private methods: 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: Method[] = [];
   private databaseItem: DatabaseItem | undefined = undefined;
   private sourceLocationPrefix: string | undefined = undefined;
   private hideModeledMethods: boolean = INITIAL_HIDE_MODELED_METHODS_VALUE;
+  private mode: Mode = INITIAL_MODE;
   private modeledMethods: Record<string, ModeledMethod[]> = {};
   private modifiedMethodSignatures: Set<string> = new Set();
 
@@ -52,21 +58,25 @@ export class MethodsUsageDataProvider
     methods: Method[],
     databaseItem: DatabaseItem,
     hideModeledMethods: boolean,
+    mode: Mode,
     modeledMethods: Record<string, ModeledMethod[]>,
     modifiedMethodSignatures: Set<string>,
   ): Promise<void> {
     if (
       this.methods !== methods ||
       this.databaseItem !== databaseItem ||
       this.hideModeledMethods !== hideModeledMethods ||
+      this.mode !== mode ||
       this.modeledMethods !== modeledMethods ||
       this.modifiedMethodSignatures !== modifiedMethodSignatures
     ) {
       this.methods = methods;
+      this.sortedMethods = sortMethodsInGroups(methods, mode);
       this.databaseItem = databaseItem;
       this.sourceLocationPrefix =
         await this.databaseItem.getSourceLocationPrefix(this.cliServer);
       this.hideModeledMethods = hideModeledMethods;
+      this.mode = mode;
       this.modeledMethods = modeledMethods;
       this.modifiedMethodSignatures = modifiedMethodSignatures;
 
@@ -133,9 +143,9 @@ export class MethodsUsageDataProvider
   getChildren(item?: MethodsUsageTreeViewItem): MethodsUsageTreeViewItem[] {
     if (item === undefined) {
       if (this.hideModeledMethods) {
-        return this.methods.filter((api) => !api.supported);
+        return this.sortedMethods.filter((api) => !api.supported);
       } else {
-        return this.methods;
+        return this.sortedMethods;
       }
     } else if (isExternalApiUsage(item)) {
       return item.usages;
@@ -183,3 +193,15 @@ function usagesAreEqual(u1: Usage, u2: Usage): boolean {
     u1.url.endColumn === u2.url.endColumn
   );
 }
+
+function sortMethodsInGroups(methods: Method[], mode: Mode): Method[] {
+  const grouped = groupMethods(methods, mode);
+
+  const sortedGroupNames = sortGroupNames(grouped);
+
+  return sortedGroupNames.flatMap((groupName) => {
+    const group = grouped[groupName];
+
+    return sortMethods(group);
+  });
+}

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

@@ -9,6 +9,7 @@ import { DatabaseItem } from "../../databases/local-databases";
 import { CodeQLCliServer } from "../../codeql-cli/cli";
 import { ModelingStore } from "../modeling-store";
 import { ModeledMethod } from "../modeled-method";
+import { Mode } from "../shared/mode";
 
 export class MethodsUsagePanel extends DisposableObject {
   private readonly dataProvider: MethodsUsageDataProvider;
@@ -34,13 +35,15 @@ export class MethodsUsagePanel extends DisposableObject {
     methods: Method[],
     databaseItem: DatabaseItem,
     hideModeledMethods: boolean,
+    mode: Mode,
     modeledMethods: Record<string, ModeledMethod[]>,
     modifiedMethodSignatures: Set<string>,
   ): Promise<void> {
     await this.dataProvider.setState(
       methods,
       databaseItem,
       hideModeledMethods,
+      mode,
       modeledMethods,
       modifiedMethodSignatures,
     );
@@ -83,6 +86,14 @@ export class MethodsUsagePanel extends DisposableObject {
       }),
     );
 
+    this.push(
+      this.modelingStore.onModeChanged(async (event) => {
+        if (event.isActiveDb) {
+          await this.handleStateChangeEvent();
+        }
+      }),
+    );
+
     this.push(
       this.modelingStore.onModifiedMethodsChanged(async (event) => {
         if (event.isActiveDb) {
@@ -99,6 +110,7 @@ export class MethodsUsagePanel extends DisposableObject {
         activeState.methods,
         activeState.databaseItem,
         activeState.hideModeledMethods,
+        activeState.mode,
         activeState.modeledMethods,
         activeState.modifiedMethodSignatures,
       );

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

@@ -14,7 +14,6 @@ import { dir } from "tmp-promise";
 import { isQueryLanguage } from "../common/query-language";
 import { DisposableObject } from "../common/disposable-object";
 import { MethodsUsagePanel } from "./methods-usage/methods-usage-panel";
-import { Mode } from "./shared/mode";
 import { Method, Usage } from "./method";
 import { setUpPack } from "./model-editor-queries";
 import { MethodModelingPanel } from "./method-modeling/method-modeling-panel";
@@ -204,7 +203,6 @@ export class ModelEditorModule extends DisposableObject {
             queryDir,
             db,
             modelFile,
-            Mode.Application,
           );
 
           this.modelingStore.onDbClosed(async (dbUri) => {

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

@@ -35,7 +35,7 @@ import { Method } from "./method";
 import { ModeledMethod } from "./modeled-method";
 import { ExtensionPack } from "./shared/extension-pack";
 import { ModelConfigListener } from "../config";
-import { Mode } from "./shared/mode";
+import { INITIAL_MODE, Mode } from "./shared/mode";
 import { loadModeledMethods, saveModeledMethods } from "./modeled-method-fs";
 import { pickExtensionPack } from "./extension-pack-picker";
 import { getLanguageDisplayName } from "../common/query-language";
@@ -66,11 +66,11 @@ export class ModelEditorView extends AbstractWebview<
     private readonly queryDir: string,
     private readonly databaseItem: DatabaseItem,
     private readonly extensionPack: ExtensionPack,
-    private mode: Mode,
+    initialMode: Mode = INITIAL_MODE,
   ) {
     super(app);
 
-    this.modelingStore.initializeStateForDb(databaseItem);
+    this.modelingStore.initializeStateForDb(databaseItem, initialMode);
     this.registerToModelingStoreEvents();
     this.registerToModelConfigEvents();
 
@@ -214,6 +214,7 @@ export class ModelEditorView extends AbstractWebview<
             this.databaseItem,
             msg.methodSignatures,
           );
+          const mode = this.modelingStore.getMode(this.databaseItem);
 
           await withProgress(
             async (progress) => {
@@ -227,7 +228,7 @@ export class ModelEditorView extends AbstractWebview<
                 this.databaseItem.language,
                 methods,
                 modeledMethods,
-                this.mode,
+                mode,
                 this.cliServer,
                 this.app.logger,
               );
@@ -288,7 +289,7 @@ export class ModelEditorView extends AbstractWebview<
         );
         break;
       case "switchMode":
-        this.mode = msg.mode;
+        this.modelingStore.setMode(this.databaseItem, msg.mode);
         this.modelingStore.setMethods(this.databaseItem, []);
         await Promise.all([
           this.postMessage({
@@ -372,7 +373,7 @@ export class ModelEditorView extends AbstractWebview<
         showFlowGeneration: this.modelConfig.flowGeneration,
         showLlmButton,
         showMultipleModels: this.modelConfig.showMultipleModels,
-        mode: this.mode,
+        mode: this.modelingStore.getMode(this.databaseItem),
       },
     });
   }
@@ -398,9 +399,11 @@ export class ModelEditorView extends AbstractWebview<
   }
 
   protected async loadMethods(progress: ProgressCallback): Promise<void> {
+    const mode = this.modelingStore.getMode(this.databaseItem);
+
     try {
       const cancellationTokenSource = new CancellationTokenSource();
-      const queryResult = await runExternalApiQueries(this.mode, {
+      const queryResult = await runExternalApiQueries(mode, {
         cliServer: this.cliServer,
         queryRunner: this.queryRunner,
         databaseItem: this.databaseItem,
@@ -434,11 +437,13 @@ export class ModelEditorView extends AbstractWebview<
       async (progress) => {
         const tokenSource = new CancellationTokenSource();
 
+        const mode = this.modelingStore.getMode(this.databaseItem);
+
         let addedDatabase: DatabaseItem | undefined;
 
         // In application mode, we need the database of a specific library to generate
         // the modeled methods. In framework mode, we'll use the current database.
-        if (this.mode === Mode.Application) {
+        if (mode === Mode.Application) {
           addedDatabase = await this.promptChooseNewOrExistingDatabase(
             progress,
           );
@@ -503,11 +508,12 @@ export class ModelEditorView extends AbstractWebview<
       this.databaseItem,
       methodSignatures,
     );
+    const mode = this.modelingStore.getMode(this.databaseItem);
     await this.autoModeler.startModeling(
       packageName,
       methods,
       modeledMethods,
-      this.mode,
+      mode,
     );
   }
 

extensions/ql-vscode/src/model-editor/modeling-store.ts

@@ -5,11 +5,13 @@ import { DatabaseItem } from "../databases/local-databases";
 import { Method, Usage } from "./method";
 import { ModeledMethod } from "./modeled-method";
 import { INITIAL_HIDE_MODELED_METHODS_VALUE } from "./shared/hide-modeled-methods";
+import { INITIAL_MODE, Mode } from "./shared/mode";
 
 interface DbModelingState {
   databaseItem: DatabaseItem;
   methods: Method[];
   hideModeledMethods: boolean;
+  mode: Mode;
   modeledMethods: Record<string, ModeledMethod[]>;
   modifiedMethodSignatures: Set<string>;
   selectedMethod: Method | undefined;
@@ -27,6 +29,11 @@ interface HideModeledMethodsChangedEvent {
   isActiveDb: boolean;
 }
 
+interface ModeChangedEvent {
+  mode: Mode;
+  isActiveDb: boolean;
+}
+
 interface ModeledMethodsChangedEvent {
   modeledMethods: Record<string, ModeledMethod[]>;
   dbUri: string;
@@ -53,6 +60,7 @@ export class ModelingStore extends DisposableObject {
   public readonly onDbClosed: AppEvent<string>;
   public readonly onMethodsChanged: AppEvent<MethodsChangedEvent>;
   public readonly onHideModeledMethodsChanged: AppEvent<HideModeledMethodsChangedEvent>;
+  public readonly onModeChanged: AppEvent<ModeChangedEvent>;
   public readonly onModeledMethodsChanged: AppEvent<ModeledMethodsChangedEvent>;
   public readonly onModifiedMethodsChanged: AppEvent<ModifiedMethodsChangedEvent>;
   public readonly onSelectedMethodChanged: AppEvent<SelectedMethodChangedEvent>;
@@ -65,6 +73,7 @@ export class ModelingStore extends DisposableObject {
   private readonly onDbClosedEventEmitter: AppEventEmitter<string>;
   private readonly onMethodsChangedEventEmitter: AppEventEmitter<MethodsChangedEvent>;
   private readonly onHideModeledMethodsChangedEventEmitter: AppEventEmitter<HideModeledMethodsChangedEvent>;
+  private readonly onModeChangedEventEmitter: AppEventEmitter<ModeChangedEvent>;
   private readonly onModeledMethodsChangedEventEmitter: AppEventEmitter<ModeledMethodsChangedEvent>;
   private readonly onModifiedMethodsChangedEventEmitter: AppEventEmitter<ModifiedMethodsChangedEvent>;
   private readonly onSelectedMethodChangedEventEmitter: AppEventEmitter<SelectedMethodChangedEvent>;
@@ -98,6 +107,11 @@ export class ModelingStore extends DisposableObject {
     this.onHideModeledMethodsChanged =
       this.onHideModeledMethodsChangedEventEmitter.event;
 
+    this.onModeChangedEventEmitter = this.push(
+      app.createEventEmitter<ModeChangedEvent>(),
+    );
+    this.onModeChanged = this.onModeChangedEventEmitter.event;
+
     this.onModeledMethodsChangedEventEmitter = this.push(
       app.createEventEmitter<ModeledMethodsChangedEvent>(),
     );
@@ -117,12 +131,16 @@ export class ModelingStore extends DisposableObject {
       this.onSelectedMethodChangedEventEmitter.event;
   }
 
-  public initializeStateForDb(databaseItem: DatabaseItem) {
+  public initializeStateForDb(
+    databaseItem: DatabaseItem,
+    mode: Mode = INITIAL_MODE,
+  ) {
     const dbUri = databaseItem.databaseUri.toString();
     this.state.set(dbUri, {
       databaseItem,
       methods: [],
       hideModeledMethods: INITIAL_HIDE_MODELED_METHODS_VALUE,
+      mode,
       modeledMethods: {},
       modifiedMethodSignatures: new Set(),
       selectedMethod: undefined,
@@ -214,6 +232,22 @@ export class ModelingStore extends DisposableObject {
     });
   }
 
+  public setMode(dbItem: DatabaseItem, mode: Mode) {
+    const dbState = this.getState(dbItem);
+    const dbUri = dbItem.databaseUri.toString();
+
+    dbState.mode = mode;
+
+    this.onModeChangedEventEmitter.fire({
+      mode,
+      isActiveDb: dbUri === this.activeDb,
+    });
+  }
+
+  public getMode(dbItem: DatabaseItem) {
+    return this.getState(dbItem).mode;
+  }
+
   /**
    * Returns the modeled methods for the given database item and method signatures.
    * If the `methodSignatures` argument is not provided or is undefined, returns all modeled methods.

extensions/ql-vscode/src/model-editor/shared/mode.ts

@@ -2,3 +2,5 @@ export enum Mode {
   Application = "application",
   Framework = "framework",
 }
+
+export const INITIAL_MODE = Mode.Application;

extensions/ql-vscode/test/__mocks__/model-editor/modelingStoreMock.ts

@@ -8,6 +8,7 @@ export function createMockModelingStore({
   onDbClosed = jest.fn(),
   onMethodsChanged = jest.fn(),
   onHideModeledMethodsChanged = jest.fn(),
+  onModeChanged = jest.fn(),
   onModeledMethodsChanged = jest.fn(),
   onModifiedMethodsChanged = jest.fn(),
 }: {
@@ -17,6 +18,7 @@ export function createMockModelingStore({
   onDbClosed?: ModelingStore["onDbClosed"];
   onMethodsChanged?: ModelingStore["onMethodsChanged"];
   onHideModeledMethodsChanged?: ModelingStore["onHideModeledMethodsChanged"];
+  onModeChanged?: ModelingStore["onModeChanged"];
   onModeledMethodsChanged?: ModelingStore["onModeledMethodsChanged"];
   onModifiedMethodsChanged?: ModelingStore["onModifiedMethodsChanged"];
 } = {}): ModelingStore {
@@ -27,6 +29,7 @@ export function createMockModelingStore({
     onDbClosed,
     onMethodsChanged,
     onHideModeledMethodsChanged,
+    onModeChanged,
     onModeledMethodsChanged,
     onModifiedMethodsChanged,
   });