Sort methods in usages panel according to model editor sort order

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";
@@ -224,7 +223,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";
@@ -63,11 +63,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();
 
@@ -211,6 +211,7 @@ export class ModelEditorView extends AbstractWebview<
             this.databaseItem,
             msg.methodSignatures,
           );
+          const mode = this.modelingStore.getMode(this.databaseItem);
 
           await withProgress(
             async (progress) => {
@@ -224,7 +225,7 @@ export class ModelEditorView extends AbstractWebview<
                 this.databaseItem.language,
                 methods,
                 modeledMethods,
-                this.mode,
+                mode,
                 this.cliServer,
                 this.app.logger,
               );
@@ -285,7 +286,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({
@@ -376,7 +377,7 @@ export class ModelEditorView extends AbstractWebview<
         showFlowGeneration: this.modelConfig.flowGeneration,
         showLlmButton,
         showMultipleModels: this.modelConfig.showMultipleModels,
-        mode: this.mode,
+        mode: this.modelingStore.getMode(this.databaseItem),
         sourceArchiveAvailable,
       },
     });
@@ -403,9 +404,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,
@@ -439,11 +442,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,
           );
@@ -508,11 +513,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,
   });