Move modeled and modified method state to store (#2861)

Author: charisk

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

@@ -500,14 +500,14 @@ interface SetMethodsMessage {
   methods: Method[];
 }
 
-interface LoadModeledMethodsMessage {
-  t: "loadModeledMethods";
-  modeledMethods: Record<string, ModeledMethod>;
+interface SetModeledMethodsMessage {
+  t: "setModeledMethods";
+  methods: Record<string, ModeledMethod>;
 }
 
-interface AddModeledMethodsMessage {
-  t: "addModeledMethods";
-  modeledMethods: Record<string, ModeledMethod>;
+interface SetModifiedMethodsMessage {
+  t: "setModifiedMethods";
+  methodSignatures: string[];
 }
 
 interface SetInProgressMethodsMessage {
@@ -570,11 +570,16 @@ interface HideModeledMethodsMessage {
   hideModeledMethods: boolean;
 }
 
+interface SetModeledMethodMessage {
+  t: "setModeledMethod";
+  method: ModeledMethod;
+}
+
 export type ToModelEditorMessage =
   | SetExtensionPackStateMessage
   | SetMethodsMessage
-  | LoadModeledMethodsMessage
-  | AddModeledMethodsMessage
+  | SetModeledMethodsMessage
+  | SetModifiedMethodsMessage
   | SetInProgressMethodsMessage;
 
 export type FromModelEditorMessage =
@@ -589,7 +594,8 @@ export type FromModelEditorMessage =
   | GenerateMethodsFromLlmMessage
   | StopGeneratingMethodsFromLlmMessage
   | ModelDependencyMessage
-  | HideModeledMethodsMessage;
+  | HideModeledMethodsMessage
+  | SetModeledMethodMessage;
 
 export type FromMethodModelingMessage =
   | TelemetryMessage

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

@@ -85,7 +85,7 @@ export class ModelEditorView extends AbstractWebview<
         });
       },
       async (modeledMethods) => {
-        await this.postMessage({ t: "addModeledMethods", modeledMethods });
+        this.addModeledMethods(modeledMethods);
       },
     );
   }
@@ -257,6 +257,12 @@ export class ModelEditorView extends AbstractWebview<
             cancellable: false,
           },
         );
+
+        this.modelingStore.removeModifiedMethods(
+          this.databaseItem,
+          Object.keys(msg.modeledMethods),
+        );
+
         void telemetryListener?.sendUIInteraction(
           "model-editor-save-modeled-methods",
         );
@@ -316,6 +322,10 @@ export class ModelEditorView extends AbstractWebview<
           "model-editor-hide-modeled-methods",
         );
         break;
+      case "setModeledMethod": {
+        this.setModeledMethod(msg.method);
+        break;
+      }
       default:
         assertNever(msg);
     }
@@ -360,10 +370,7 @@ export class ModelEditorView extends AbstractWebview<
         this.cliServer,
         this.app.logger,
       );
-      await this.postMessage({
-        t: "loadModeledMethods",
-        modeledMethods,
-      });
+      this.modelingStore.setModeledMethods(this.databaseItem, modeledMethods);
     } catch (e: unknown) {
       void showAndLogErrorMessage(
         this.app.logger,
@@ -441,10 +448,7 @@ export class ModelEditorView extends AbstractWebview<
                 modeledMethodsByName[modeledMethod.signature] = modeledMethod;
               }
 
-              await this.postMessage({
-                t: "addModeledMethods",
-                modeledMethods: modeledMethodsByName,
-              });
+              this.addModeledMethods(modeledMethodsByName);
             },
             progress,
             token: tokenSource.token,
@@ -597,5 +601,46 @@ export class ModelEditorView extends AbstractWebview<
         }
       }),
     );
+
+    this.push(
+      this.modelingStore.onModeledMethodsChanged(async (event) => {
+        if (event.dbUri === this.databaseItem.databaseUri.toString()) {
+          await this.postMessage({
+            t: "setModeledMethods",
+            methods: event.modeledMethods,
+          });
+        }
+      }),
+    );
+
+    this.push(
+      this.modelingStore.onModifiedMethodsChanged(async (event) => {
+        if (event.dbUri === this.databaseItem.databaseUri.toString()) {
+          await this.postMessage({
+            t: "setModifiedMethods",
+            methodSignatures: [...event.modifiedMethods],
+          });
+        }
+      }),
+    );
+  }
+
+  private addModeledMethods(modeledMethods: Record<string, ModeledMethod>) {
+    this.modelingStore.addModeledMethods(this.databaseItem, modeledMethods);
+
+    this.modelingStore.addModifiedMethods(
+      this.databaseItem,
+      new Set(Object.keys(modeledMethods)),
+    );
+  }
+
+  private setModeledMethod(method: ModeledMethod) {
+    const state = this.modelingStore.getStateForActiveDb();
+    if (!state) {
+      throw new Error("Attempting to set modeled method without active db");
+    }
+
+    this.modelingStore.updateModeledMethod(state.databaseItem, method);
+    this.modelingStore.addModifiedMethod(state.databaseItem, method.signature);
   }
 }

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

@@ -3,12 +3,15 @@ import { DisposableObject } from "../common/disposable-object";
 import { AppEvent, AppEventEmitter } from "../common/events";
 import { DatabaseItem } from "../databases/local-databases";
 import { Method } from "./method";
+import { ModeledMethod } from "./modeled-method";
 import { INITIAL_HIDE_MODELED_METHODS_VALUE } from "./shared/hide-modeled-methods";
 
 interface DbModelingState {
   databaseItem: DatabaseItem;
   methods: Method[];
   hideModeledMethods: boolean;
+  modeledMethods: Record<string, ModeledMethod>;
+  modifiedMethodSignatures: Set<string>;
 }
 
 interface MethodsChangedEvent {
@@ -22,11 +25,25 @@ interface HideModeledMethodsChangedEvent {
   isActiveDb: boolean;
 }
 
+interface ModeledMethodsChangedEvent {
+  modeledMethods: Record<string, ModeledMethod>;
+  dbUri: string;
+  isActiveDb: boolean;
+}
+
+interface ModifiedMethodsChangedEvent {
+  modifiedMethods: Set<string>;
+  dbUri: string;
+  isActiveDb: boolean;
+}
+
 export class ModelingStore extends DisposableObject {
   public readonly onActiveDbChanged: AppEvent<void>;
   public readonly onDbClosed: AppEvent<string>;
   public readonly onMethodsChanged: AppEvent<MethodsChangedEvent>;
   public readonly onHideModeledMethodsChanged: AppEvent<HideModeledMethodsChangedEvent>;
+  public readonly onModeledMethodsChanged: AppEvent<ModeledMethodsChangedEvent>;
+  public readonly onModifiedMethodsChanged: AppEvent<ModifiedMethodsChangedEvent>;
 
   private readonly state: Map<string, DbModelingState>;
   private activeDb: string | undefined;
@@ -35,6 +52,8 @@ export class ModelingStore extends DisposableObject {
   private readonly onDbClosedEventEmitter: AppEventEmitter<string>;
   private readonly onMethodsChangedEventEmitter: AppEventEmitter<MethodsChangedEvent>;
   private readonly onHideModeledMethodsChangedEventEmitter: AppEventEmitter<HideModeledMethodsChangedEvent>;
+  private readonly onModeledMethodsChangedEventEmitter: AppEventEmitter<ModeledMethodsChangedEvent>;
+  private readonly onModifiedMethodsChangedEventEmitter: AppEventEmitter<ModifiedMethodsChangedEvent>;
 
   constructor(app: App) {
     super();
@@ -61,6 +80,18 @@ export class ModelingStore extends DisposableObject {
     );
     this.onHideModeledMethodsChanged =
       this.onHideModeledMethodsChangedEventEmitter.event;
+
+    this.onModeledMethodsChangedEventEmitter = this.push(
+      app.createEventEmitter<ModeledMethodsChangedEvent>(),
+    );
+    this.onModeledMethodsChanged =
+      this.onModeledMethodsChangedEventEmitter.event;
+
+    this.onModifiedMethodsChangedEventEmitter = this.push(
+      app.createEventEmitter<ModifiedMethodsChangedEvent>(),
+    );
+    this.onModifiedMethodsChanged =
+      this.onModifiedMethodsChangedEventEmitter.event;
   }
 
   public initializeStateForDb(databaseItem: DatabaseItem) {
@@ -69,6 +100,8 @@ export class ModelingStore extends DisposableObject {
       databaseItem,
       methods: [],
       hideModeledMethods: INITIAL_HIDE_MODELED_METHODS_VALUE,
+      modeledMethods: {},
+      modifiedMethodSignatures: new Set(),
     });
   }
 
@@ -129,6 +162,73 @@ export class ModelingStore extends DisposableObject {
     });
   }
 
+  public addModeledMethods(
+    dbItem: DatabaseItem,
+    methods: Record<string, ModeledMethod>,
+  ) {
+    this.changeModeledMethods(dbItem, (state) => {
+      const newModeledMethods = {
+        ...methods,
+        ...Object.fromEntries(
+          Object.entries(state.modeledMethods).filter(
+            ([_, value]) => value.type !== "none",
+          ),
+        ),
+      };
+      state.modeledMethods = newModeledMethods;
+    });
+  }
+
+  public setModeledMethods(
+    dbItem: DatabaseItem,
+    methods: Record<string, ModeledMethod>,
+  ) {
+    this.changeModeledMethods(dbItem, (state) => {
+      state.modeledMethods = methods;
+    });
+  }
+
+  public updateModeledMethod(dbItem: DatabaseItem, method: ModeledMethod) {
+    this.changeModeledMethods(dbItem, (state) => {
+      state.modeledMethods[method.signature] = method;
+    });
+  }
+
+  public setModifiedMethods(
+    dbItem: DatabaseItem,
+    methodSignatures: Set<string>,
+  ) {
+    this.changeModifiedMethods(dbItem, (state) => {
+      state.modifiedMethodSignatures = methodSignatures;
+    });
+  }
+
+  public addModifiedMethods(
+    dbItem: DatabaseItem,
+    methodSignatures: Iterable<string>,
+  ) {
+    this.changeModifiedMethods(dbItem, (state) => {
+      for (const signature of methodSignatures) {
+        state.modifiedMethodSignatures.add(signature);
+      }
+    });
+  }
+
+  public addModifiedMethod(dbItem: DatabaseItem, methodSignature: string) {
+    this.addModifiedMethods(dbItem, [methodSignature]);
+  }
+
+  public removeModifiedMethods(
+    dbItem: DatabaseItem,
+    methodSignatures: string[],
+  ) {
+    this.changeModifiedMethods(dbItem, (state) => {
+      methodSignatures.forEach((signature) => {
+        state.modifiedMethodSignatures.delete(signature);
+      });
+    });
+  }
+
   private getState(databaseItem: DatabaseItem): DbModelingState {
     if (!this.state.has(databaseItem.databaseUri.toString())) {
       throw Error(
@@ -138,4 +238,34 @@ export class ModelingStore extends DisposableObject {
 
     return this.state.get(databaseItem.databaseUri.toString())!;
   }
+
+  private changeModifiedMethods(
+    dbItem: DatabaseItem,
+    updateState: (state: DbModelingState) => void,
+  ) {
+    const state = this.getState(dbItem);
+
+    updateState(state);
+
+    this.onModifiedMethodsChangedEventEmitter.fire({
+      modifiedMethods: state.modifiedMethodSignatures,
+      dbUri: dbItem.databaseUri.toString(),
+      isActiveDb: dbItem.databaseUri.toString() === this.activeDb,
+    });
+  }
+
+  private changeModeledMethods(
+    dbItem: DatabaseItem,
+    updateState: (state: DbModelingState) => void,
+  ) {
+    const state = this.getState(dbItem);
+
+    updateState(state);
+
+    this.onModeledMethodsChangedEventEmitter.fire({
+      modeledMethods: state.modeledMethods,
+      dbUri: dbItem.databaseUri.toString(),
+      isActiveDb: dbItem.databaseUri.toString() === this.activeDb,
+    });
+  }
 }