Move methods, hideModeledMethods and active editor state to the modeling store

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

@@ -7,12 +7,16 @@ import {
 import { Method, Usage } from "../method";
 import { DatabaseItem } from "../../databases/local-databases";
 import { CodeQLCliServer } from "../../codeql-cli/cli";
+import { ModelingStore } from "../modeling-store";
 
 export class MethodsUsagePanel extends DisposableObject {
   private readonly dataProvider: MethodsUsageDataProvider;
   private readonly treeView: TreeView<MethodsUsageTreeViewItem>;
 
-  public constructor(cliServer: CodeQLCliServer) {
+  public constructor(
+    private readonly modelingStore: ModelingStore,
+    cliServer: CodeQLCliServer,
+  ) {
     super();
 
     this.dataProvider = new MethodsUsageDataProvider(cliServer);
@@ -21,6 +25,8 @@ export class MethodsUsagePanel extends DisposableObject {
       treeDataProvider: this.dataProvider,
     });
     this.push(this.treeView);
+
+    this.registerToModelingStoreEvents();
   }
 
   public async setState(
@@ -44,4 +50,39 @@ export class MethodsUsagePanel extends DisposableObject {
       await this.treeView.reveal(canonicalUsage);
     }
   }
+
+  private registerToModelingStoreEvents(): void {
+    this.push(
+      this.modelingStore.onActiveDbChanged(async () => {
+        await this.handleStateChangeEvent();
+      }),
+    );
+
+    this.push(
+      this.modelingStore.onMethodsChanged(async (event) => {
+        if (event.isActiveDb) {
+          await this.handleStateChangeEvent();
+        }
+      }),
+    );
+
+    this.push(
+      this.modelingStore.onHideModeledMethodsChanged(async (event) => {
+        if (event.isActiveDb) {
+          await this.handleStateChangeEvent();
+        }
+      }),
+    );
+  }
+
+  private async handleStateChangeEvent(): Promise<void> {
+    const activeState = this.modelingStore.getStateForActiveDb();
+    if (activeState !== undefined) {
+      await this.setState(
+        activeState.methods,
+        activeState.databaseItem,
+        activeState.hideModeledMethods,
+      );
+    }
+  }
 }

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

@@ -29,8 +29,6 @@ export class ModelEditorModule extends DisposableObject {
   private readonly methodsUsagePanel: MethodsUsagePanel;
   private readonly methodModelingPanel: MethodModelingPanel;
 
-  private mostRecentlyActiveView: ModelEditorView | undefined = undefined;
-
   private constructor(
     private readonly app: App,
     private readonly databaseManager: DatabaseManager,
@@ -41,24 +39,12 @@ export class ModelEditorModule extends DisposableObject {
     super();
     this.queryStorageDir = join(baseQueryStorageDir, "model-editor-results");
     this.modelingStore = new ModelingStore(app);
-    this.methodsUsagePanel = this.push(new MethodsUsagePanel(cliServer));
+    this.methodsUsagePanel = this.push(
+      new MethodsUsagePanel(this.modelingStore, cliServer),
+    );
     this.methodModelingPanel = this.push(new MethodModelingPanel(app));
   }
 
-  private handleViewBecameActive(view: ModelEditorView): void {
-    this.mostRecentlyActiveView = view;
-  }
-
-  private handleViewWasDisposed(view: ModelEditorView): void {
-    if (this.mostRecentlyActiveView === view) {
-      this.mostRecentlyActiveView = undefined;
-    }
-  }
-
-  private isMostRecentlyActiveView(view: ModelEditorView): boolean {
-    return this.mostRecentlyActiveView === view;
-  }
-
   public static async initialize(
     app: App,
     databaseManager: DatabaseManager,
@@ -165,16 +151,15 @@ export class ModelEditorModule extends DisposableObject {
               db,
               modelFile,
               Mode.Application,
-              this.methodsUsagePanel.setState.bind(this.methodsUsagePanel),
               this.showMethod.bind(this),
-              this.handleViewBecameActive.bind(this),
-              (view) => {
-                this.handleViewWasDisposed(view);
-                void cleanupQueryDir();
-              },
-              this.isMostRecentlyActiveView.bind(this),
             );
 
+            this.modelingStore.onDbClosed(async (dbUri) => {
+              if (dbUri === db.databaseUri.toString()) {
+                await cleanupQueryDir();
+              }
+            });
+
             this.push(view);
             this.push({
               dispose(): void {

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

@@ -41,7 +41,6 @@ import { loadModeledMethods, saveModeledMethods } from "./modeled-method-fs";
 import { pickExtensionPack } from "./extension-pack-picker";
 import { getLanguageDisplayName } from "../common/query-language";
 import { AutoModeler } from "./auto-modeler";
-import { INITIAL_HIDE_MODELED_METHODS_VALUE } from "./shared/hide-modeled-methods";
 import { telemetryListener } from "../common/vscode/telemetry";
 import { ModelingStore } from "./modeling-store";
 
@@ -51,9 +50,6 @@ export class ModelEditorView extends AbstractWebview<
 > {
   private readonly autoModeler: AutoModeler;
 
-  private methods: Method[];
-  private hideModeledMethods: boolean;
-
   public constructor(
     protected readonly app: App,
     private readonly modelingStore: ModelingStore,
@@ -65,24 +61,15 @@ export class ModelEditorView extends AbstractWebview<
     private readonly databaseItem: DatabaseItem,
     private readonly extensionPack: ExtensionPack,
     private mode: Mode,
-    private readonly updateMethodsUsagePanelState: (
-      methods: Method[],
-      databaseItem: DatabaseItem,
-      hideModeledMethods: boolean,
-    ) => Promise<void>,
     private readonly showMethod: (
       method: Method,
       usage: Usage,
     ) => Promise<void>,
-    private readonly handleViewBecameActive: (view: ModelEditorView) => void,
-    private readonly handleViewWasDisposed: (view: ModelEditorView) => void,
-    private readonly isMostRecentlyActiveView: (
-      view: ModelEditorView,
-    ) => boolean,
   ) {
     super(app);
 
     this.modelingStore.initializeStateForDb(databaseItem);
+    this.registerToModelingStoreEvents();
 
     this.autoModeler = new AutoModeler(
       app,
@@ -101,8 +88,6 @@ export class ModelEditorView extends AbstractWebview<
         await this.postMessage({ t: "addModeledMethods", modeledMethods });
       },
     );
-    this.methods = [];
-    this.hideModeledMethods = INITIAL_HIDE_MODELED_METHODS_VALUE;
   }
 
   public async openView() {
@@ -111,12 +96,7 @@ export class ModelEditorView extends AbstractWebview<
 
     panel.onDidChangeViewState(async () => {
       if (panel.active) {
-        this.handleViewBecameActive(this);
-        await this.updateMethodsUsagePanelState(
-          this.methods,
-          this.databaseItem,
-          this.hideModeledMethods,
-        );
+        this.modelingStore.setActiveDb(this.databaseItem);
         await this.markModelEditorAsActive();
       } else {
         await this.updateModelEditorActiveContext();
@@ -124,7 +104,7 @@ export class ModelEditorView extends AbstractWebview<
     });
 
     panel.onDidDispose(() => {
-      this.handleViewWasDisposed(this);
+      this.modelingStore.removeDb(this.databaseItem);
       // onDidDispose is called after the tab has been closed,
       // so we want to check if there are any others still open.
       void this.app.commands.execute(
@@ -313,11 +293,11 @@ export class ModelEditorView extends AbstractWebview<
         break;
       case "switchMode":
         this.mode = msg.mode;
-        this.methods = [];
+        this.modelingStore.setMethods(this.databaseItem, []);
         await Promise.all([
           this.postMessage({
             t: "setMethods",
-            methods: this.methods,
+            methods: [],
           }),
           this.setViewState(),
           withProgress((progress) => this.loadMethods(progress), {
@@ -328,11 +308,9 @@ export class ModelEditorView extends AbstractWebview<
 
         break;
       case "hideModeledMethods":
-        this.hideModeledMethods = msg.hideModeledMethods;
-        await this.updateMethodsUsagePanelState(
-          this.methods,
+        this.modelingStore.setHideModeledMethods(
           this.databaseItem,
-          this.hideModeledMethods,
+          msg.hideModeledMethods,
         );
         void telemetryListener?.sendUIInteraction(
           "model-editor-hide-modeled-methods",
@@ -413,19 +391,8 @@ export class ModelEditorView extends AbstractWebview<
       if (!queryResult) {
         return;
       }
-      this.methods = queryResult;
 
-      await this.postMessage({
-        t: "setMethods",
-        methods: this.methods,
-      });
-      if (this.isMostRecentlyActiveView(this)) {
-        await this.updateMethodsUsagePanelState(
-          this.methods,
-          this.databaseItem,
-          this.hideModeledMethods,
-        );
-      }
+      this.modelingStore.setMethods(this.databaseItem, queryResult);
     } catch (err) {
       void showAndLogExceptionWithTelemetry(
         this.app.logger,
@@ -540,11 +507,7 @@ export class ModelEditorView extends AbstractWebview<
         addedDatabase,
         modelFile,
         Mode.Framework,
-        this.updateMethodsUsagePanelState,
         this.showMethod,
-        this.handleViewBecameActive,
-        this.handleViewWasDisposed,
-        this.isMostRecentlyActiveView,
       );
       await view.openView();
     });
@@ -622,4 +585,17 @@ export class ModelEditorView extends AbstractWebview<
 
     return addedDatabase;
   }
+
+  private registerToModelingStoreEvents() {
+    this.push(
+      this.modelingStore.onMethodsChanged(async (event) => {
+        if (event.dbUri === this.databaseItem.databaseUri.toString()) {
+          await this.postMessage({
+            t: "setMethods",
+            methods: event.methods,
+          });
+        }
+      }),
+    );
+  }
 }

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

@@ -2,26 +2,44 @@ import { App } from "../common/app";
 import { DisposableObject } from "../common/disposable-object";
 import { AppEvent, AppEventEmitter } from "../common/events";
 import { DatabaseItem } from "../databases/local-databases";
+import { Method } from "./method";
+import { INITIAL_HIDE_MODELED_METHODS_VALUE } from "./shared/hide-modeled-methods";
 
 interface DbModelingState {
-  // Currently empty but will soon contain information about methods, etc.
+  databaseItem: DatabaseItem;
+  methods: Method[];
+  hideModeledMethods: boolean;
+}
+
+interface MethodsChangedEvent {
+  methods: Method[];
+  dbUri: string;
+  isActiveDb: boolean;
+}
+
+interface HideModeledMethodsChangedEvent {
+  hideModeledMethods: boolean;
+  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>;
 
-  private state: Map<string, DbModelingState>;
+  private readonly state: Map<string, DbModelingState>;
   private activeDb: string | undefined;
 
   private readonly onActiveDbChangedEventEmitter: AppEventEmitter<void>;
   private readonly onDbClosedEventEmitter: AppEventEmitter<string>;
+  private readonly onMethodsChangedEventEmitter: AppEventEmitter<MethodsChangedEvent>;
+  private readonly onHideModeledMethodsChangedEventEmitter: AppEventEmitter<HideModeledMethodsChangedEvent>;
 
   constructor(app: App) {
     super();
 
     // State initialization
-    this.activeDb = undefined;
     this.state = new Map<string, DbModelingState>();
 
     // Event initialization
@@ -32,12 +50,25 @@ export class ModelingStore extends DisposableObject {
 
     this.onDbClosedEventEmitter = this.push(app.createEventEmitter<string>());
     this.onDbClosed = this.onDbClosedEventEmitter.event;
+
+    this.onMethodsChangedEventEmitter = this.push(
+      app.createEventEmitter<MethodsChangedEvent>(),
+    );
+    this.onMethodsChanged = this.onMethodsChangedEventEmitter.event;
+
+    this.onHideModeledMethodsChangedEventEmitter = this.push(
+      app.createEventEmitter<HideModeledMethodsChangedEvent>(),
+    );
+    this.onHideModeledMethodsChanged =
+      this.onHideModeledMethodsChangedEventEmitter.event;
   }
 
   public initializeStateForDb(databaseItem: DatabaseItem) {
     const dbUri = databaseItem.databaseUri.toString();
     this.state.set(dbUri, {
       databaseItem,
+      methods: [],
+      hideModeledMethods: INITIAL_HIDE_MODELED_METHODS_VALUE,
     });
   }
 
@@ -61,4 +92,50 @@ export class ModelingStore extends DisposableObject {
     this.state.delete(dbUri);
     this.onDbClosedEventEmitter.fire(dbUri);
   }
+
+  public getStateForActiveDb(): DbModelingState | undefined {
+    if (!this.activeDb) {
+      return undefined;
+    }
+
+    return this.state.get(this.activeDb);
+  }
+
+  public setMethods(dbItem: DatabaseItem, methods: Method[]) {
+    const dbState = this.getState(dbItem);
+    const dbUri = dbItem.databaseUri.toString();
+
+    dbState.methods = methods;
+
+    this.onMethodsChangedEventEmitter.fire({
+      methods,
+      dbUri,
+      isActiveDb: dbUri === this.activeDb,
+    });
+  }
+
+  public setHideModeledMethods(
+    dbItem: DatabaseItem,
+    hideModeledMethods: boolean,
+  ) {
+    const dbState = this.getState(dbItem);
+    const dbUri = dbItem.databaseUri.toString();
+
+    dbState.hideModeledMethods = hideModeledMethods;
+
+    this.onHideModeledMethodsChangedEventEmitter.fire({
+      hideModeledMethods,
+      isActiveDb: dbUri === this.activeDb,
+    });
+  }
+
+  private getState(databaseItem: DatabaseItem): DbModelingState {
+    if (!this.state.has(databaseItem.databaseUri.toString())) {
+      throw Error(
+        "Cannot get state for a database that has not been initialized",
+      );
+    }
+
+    return this.state.get(databaseItem.databaseUri.toString())!;
+  }
 }

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

@@ -0,0 +1,27 @@
+import { mockedObject } from "../../vscode-tests/utils/mocking.helpers";
+import { ModelingStore } from "../../../src/model-editor/modeling-store";
+
+export function createMockModelingStore({
+  initializeStateForDb = jest.fn(),
+  getStateForActiveDb = jest.fn(),
+  onActiveDbChanged = jest.fn(),
+  onDbClosed = jest.fn(),
+  onMethodsChanged = jest.fn(),
+  onHideModeledMethodsChanged = jest.fn(),
+}: {
+  initializeStateForDb?: ModelingStore["initializeStateForDb"];
+  getStateForActiveDb?: ModelingStore["getStateForActiveDb"];
+  onActiveDbChanged?: ModelingStore["onActiveDbChanged"];
+  onDbClosed?: ModelingStore["onDbClosed"];
+  onMethodsChanged?: ModelingStore["onMethodsChanged"];
+  onHideModeledMethodsChanged?: ModelingStore["onHideModeledMethodsChanged"];
+} = {}): ModelingStore {
+  return mockedObject<ModelingStore>({
+    initializeStateForDb,
+    getStateForActiveDb,
+    onActiveDbChanged,
+    onDbClosed,
+    onMethodsChanged,
+    onHideModeledMethodsChanged,
+  });
+}

extensions/ql-vscode/test/vscode-tests/no-workspace/model-editor/methods-usage/methods-usage-panel.test.ts

@@ -8,6 +8,8 @@ import {
   createMethod,
   createUsage,
 } from "../../../../factories/model-editor/method-factories";
+import { ModelingStore } from "../../../../../src/model-editor/modeling-store";
+import { createMockModelingStore } from "../../../../__mocks__/model-editor/modelingStoreMock";
 
 describe("MethodsUsagePanel", () => {
   const mockCliServer = mockedObject<CodeQLCliServer>({});
@@ -25,7 +27,9 @@ describe("MethodsUsagePanel", () => {
       } as TreeView<unknown>;
       jest.spyOn(window, "createTreeView").mockReturnValue(mockTreeView);
 
-      const panel = new MethodsUsagePanel(mockCliServer);
+      const modelingStore = createMockModelingStore();
+
+      const panel = new MethodsUsagePanel(modelingStore, mockCliServer);
       await panel.setState(methods, dbItem, hideModeledMethods);
 
       expect(mockTreeView.badge?.value).toBe(1);
@@ -34,6 +38,7 @@ describe("MethodsUsagePanel", () => {
 
   describe("revealItem", () => {
     let mockTreeView: TreeView<unknown>;
+    let modelingStore: ModelingStore;
 
     const hideModeledMethods: boolean = false;
     const usage = createUsage();
@@ -43,6 +48,8 @@ describe("MethodsUsagePanel", () => {
         reveal: jest.fn(),
       });
       jest.spyOn(window, "createTreeView").mockReturnValue(mockTreeView);
+
+      modelingStore = createMockModelingStore();
     });
 
     it("should reveal the correct item in the tree view", async () => {
@@ -52,7 +59,7 @@ describe("MethodsUsagePanel", () => {
         }),
       ];
 
-      const panel = new MethodsUsagePanel(mockCliServer);
+      const panel = new MethodsUsagePanel(modelingStore, mockCliServer);
       await panel.setState(methods, dbItem, hideModeledMethods);
 
       await panel.revealItem(usage);
@@ -62,7 +69,7 @@ describe("MethodsUsagePanel", () => {
 
     it("should do nothing if usage cannot be found", async () => {
       const methods = [createMethod({})];
-      const panel = new MethodsUsagePanel(mockCliServer);
+      const panel = new MethodsUsagePanel(modelingStore, mockCliServer);
       await panel.setState(methods, dbItem, hideModeledMethods);
 
       await panel.revealItem(usage);

extensions/ql-vscode/test/vscode-tests/no-workspace/model-editor/model-editor-view.test.ts

@@ -8,11 +8,11 @@ import { createMockApp } from "../../../__mocks__/appMock";
 import { mockEmptyDatabaseManager } from "../query-testing/test-runner-helpers";
 import { QueryRunner } from "../../../../src/query-server";
 import { ExtensionPack } from "../../../../src/model-editor/shared/extension-pack";
-import { ModelingStore } from "../../../../src/model-editor/modeling-store";
+import { createMockModelingStore } from "../../../__mocks__/model-editor/modelingStoreMock";
 
 describe("ModelEditorView", () => {
   const app = createMockApp({});
-  const modelingStore = mockedObject<ModelingStore>({});
+  const modelingStore = createMockModelingStore();
   const databaseManager = mockEmptyDatabaseManager();
   const cliServer = mockedObject<CodeQLCliServer>({});
   const queryRunner = mockedObject<QueryRunner>({});
@@ -31,11 +31,7 @@ describe("ModelEditorView", () => {
     dataExtensions: ["models/**/*.yml"],
   };
   const mode = Mode.Application;
-  const updateMethodsUsagePanelState = jest.fn();
   const showMethod = jest.fn();
-  const handleViewBecameActive = jest.fn();
-  const handleViewWasDisposed = jest.fn();
-  const isMostRecentlyActiveView = jest.fn();
 
   let view: ModelEditorView;
 
@@ -51,11 +47,7 @@ describe("ModelEditorView", () => {
       databaseItem,
       extensionPack,
       mode,
-      updateMethodsUsagePanelState,
       showMethod,
-      handleViewBecameActive,
-      handleViewWasDisposed,
-      isMostRecentlyActiveView,
     );
   });