Move in-progress methods state to modeling store (#2971)

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

@@ -25,6 +25,7 @@ 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
@@ -517,8 +518,7 @@ interface SetModifiedMethodsMessage {
 
 interface SetInProgressMethodsMessage {
   t: "setInProgressMethods";
-  packageName: string;
-  inProgressMethods: string[];
+  methods: InProgressMethods;
 }
 
 interface SwitchModeMessage {

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

@@ -16,6 +16,7 @@ import { QueryRunner } from "../query-server";
 import { DatabaseItem } from "../databases/local-databases";
 import { Mode } from "./shared/mode";
 import { CancellationTokenSource } from "vscode";
+import { ModelingStore } from "./modeling-store";
 
 // Limit the number of candidates we send to the model in each request
 // to avoid long requests.
@@ -35,12 +36,9 @@ export class AutoModeler {
     private readonly app: App,
     private readonly cliServer: CodeQLCliServer,
     private readonly queryRunner: QueryRunner,
+    private readonly modelingStore: ModelingStore,
     private readonly queryStorageDir: string,
     private readonly databaseItem: DatabaseItem,
-    private readonly setInProgressMethods: (
-      packageName: string,
-      inProgressMethods: string[],
-    ) => Promise<void>,
     private readonly addModeledMethods: (
       modeledMethods: Record<string, ModeledMethod[]>,
     ) => Promise<void>,
@@ -137,7 +135,8 @@ export class AutoModeler {
           const candidatesToProcess = allCandidateMethods.slice(start, end);
 
           // Let the UI know which candidates we are modeling
-          await this.setInProgressMethods(
+          this.modelingStore.setInProgressMethods(
+            this.databaseItem,
             packageName,
             candidatesToProcess.map((c) => c.signature),
           );
@@ -152,7 +151,11 @@ export class AutoModeler {
         }
       } finally {
         // Clear out in progress methods
-        await this.setInProgressMethods(packageName, []);
+        this.modelingStore.setInProgressMethods(
+          this.databaseItem,
+          packageName,
+          [],
+        );
       }
     });
   }

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

@@ -76,15 +76,9 @@ export class ModelEditorView extends AbstractWebview<
       app,
       cliServer,
       queryRunner,
+      modelingStore,
       queryStorageDir,
       databaseItem,
-      async (packageName, inProgressMethods) => {
-        await this.postMessage({
-          t: "setInProgressMethods",
-          packageName,
-          inProgressMethods,
-        });
-      },
       async (modeledMethods) => {
         this.addModeledMethods(modeledMethods);
       },
@@ -683,6 +677,17 @@ export class ModelEditorView extends AbstractWebview<
         }
       }),
     );
+
+    this.push(
+      this.modelingStore.onInProgressMethodsChanged(async (event) => {
+        if (event.dbUri === this.databaseItem.databaseUri.toString()) {
+          await this.postMessage({
+            t: "setInProgressMethods",
+            methods: event.methods,
+          });
+        }
+      }),
+    );
   }
 
   private registerToModelConfigEvents() {

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

@@ -5,6 +5,7 @@ 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 } from "./shared/in-progress-methods";
 import { INITIAL_MODE, Mode } from "./shared/mode";
 
 interface InternalDbModelingState {
@@ -14,6 +15,7 @@ interface InternalDbModelingState {
   mode: Mode;
   modeledMethods: Record<string, ModeledMethod[]>;
   modifiedMethodSignatures: Set<string>;
+  inProgressMethods: InProgressMethods;
   selectedMethod: Method | undefined;
   selectedUsage: Usage | undefined;
 }
@@ -73,6 +75,11 @@ interface SelectedMethodChangedEvent {
   readonly isModified: boolean;
 }
 
+interface InProgressMethodsChangedEvent {
+  readonly dbUri: string;
+  readonly methods: InProgressMethods;
+}
+
 export class ModelingStore extends DisposableObject {
   public readonly onActiveDbChanged: AppEvent<void>;
   public readonly onDbOpened: AppEvent<string>;
@@ -83,6 +90,7 @@ export class ModelingStore extends DisposableObject {
   public readonly onModeledMethodsChanged: AppEvent<ModeledMethodsChangedEvent>;
   public readonly onModifiedMethodsChanged: AppEvent<ModifiedMethodsChangedEvent>;
   public readonly onSelectedMethodChanged: AppEvent<SelectedMethodChangedEvent>;
+  public readonly onInProgressMethodsChanged: AppEvent<InProgressMethodsChangedEvent>;
 
   private readonly state: Map<string, InternalDbModelingState>;
   private activeDb: string | undefined;
@@ -96,6 +104,7 @@ export class ModelingStore extends DisposableObject {
   private readonly onModeledMethodsChangedEventEmitter: AppEventEmitter<ModeledMethodsChangedEvent>;
   private readonly onModifiedMethodsChangedEventEmitter: AppEventEmitter<ModifiedMethodsChangedEvent>;
   private readonly onSelectedMethodChangedEventEmitter: AppEventEmitter<SelectedMethodChangedEvent>;
+  private readonly onInProgressMethodsChangedEventEmitter: AppEventEmitter<InProgressMethodsChangedEvent>;
 
   constructor(app: App) {
     super();
@@ -148,6 +157,12 @@ export class ModelingStore extends DisposableObject {
     );
     this.onSelectedMethodChanged =
       this.onSelectedMethodChangedEventEmitter.event;
+
+    this.onInProgressMethodsChangedEventEmitter = this.push(
+      app.createEventEmitter<InProgressMethodsChangedEvent>(),
+    );
+    this.onInProgressMethodsChanged =
+      this.onInProgressMethodsChangedEventEmitter.event;
   }
 
   public initializeStateForDb(
@@ -164,6 +179,7 @@ export class ModelingStore extends DisposableObject {
       modifiedMethodSignatures: new Set(),
       selectedMethod: undefined,
       selectedUsage: undefined,
+      inProgressMethods: {},
     });
 
     this.onDbOpenedEventEmitter.fire(dbUri);
@@ -396,6 +412,24 @@ export class ModelingStore extends DisposableObject {
     });
   }
 
+  public setInProgressMethods(
+    dbItem: DatabaseItem,
+    packageName: string,
+    inProgressMethods: string[],
+  ) {
+    const dbState = this.getState(dbItem);
+
+    dbState.inProgressMethods = {
+      ...dbState.inProgressMethods,
+      [packageName]: inProgressMethods,
+    };
+
+    this.onInProgressMethodsChangedEventEmitter.fire({
+      dbUri: dbItem.databaseUri.toString(),
+      methods: dbState.inProgressMethods,
+    });
+  }
+
   public getSelectedMethodDetails(): SelectedMethodDetails | undefined {
     const dbState = this.getInternalStateForActiveDb();
     if (!dbState) {

extensions/ql-vscode/src/model-editor/shared/in-progress-methods.ts

@@ -1,18 +1,7 @@
 /**
  * An interface to help keep track of which methods are in progress for each package.
  */
-export type InProgressMethods = Record<string, string[]>;
-
-export function setPackageInProgressMethods(
-  inProgressMethods: InProgressMethods,
-  packageName: string,
-  methods: string[],
-): InProgressMethods {
-  return {
-    ...inProgressMethods,
-    [packageName]: methods,
-  };
-}
+export type InProgressMethods = Readonly<Record<string, readonly string[]>>;
 
 export function hasInProgressMethod(
   inProgressMethods: InProgressMethods,

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

@@ -17,10 +17,7 @@ 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,
-  setPackageInProgressMethods,
-} from "../../model-editor/shared/in-progress-methods";
+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";
 
@@ -136,15 +133,10 @@ export function ModelEditor({
           case "setModifiedMethods":
             setModifiedSignatures(new Set(msg.methodSignatures));
             break;
-          case "setInProgressMethods":
-            setInProgressMethods((oldInProgressMethods) =>
-              setPackageInProgressMethods(
-                oldInProgressMethods,
-                msg.packageName,
-                msg.inProgressMethods,
-              ),
-            );
+          case "setInProgressMethods": {
+            setInProgressMethods(msg.methods);
             break;
+          }
           case "revealMethod":
             setRevealedMethodSignature(msg.methodSignature);
 

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

@@ -11,6 +11,7 @@ export function createMockModelingStore({
   onModeChanged = jest.fn(),
   onModeledMethodsChanged = jest.fn(),
   onModifiedMethodsChanged = jest.fn(),
+  onInProgressMethodsChanged = jest.fn(),
 }: {
   initializeStateForDb?: ModelingStore["initializeStateForDb"];
   getStateForActiveDb?: ModelingStore["getStateForActiveDb"];
@@ -21,6 +22,7 @@ export function createMockModelingStore({
   onModeChanged?: ModelingStore["onModeChanged"];
   onModeledMethodsChanged?: ModelingStore["onModeledMethodsChanged"];
   onModifiedMethodsChanged?: ModelingStore["onModifiedMethodsChanged"];
+  onInProgressMethodsChanged?: ModelingStore["onInProgressMethodsChanged"];
 } = {}): ModelingStore {
   return mockedObject<ModelingStore>({
     initializeStateForDb,
@@ -32,5 +34,6 @@ export function createMockModelingStore({
     onModeChanged,
     onModeledMethodsChanged,
     onModifiedMethodsChanged,
+    onInProgressMethodsChanged,
   });
 }