Simplify modeling in progress state

extensions/ql-vscode/src/common/interface-types.ts

@@ -25,7 +25,6 @@ import {
 } from "../model-editor/shared/view-state";
 import { Mode } from "../model-editor/shared/mode";
 import { QueryLanguage } from "./query-language";
-import { InProgressMethods } from "../model-editor/shared/in-progress-methods";
 
 /**
  * This module contains types and code that are shared between
@@ -518,7 +517,7 @@ interface SetModifiedMethodsMessage {
 
 interface SetInProgressMethodsMessage {
   t: "setInProgressMethods";
-  methods: InProgressMethods;
+  methods: string[];
 }
 
 interface SwitchModeMessage {

extensions/ql-vscode/src/model-editor/auto-modeler.ts

@@ -133,12 +133,14 @@ export class AutoModeler {
           const start = i * candidateBatchSize;
           const end = start + candidateBatchSize;
           const candidatesToProcess = allCandidateMethods.slice(start, end);
+          const candidateSignatures = candidatesToProcess.map(
+            (c) => c.signature,
+          );
 
           // Let the UI know which candidates we are modeling
-          this.modelingStore.setInProgressMethods(
+          this.modelingStore.addInProgressMethods(
             this.databaseItem,
-            packageName,
-            candidatesToProcess.map((c) => c.signature),
+            candidateSignatures,
           );
 
           // Kick off the process to model the slice of candidates
@@ -148,13 +150,18 @@ export class AutoModeler {
             progress,
             cancellationTokenSource,
           );
+
+          // Let the UI know which candidates we are done modeling
+          this.modelingStore.removeInProgressMethods(
+            this.databaseItem,
+            candidateSignatures,
+          );
         }
       } finally {
-        // Clear out in progress methods
-        this.modelingStore.setInProgressMethods(
+        // Clear out in progress methods in case anything went wrong
+        this.modelingStore.removeInProgressMethods(
           this.databaseItem,
-          packageName,
-          [],
+          allCandidateMethods.map((c) => c.signature),
         );
       }
     });

extensions/ql-vscode/src/model-editor/method-modeling/method-modeling-view-provider.ts

@@ -14,7 +14,6 @@ import { assertNever } from "../../common/helpers-pure";
 import { ModelEditorViewTracker } from "../model-editor-view-tracker";
 import { ModelConfigListener } from "../../config";
 import { DatabaseItem } from "../../databases/local-databases";
-import { hasInProgressMethodSignature } from "../shared/in-progress-methods";
 
 export class MethodModelingViewProvider extends AbstractWebviewViewProvider<
   ToMethodModelingMessage,
@@ -225,10 +224,7 @@ export class MethodModelingViewProvider extends AbstractWebviewViewProvider<
         if (this.method && this.databaseItem) {
           const dbUri = this.databaseItem.databaseUri.toString();
           if (e.dbUri === dbUri) {
-            const inProgress = hasInProgressMethodSignature(
-              e.methods,
-              this.method.signature,
-            );
+            const inProgress = e.methods.has(this.method.signature);
             await this.postMessage({
               t: "setInProgress",
               inProgress,

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

@@ -693,7 +693,7 @@ export class ModelEditorView extends AbstractWebview<
         if (event.dbUri === this.databaseItem.databaseUri.toString()) {
           await this.postMessage({
             t: "setInProgressMethods",
-            methods: event.methods,
+            methods: Array.from(event.methods),
           });
         }
       }),

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

@@ -5,10 +5,6 @@ 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 {
-  InProgressMethods,
-  hasInProgressMethodSignature,
-} from "./shared/in-progress-methods";
 import { INITIAL_MODE, Mode } from "./shared/mode";
 
 interface InternalDbModelingState {
@@ -18,7 +14,7 @@ interface InternalDbModelingState {
   mode: Mode;
   modeledMethods: Record<string, ModeledMethod[]>;
   modifiedMethodSignatures: Set<string>;
-  inProgressMethods: InProgressMethods;
+  inProgressMethods: Set<string>;
   selectedMethod: Method | undefined;
   selectedUsage: Usage | undefined;
 }
@@ -30,6 +26,7 @@ interface DbModelingState {
   readonly mode: Mode;
   readonly modeledMethods: Readonly<Record<string, readonly ModeledMethod[]>>;
   readonly modifiedMethodSignatures: ReadonlySet<string>;
+  readonly inProgressMethods: ReadonlySet<string>;
   readonly selectedMethod: Method | undefined;
   readonly selectedUsage: Usage | undefined;
 }
@@ -82,7 +79,7 @@ interface SelectedMethodChangedEvent {
 
 interface InProgressMethodsChangedEvent {
   readonly dbUri: string;
-  readonly methods: InProgressMethods;
+  readonly methods: ReadonlySet<string>;
 }
 
 export class ModelingStore extends DisposableObject {
@@ -184,7 +181,7 @@ export class ModelingStore extends DisposableObject {
       modifiedMethodSignatures: new Set(),
       selectedMethod: undefined,
       selectedUsage: undefined,
-      inProgressMethods: {},
+      inProgressMethods: new Set(),
     });
 
     this.onDbOpenedEventEmitter.fire(dbUri);
@@ -425,28 +422,32 @@ export class ModelingStore extends DisposableObject {
       usage,
       modeledMethods: dbState.modeledMethods[method.signature] ?? [],
       isModified: dbState.modifiedMethodSignatures.has(method.signature),
-      isInProgress: hasInProgressMethodSignature(
-        dbState.inProgressMethods,
-        method.signature,
-      ),
+      isInProgress: dbState.inProgressMethods.has(method.signature),
     });
   }
 
-  public setInProgressMethods(
+  public addInProgressMethods(
     dbItem: DatabaseItem,
-    packageName: string,
     inProgressMethods: string[],
   ) {
-    const dbState = this.getState(dbItem);
-
-    dbState.inProgressMethods = {
-      ...dbState.inProgressMethods,
-      [packageName]: inProgressMethods,
-    };
+    this.changeInProgressMethods(dbItem, (state) => {
+      state.inProgressMethods = new Set([
+        ...state.inProgressMethods,
+        ...inProgressMethods,
+      ]);
+    });
+  }
 
-    this.onInProgressMethodsChangedEventEmitter.fire({
-      dbUri: dbItem.databaseUri.toString(),
-      methods: dbState.inProgressMethods,
+  public removeInProgressMethods(
+    dbItem: DatabaseItem,
+    methodSignatures: string[],
+  ) {
+    this.changeInProgressMethods(dbItem, (state) => {
+      state.inProgressMethods = new Set(
+        Array.from(state.inProgressMethods).filter(
+          (s) => !methodSignatures.includes(s),
+        ),
+      );
     });
   }
 
@@ -469,10 +470,7 @@ export class ModelingStore extends DisposableObject {
       isModified: dbState.modifiedMethodSignatures.has(
         selectedMethod.signature,
       ),
-      isInProgress: hasInProgressMethodSignature(
-        dbState.inProgressMethods,
-        selectedMethod.signature,
-      ),
+      isInProgress: dbState.inProgressMethods.has(selectedMethod.signature),
     };
   }
 
@@ -515,4 +513,18 @@ export class ModelingStore extends DisposableObject {
       isActiveDb: dbItem.databaseUri.toString() === this.activeDb,
     });
   }
+
+  private changeInProgressMethods(
+    dbItem: DatabaseItem,
+    updateState: (state: InternalDbModelingState) => void,
+  ) {
+    const state = this.getState(dbItem);
+
+    updateState(state);
+
+    this.onInProgressMethodsChangedEventEmitter.fire({
+      dbUri: dbItem.databaseUri.toString(),
+      methods: state.inProgressMethods,
+    });
+  }
 }

extensions/ql-vscode/src/stories/model-editor/LibraryRow.stories.tsx

@@ -218,7 +218,7 @@ LibraryRow.args = {
     ],
   },
   modifiedSignatures: new Set(["org.sql2o.Sql2o#Sql2o(String)"]),
-  inProgressMethods: {},
+  inProgressMethods: new Set(),
   viewState: {
     extensionPack: createMockExtensionPack(),
     showFlowGeneration: true,

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

@@ -14,10 +14,6 @@ import {
   VSCodeTag,
 } from "@vscode/webview-ui-toolkit/react";
 import { ModelEditorViewState } from "../../model-editor/shared/view-state";
-import {
-  InProgressMethods,
-  hasInProgressMethod,
-} from "../../model-editor/shared/in-progress-methods";
 
 const LibraryContainer = styled.div`
   background-color: var(--vscode-peekViewResult-background);
@@ -76,7 +72,7 @@ export type LibraryRowProps = {
   methods: Method[];
   modeledMethodsMap: Record<string, ModeledMethod[]>;
   modifiedSignatures: Set<string>;
-  inProgressMethods: InProgressMethods;
+  inProgressMethods: Set<string>;
   viewState: ModelEditorViewState;
   hideModeledMethods: boolean;
   revealedMethodSignature: string | null;
@@ -179,10 +175,8 @@ export const LibraryRow = ({
   }, [methods, modifiedSignatures]);
 
   const canStopAutoModeling = useMemo(() => {
-    return methods.some((method) =>
-      hasInProgressMethod(inProgressMethods, title, method.signature),
-    );
-  }, [methods, title, inProgressMethods]);
+    return methods.some((method) => inProgressMethods.has(method.signature));
+  }, [methods, inProgressMethods]);
 
   return (
     <LibraryContainer>

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

@@ -17,7 +17,6 @@ import { ModelEditorViewState } from "../../model-editor/shared/view-state";
 import { ModeledMethodsList } from "./ModeledMethodsList";
 import { percentFormatter } from "./formatters";
 import { Mode } from "../../model-editor/shared/mode";
-import { InProgressMethods } from "../../model-editor/shared/in-progress-methods";
 import { getLanguageDisplayName } from "../../common/query-language";
 import { INITIAL_HIDE_MODELED_METHODS_VALUE } from "../../model-editor/shared/hide-modeled-methods";
 
@@ -93,8 +92,8 @@ export function ModelEditor({
     new Set(),
   );
 
-  const [inProgressMethods, setInProgressMethods] = useState<InProgressMethods>(
-    {},
+  const [inProgressMethods, setInProgressMethods] = useState<Set<string>>(
+    new Set(),
   );
 
   const [hideModeledMethods, setHideModeledMethods] = useState(
@@ -134,7 +133,7 @@ export function ModelEditor({
             setModifiedSignatures(new Set(msg.methodSignatures));
             break;
           case "setInProgressMethods": {
-            setInProgressMethods(msg.methods);
+            setInProgressMethods(new Set(msg.methods));
             break;
           }
           case "revealMethod":

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

@@ -9,10 +9,6 @@ import { Method, canMethodBeModeled } from "../../model-editor/method";
 import { ModeledMethod } from "../../model-editor/modeled-method";
 import { useMemo } from "react";
 import { sortMethods } from "../../model-editor/shared/sorting";
-import {
-  InProgressMethods,
-  hasInProgressMethod,
-} from "../../model-editor/shared/in-progress-methods";
 import { HiddenMethodsRow } from "./HiddenMethodsRow";
 import { ModelEditorViewState } from "../../model-editor/shared/view-state";
 import { ScreenReaderOnly } from "../common/ScreenReaderOnly";
@@ -27,7 +23,7 @@ export type ModeledMethodDataGridProps = {
   methods: Method[];
   modeledMethodsMap: Record<string, ModeledMethod[]>;
   modifiedSignatures: Set<string>;
-  inProgressMethods: InProgressMethods;
+  inProgressMethods: Set<string>;
   viewState: ModelEditorViewState;
   hideModeledMethods: boolean;
   revealedMethodSignature: string | null;
@@ -110,11 +106,7 @@ export const ModeledMethodDataGrid = ({
                 methodCanBeModeled={methodCanBeModeled}
                 modeledMethods={modeledMethods}
                 methodIsUnsaved={modifiedSignatures.has(method.signature)}
-                modelingInProgress={hasInProgressMethod(
-                  inProgressMethods,
-                  packageName,
-                  method.signature,
-                )}
+                modelingInProgress={inProgressMethods.has(method.signature)}
                 viewState={viewState}
                 revealedMethodSignature={revealedMethodSignature}
                 onChange={onChange}

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

@@ -9,13 +9,12 @@ import {
   sortGroupNames,
 } from "../../model-editor/shared/sorting";
 import { ModelEditorViewState } from "../../model-editor/shared/view-state";
-import { InProgressMethods } from "../../model-editor/shared/in-progress-methods";
 
 export type ModeledMethodsListProps = {
   methods: Method[];
   modeledMethodsMap: Record<string, ModeledMethod[]>;
   modifiedSignatures: Set<string>;
-  inProgressMethods: InProgressMethods;
+  inProgressMethods: Set<string>;
   revealedMethodSignature: string | null;
   viewState: ModelEditorViewState;
   hideModeledMethods: boolean;

extensions/ql-vscode/src/view/model-editor/__tests__/LibraryRow.spec.tsx

@@ -43,7 +43,7 @@ describe(LibraryRow.name, () => {
           ],
         }}
         modifiedSignatures={new Set([method.signature])}
-        inProgressMethods={{}}
+        inProgressMethods={new Set()}
         viewState={viewState}
         hideModeledMethods={false}
         revealedMethodSignature={null}

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

@@ -69,7 +69,7 @@ describe(ModeledMethodDataGrid.name, () => {
           ],
         }}
         modifiedSignatures={new Set([method1.signature])}
-        inProgressMethods={{}}
+        inProgressMethods={new Set()}
         viewState={viewState}
         hideModeledMethods={false}
         revealedMethodSignature={null}

extensions/ql-vscode/src/view/model-editor/__tests__/ModeledMethodsList.spec.tsx

@@ -69,7 +69,7 @@ describe(ModeledMethodsList.name, () => {
           ],
         }}
         modifiedSignatures={new Set([method1.signature])}
-        inProgressMethods={{}}
+        inProgressMethods={new Set()}
         viewState={viewState}
         hideModeledMethods={false}
         revealedMethodSignature={null}