Add reveal in editor button to method modeling panel

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

@@ -575,12 +575,18 @@ interface SetModeledMethodMessage {
   method: ModeledMethod;
 }
 
+interface RevealMethodMessage {
+  t: "revealMethod";
+  method: Method;
+}
+
 export type ToModelEditorMessage =
   | SetExtensionPackStateMessage
   | SetMethodsMessage
   | SetModeledMethodsMessage
   | SetModifiedMethodsMessage
-  | SetInProgressMethodsMessage;
+  | SetInProgressMethodsMessage
+  | RevealMethodMessage;
 
 export type FromModelEditorMessage =
   | ViewLoadedMsg
@@ -597,9 +603,15 @@ export type FromModelEditorMessage =
   | HideModeledMethodsMessage
   | SetModeledMethodMessage;
 
+interface RevealInEditorMessage {
+  t: "revealInModelEditor";
+  method: Method;
+}
+
 export type FromMethodModelingMessage =
   | CommonFromViewMessages
-  | SetModeledMethodMessage;
+  | SetModeledMethodMessage
+  | RevealInEditorMessage;
 
 interface SetMethodMessage {
   t: "setMethod";

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

@@ -4,14 +4,23 @@ import { DisposableObject } from "../../common/disposable-object";
 import { MethodModelingViewProvider } from "./method-modeling-view-provider";
 import { Method } from "../method";
 import { ModelingStore } from "../modeling-store";
+import { ModelEditorViewTracker } from "../model-editor-view-tracker";
 
 export class MethodModelingPanel extends DisposableObject {
   private readonly provider: MethodModelingViewProvider;
 
-  constructor(app: App, modelingStore: ModelingStore) {
+  constructor(
+    app: App,
+    modelingStore: ModelingStore,
+    editorViewTracker: ModelEditorViewTracker,
+  ) {
     super();
 
-    this.provider = new MethodModelingViewProvider(app, modelingStore);
+    this.provider = new MethodModelingViewProvider(
+      app,
+      modelingStore,
+      editorViewTracker,
+    );
     this.push(
       window.registerWebviewViewProvider(
         MethodModelingViewProvider.viewType,

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

@@ -8,8 +8,10 @@ import { extLogger } from "../../common/logging/vscode/loggers";
 import { App } from "../../common/app";
 import { redactableError } from "../../common/errors";
 import { Method } from "../method";
-import { ModelingStore } from "../modeling-store";
+import { DbModelingState, ModelingStore } from "../modeling-store";
 import { AbstractWebviewViewProvider } from "../../common/vscode/abstract-webview-view-provider";
+import { assertNever } from "../../common/helpers-pure";
+import { ModelEditorViewTracker } from "../model-editor-view-tracker";
 
 export class MethodModelingViewProvider extends AbstractWebviewViewProvider<
   ToMethodModelingMessage,
@@ -22,6 +24,7 @@ export class MethodModelingViewProvider extends AbstractWebviewViewProvider<
   constructor(
     app: App,
     private readonly modelingStore: ModelingStore,
+    private readonly editorViewTracker: ModelEditorViewTracker,
   ) {
     super(app, "method-modeling");
   }
@@ -77,17 +80,43 @@ export class MethodModelingViewProvider extends AbstractWebviewViewProvider<
         break;
 
       case "setModeledMethod": {
-        const activeState = this.modelingStore.getStateForActiveDb();
-        if (!activeState) {
-          throw new Error("No active state found in modeling store");
-        }
+        const activeState = this.ensureActiveState();
+
         this.modelingStore.updateModeledMethod(
           activeState.databaseItem,
           msg.method,
         );
         break;
       }
+      case "revealInModelEditor":
+        await this.revealInModelEditor(msg.method);
+
+        break;
+      default:
+        assertNever(msg);
+    }
+  }
+
+  private async revealInModelEditor(method: Method): Promise<void> {
+    const activeState = this.ensureActiveState();
+
+    const views = this.editorViewTracker.getViews(
+      activeState.databaseItem.databaseUri.toString(),
+    );
+    if (views.length === 0) {
+      return;
     }
+
+    await Promise.all(views.map((view) => view.revealMethod(method)));
+  }
+
+  private ensureActiveState(): DbModelingState {
+    const activeState = this.modelingStore.getStateForActiveDb();
+    if (!activeState) {
+      throw new Error("No active state found in modeling store");
+    }
+
+    return activeState;
   }
 
   private registerToModelingStoreEvents(): void {

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

@@ -20,12 +20,14 @@ import { setUpPack } from "./model-editor-queries";
 import { MethodModelingPanel } from "./method-modeling/method-modeling-panel";
 import { ModelingStore } from "./modeling-store";
 import { showResolvableLocation } from "../databases/local-databases/locations";
+import { ModelEditorViewTracker } from "./model-editor-view-tracker";
 
 const SUPPORTED_LANGUAGES: string[] = ["java", "csharp"];
 
 export class ModelEditorModule extends DisposableObject {
   private readonly queryStorageDir: string;
   private readonly modelingStore: ModelingStore;
+  private readonly editorViewTracker: ModelEditorViewTracker<ModelEditorView>;
   private readonly methodsUsagePanel: MethodsUsagePanel;
   private readonly methodModelingPanel: MethodModelingPanel;
 
@@ -39,11 +41,12 @@ export class ModelEditorModule extends DisposableObject {
     super();
     this.queryStorageDir = join(baseQueryStorageDir, "model-editor-results");
     this.modelingStore = new ModelingStore(app);
+    this.editorViewTracker = new ModelEditorViewTracker();
     this.methodsUsagePanel = this.push(
       new MethodsUsagePanel(this.modelingStore, cliServer),
     );
     this.methodModelingPanel = this.push(
-      new MethodModelingPanel(app, this.modelingStore),
+      new MethodModelingPanel(app, this.modelingStore, this.editorViewTracker),
     );
 
     this.registerToModelingStoreEvents();
@@ -148,6 +151,7 @@ export class ModelEditorModule extends DisposableObject {
             const view = new ModelEditorView(
               this.app,
               this.modelingStore,
+              this.editorViewTracker,
               this.databaseManager,
               this.cliServer,
               this.queryRunner,

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

@@ -0,0 +1,41 @@
+import { Method } from "./method";
+
+interface ModelEditorViewInterface {
+  databaseUri: string;
+
+  revealMethod(method: Method): Promise<void>;
+}
+
+export class ModelEditorViewTracker<
+  T extends ModelEditorViewInterface = ModelEditorViewInterface,
+> {
+  private readonly views = new Map<string, T[]>();
+
+  constructor() {}
+
+  public registerView(view: T): void {
+    const databaseUri = view.databaseUri;
+
+    if (!this.views.has(databaseUri)) {
+      this.views.set(databaseUri, []);
+    }
+
+    this.views.get(databaseUri)?.push(view);
+  }
+
+  public unregisterView(view: T): void {
+    const views = this.views.get(view.databaseUri);
+    if (!views) {
+      return;
+    }
+
+    const index = views.indexOf(view);
+    if (index !== -1) {
+      views.splice(index, 1);
+    }
+  }
+
+  public getViews(databaseUri: string): T[] {
+    return this.views.get(databaseUri) ?? [];
+  }
+}

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

@@ -42,6 +42,7 @@ import { getLanguageDisplayName } from "../common/query-language";
 import { AutoModeler } from "./auto-modeler";
 import { telemetryListener } from "../common/vscode/telemetry";
 import { ModelingStore } from "./modeling-store";
+import { ModelEditorViewTracker } from "./model-editor-view-tracker";
 
 export class ModelEditorView extends AbstractWebview<
   ToModelEditorMessage,
@@ -52,6 +53,7 @@ export class ModelEditorView extends AbstractWebview<
   public constructor(
     protected readonly app: App,
     private readonly modelingStore: ModelingStore,
+    private readonly viewTracker: ModelEditorViewTracker<ModelEditorView>,
     private readonly databaseManager: DatabaseManager,
     private readonly cliServer: CodeQLCliServer,
     private readonly queryRunner: QueryRunner,
@@ -66,6 +68,8 @@ export class ModelEditorView extends AbstractWebview<
     this.modelingStore.initializeStateForDb(databaseItem);
     this.registerToModelingStoreEvents();
 
+    this.viewTracker.registerView(this);
+
     this.autoModeler = new AutoModeler(
       app,
       cliServer,
@@ -181,7 +185,7 @@ export class ModelEditorView extends AbstractWebview<
   }
 
   protected onPanelDispose(): void {
-    // Nothing to do here
+    this.viewTracker.unregisterView(this);
   }
 
   protected async onMessage(msg: FromModelEditorMessage): Promise<void> {
@@ -338,6 +342,19 @@ export class ModelEditorView extends AbstractWebview<
     ]);
   }
 
+  public get databaseUri(): string {
+    return this.databaseItem.databaseUri.toString();
+  }
+
+  public async revealMethod(method: Method): Promise<void> {
+    this.panel?.reveal();
+
+    await this.postMessage({
+      t: "revealMethod",
+      method,
+    });
+  }
+
   private async setViewState(): Promise<void> {
     const showLlmButton =
       this.databaseItem.language === "java" && showLlmGeneration();
@@ -497,6 +514,7 @@ export class ModelEditorView extends AbstractWebview<
       const view = new ModelEditorView(
         this.app,
         this.modelingStore,
+        this.viewTracker,
         this.databaseManager,
         this.cliServer,
         this.queryRunner,

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

@@ -6,7 +6,7 @@ import { Method, Usage } from "./method";
 import { ModeledMethod } from "./modeled-method";
 import { INITIAL_HIDE_MODELED_METHODS_VALUE } from "./shared/hide-modeled-methods";
 
-interface DbModelingState {
+export interface DbModelingState {
   databaseItem: DatabaseItem;
   methods: Method[];
   hideModeledMethods: boolean;

extensions/ql-vscode/src/view/method-modeling/MethodModeling.tsx

@@ -7,6 +7,7 @@ import { MethodName } from "../model-editor/MethodName";
 import { ModeledMethod } from "../../model-editor/modeled-method";
 import { MethodModelingInputs } from "./MethodModelingInputs";
 import { VSCodeTag } from "@vscode/webview-ui-toolkit/react";
+import { ReviewInEditorButton } from "./ReviewInEditorButton";
 
 const Container = styled.div`
   padding: 0.3rem;
@@ -64,6 +65,7 @@ export const MethodModeling = ({
         modeledMethod={modeledMethod}
         onChange={onChange}
       />
+      <ReviewInEditorButton method={method} />
     </Container>
   );
 };

extensions/ql-vscode/src/view/method-modeling/ReviewInEditorButton.tsx

@@ -0,0 +1,25 @@
+import * as React from "react";
+import { useCallback } from "react";
+import { styled } from "styled-components";
+import { vscode } from "../vscode-api";
+import TextButton from "../common/TextButton";
+import { Method } from "../../model-editor/method";
+
+const Button = styled(TextButton)`
+  margin-top: 0.5rem;
+`;
+
+type Props = {
+  method: Method;
+};
+
+export const ReviewInEditorButton = ({ method }: Props) => {
+  const handleClick = useCallback(() => {
+    vscode.postMessage({
+      t: "revealInModelEditor",
+      method,
+    });
+  }, [method]);
+
+  return <Button onClick={handleClick}>Review in editor</Button>;
+};

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

@@ -1,5 +1,5 @@
 import * as React from "react";
-import { useCallback, useMemo, useState } from "react";
+import { useCallback, useEffect, useMemo, useState } from "react";
 import { styled } from "styled-components";
 import { Method } from "../../model-editor/method";
 import { ModeledMethod } from "../../model-editor/modeled-method";
@@ -76,6 +76,7 @@ export type LibraryRowProps = {
   inProgressMethods: InProgressMethods;
   viewState: ModelEditorViewState;
   hideModeledMethods: boolean;
+  revealedMethodSignature: string | null;
   onChange: (modeledMethod: ModeledMethod) => void;
   onSaveModelClick: (
     methods: Method[],
@@ -100,6 +101,7 @@ export const LibraryRow = ({
   inProgressMethods,
   viewState,
   hideModeledMethods,
+  revealedMethodSignature,
   onChange,
   onSaveModelClick,
   onGenerateFromLlmClick,
@@ -117,6 +119,14 @@ export const LibraryRow = ({
     setExpanded((oldIsExpanded) => !oldIsExpanded);
   }, []);
 
+  useEffect(() => {
+    // If any of the methods in this group is the one that should be revealed, we should expand
+    // this group so the method can highlight itself.
+    if (methods.some((m) => m.signature === revealedMethodSignature)) {
+      setExpanded(true);
+    }
+  }, [methods, revealedMethodSignature]);
+
   const handleModelWithAI = useCallback(
     async (e: React.MouseEvent) => {
       onGenerateFromLlmClick(title, methods, modeledMethods);
@@ -227,6 +237,7 @@ export const LibraryRow = ({
             inProgressMethods={inProgressMethods}
             mode={viewState.mode}
             hideModeledMethods={hideModeledMethods}
+            revealedMethodSignature={revealedMethodSignature}
             onChange={onChange}
           />
           <SectionDivider />