Move all sorting to the modeling store, and then all views preserve the sorting

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

@@ -16,7 +16,7 @@ import { INITIAL_HIDE_MODELED_METHODS_VALUE } from "../shared/hide-modeled-metho
 import { getModelingStatus } from "../shared/modeling-status";
 import { assertNever } from "../../common/helpers-pure";
 import type { ModeledMethod } from "../modeled-method";
-import { groupMethods, sortGroupNames, sortMethods } from "../shared/sorting";
+import { groupMethods, sortGroupNames } from "../shared/sorting";
 import type { Mode } from "../shared/mode";
 import { INITIAL_MODE } from "../shared/mode";
 import type { UrlValueResolvable } from "../../common/raw-result-types";
@@ -63,7 +63,6 @@ export class MethodsUsageDataProvider
     mode: Mode,
     modeledMethods: Readonly<Record<string, readonly ModeledMethod[]>>,
     modifiedMethodSignatures: ReadonlySet<string>,
-    processedByAutoModelMethods: ReadonlySet<string>,
   ): Promise<void> {
     if (
       this.methods !== methods ||
@@ -74,15 +73,7 @@ export class MethodsUsageDataProvider
       this.modifiedMethodSignatures !== modifiedMethodSignatures
     ) {
       this.methods = methods;
-      this.sortedTreeItems = createTreeItems(
-        sortMethodsInGroups(
-          methods,
-          modeledMethods,
-          mode,
-          modifiedMethodSignatures,
-          processedByAutoModelMethods,
-        ),
-      );
+      this.sortedTreeItems = createTreeItems(createGroups(methods, mode));
       this.databaseItem = databaseItem;
       this.sourceLocationPrefix =
         await this.databaseItem.getSourceLocationPrefix(this.cliServer);
@@ -253,27 +244,9 @@ function urlValueResolvablesAreEqual(
   return false;
 }
 
-function sortMethodsInGroups(
-  methods: readonly Method[],
-  modeledMethods: Readonly<Record<string, readonly ModeledMethod[]>>,
-  mode: Mode,
-  modifiedMethodSignatures: ReadonlySet<string>,
-  processedByAutoModelMethods: ReadonlySet<string>,
-): Method[] {
+function createGroups(methods: readonly Method[], mode: Mode): Method[] {
   const grouped = groupMethods(methods, mode);
-
-  const sortedGroupNames = sortGroupNames(grouped);
-
-  return sortedGroupNames.flatMap((groupName) => {
-    const group = grouped[groupName];
-
-    return sortMethods(
-      group,
-      modeledMethods,
-      modifiedMethodSignatures,
-      processedByAutoModelMethods,
-    );
-  });
+  return sortGroupNames(grouped).flatMap((groupName) => grouped[groupName]);
 }
 
 function createTreeItems(methods: readonly Method[]): MethodTreeViewItem[] {

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

@@ -39,7 +39,6 @@ export class MethodsUsagePanel extends DisposableObject {
     mode: Mode,
     modeledMethods: Readonly<Record<string, readonly ModeledMethod[]>>,
     modifiedMethodSignatures: ReadonlySet<string>,
-    processedByAutoModelMethods: ReadonlySet<string>,
   ): Promise<void> {
     await this.dataProvider.setState(
       methods,
@@ -48,7 +47,6 @@ export class MethodsUsagePanel extends DisposableObject {
       mode,
       modeledMethods,
       modifiedMethodSignatures,
-      processedByAutoModelMethods,
     );
     const numOfApis = hideModeledMethods
       ? methods.filter((api) => !api.supported).length
@@ -122,7 +120,6 @@ export class MethodsUsagePanel extends DisposableObject {
         activeState.mode,
         activeState.modeledMethods,
         activeState.modifiedMethodSignatures,
-        activeState.processedByAutoModelMethods,
       );
     }
   }

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

@@ -273,6 +273,8 @@ export class ModelEditorView extends AbstractWebview<
             Object.keys(modeledMethods),
           );
 
+          this.modelingStore.updateMethodSorting(this.databaseItem);
+
           void telemetryListener?.sendUIInteraction(
             "model-editor-save-modeled-methods",
           );

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

@@ -6,6 +6,7 @@ import type { ModeledMethod } from "./modeled-method";
 import type { ModelingEvents } from "./modeling-events";
 import { INITIAL_HIDE_MODELED_METHODS_VALUE } from "./shared/hide-modeled-methods";
 import type { Mode } from "./shared/mode";
+import { sortMethods } from "./shared/sorting";
 
 interface InternalDbModelingState {
   databaseItem: DatabaseItem;
@@ -155,17 +156,25 @@ export class ModelingStore extends DisposableObject {
   }
 
   public setMethods(dbItem: DatabaseItem, methods: Method[]) {
-    const dbState = this.getState(dbItem);
-    const dbUri = dbItem.databaseUri.toString();
-
-    dbState.methods = [...methods];
+    this.changeMethods(dbItem, (state) => {
+      state.methods = sortMethods(
+        methods,
+        state.modeledMethods,
+        state.modifiedMethodSignatures,
+        state.processedByAutoModelMethods,
+      );
+    });
+  }
 
-    this.modelingEvents.fireMethodsChangedEvent(
-      methods,
-      dbUri,
-      dbItem,
-      dbUri === this.activeDb,
-    );
+  public updateMethodSorting(dbItem: DatabaseItem) {
+    this.changeMethods(dbItem, (state) => {
+      state.methods = sortMethods(
+        state.methods,
+        state.modeledMethods,
+        state.modifiedMethodSignatures,
+        state.processedByAutoModelMethods,
+      );
+    });
   }
 
   public setHideModeledMethods(
@@ -374,6 +383,7 @@ export class ModelingStore extends DisposableObject {
         ...processedByAutoModelMethods,
       ]);
     });
+    this.updateMethodSorting(dbItem);
   }
 
   public updateModelEvaluationRun(
@@ -421,6 +431,22 @@ export class ModelingStore extends DisposableObject {
     return this.state.get(databaseItem.databaseUri.toString())!;
   }
 
+  private changeMethods(
+    dbItem: DatabaseItem,
+    updateState: (state: InternalDbModelingState) => void,
+  ) {
+    const state = this.getState(dbItem);
+
+    updateState(state);
+
+    this.modelingEvents.fireMethodsChangedEvent(
+      state.methods,
+      dbItem.databaseUri.toString(),
+      dbItem,
+      dbItem.databaseUri.toString() === this.activeDb,
+    );
+  }
+
   private changeModifiedMethods(
     dbItem: DatabaseItem,
     updateState: (state: InternalDbModelingState) => void,

extensions/ql-vscode/src/model-editor/shared/auto-model-candidates.ts

@@ -1,7 +1,7 @@
 import type { Method, MethodSignature } from "../method";
 import type { ModeledMethod } from "../modeled-method";
 import type { Mode } from "./mode";
-import { groupMethods, sortGroupNames, sortMethods } from "./sorting";
+import { groupMethods, sortGroupNames } from "./sorting";
 
 /**
  * Return the candidates that the model should be run on. This includes limiting the number of
@@ -44,10 +44,5 @@ export function getCandidates(
 
   // Sort the same way as the UI so we send the first ones listed in the UI first
   const grouped = groupMethods(candidateMethods, mode);
-  const sortedGroupNames = sortGroupNames(grouped);
-  return sortedGroupNames.flatMap((name) =>
-    // We can safely pass empty sets for `modifiedSignatures` and `processedByAutoModelMethods`
-    // because we've filtered out all methods that are already modeled or have already been processed by auto-model.
-    sortMethods(grouped[name], modeledMethodsBySignature, new Set(), new Set()),
-  );
+  return sortGroupNames(grouped).flatMap((name) => grouped[name]);
 }

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

@@ -4,6 +4,10 @@ import type { ModeledMethod } from "../modeled-method";
 import { Mode } from "./mode";
 import { calculateModeledPercentage } from "./modeled-percentage";
 
+/**
+ * Groups methods by library or package name.
+ * Does not change the order of methods within a group.
+ */
 export function groupMethods(
   methods: readonly Method[],
   mode: Mode,

extensions/ql-vscode/src/view/model-editor/ModeledMethodDataGrid.tsx

@@ -3,7 +3,6 @@ import type { Method } from "../../model-editor/method";
 import { canMethodBeModeled } from "../../model-editor/method";
 import type { ModeledMethod } from "../../model-editor/modeled-method";
 import { useMemo } from "react";
-import { sortMethods } from "../../model-editor/shared/sorting";
 import { HiddenMethodsRow } from "./HiddenMethodsRow";
 import type { ModelEditorViewState } from "../../model-editor/shared/view-state";
 import { ScreenReaderOnly } from "../common/ScreenReaderOnly";
@@ -48,12 +47,7 @@ export const ModeledMethodDataGrid = ({
   ] = useMemo(() => {
     const methodsWithModelability = [];
     let numHiddenMethods = 0;
-    for (const method of sortMethods(
-      methods,
-      modeledMethodsMap,
-      modifiedSignatures,
-      processedByAutoModelMethods,
-    )) {
+    for (const method of methods) {
       const modeledMethods = modeledMethodsMap[method.signature] ?? [];
       const methodIsUnsaved = modifiedSignatures.has(method.signature);
       const methodCanBeModeled = canMethodBeModeled(
@@ -69,13 +63,7 @@ export const ModeledMethodDataGrid = ({
       }
     }
     return [methodsWithModelability, numHiddenMethods];
-  }, [
-    hideModeledMethods,
-    methods,
-    modeledMethodsMap,
-    modifiedSignatures,
-    processedByAutoModelMethods,
-  ]);
+  }, [hideModeledMethods, methods, modeledMethodsMap, modifiedSignatures]);
 
   const someMethodsAreVisible = methodsWithModelability.length > 0;