Merge pull request #2952 from github/robertbrignull/save_multiple_modeled_methods

robertbrignull · web-flow · commit 1993db5122d1 · 2023-10-12T09:19:16.000+01:00
Implement editing / saving multiple modeled methods from the model editor
diff --git a/extensions/ql-vscode/src/common/interface-types.ts b/extensions/ql-vscode/src/common/interface-types.ts
@@ -614,7 +614,7 @@ export type FromModelEditorMessage =
   | StopGeneratingMethodsFromLlmMessage
   | ModelDependencyMessage
   | HideModeledMethodsMessage
-  | SetModeledMethodMessage;
+  | SetMultipleModeledMethodsMessage;
 
 interface RevealInEditorMessage {
   t: "revealInModelEditor";
diff --git a/extensions/ql-vscode/src/model-editor/model-editor-view.ts b/extensions/ql-vscode/src/model-editor/model-editor-view.ts
@@ -43,7 +43,6 @@ import { AutoModeler } from "./auto-modeler";
 import { telemetryListener } from "../common/vscode/telemetry";
 import { ModelingStore } from "./modeling-store";
 import { ModelEditorViewTracker } from "./model-editor-view-tracker";
-import { convertFromLegacyModeledMethod } from "./shared/modeled-methods-legacy";
 
 export class ModelEditorView extends AbstractWebview<
   ToModelEditorMessage,
@@ -310,11 +309,8 @@ export class ModelEditorView extends AbstractWebview<
           "model-editor-hide-modeled-methods",
         );
         break;
-      case "setModeledMethod": {
-        this.setModeledMethods(
-          msg.method.signature,
-          convertFromLegacyModeledMethod(msg.method),
-        );
+      case "setMultipleModeledMethods": {
+        this.setModeledMethods(msg.methodSignature, msg.modeledMethods);
         break;
       }
       case "telemetry":
diff --git a/extensions/ql-vscode/src/model-editor/shared/modeling-status.ts b/extensions/ql-vscode/src/model-editor/shared/modeling-status.ts
@@ -3,13 +3,13 @@ import { ModeledMethod } from "../modeled-method";
 export type ModelingStatus = "unmodeled" | "unsaved" | "saved";
 
 export function getModelingStatus(
-  modeledMethods: ModeledMethod[],
+  modeledMethods: Array<ModeledMethod | undefined>,
   methodIsUnsaved: boolean,
 ): ModelingStatus {
   if (modeledMethods.length > 0) {
     if (methodIsUnsaved) {
       return "unsaved";
-    } else if (modeledMethods.some((m) => m.type !== "none")) {
+    } else if (modeledMethods.some((m) => m && m.type !== "none")) {
       return "saved";
     }
   }
diff --git a/extensions/ql-vscode/src/view/model-editor/LibraryRow.tsx b/extensions/ql-vscode/src/view/model-editor/LibraryRow.tsx
@@ -77,7 +77,7 @@ export type LibraryRowProps = {
   viewState: ModelEditorViewState;
   hideModeledMethods: boolean;
   revealedMethodSignature: string | null;
-  onChange: (modeledMethod: ModeledMethod) => void;
+  onChange: (methodSignature: string, modeledMethods: ModeledMethod[]) => void;
   onSaveModelClick: (methodSignatures: string[]) => void;
   onGenerateFromLlmClick: (
     dependencyName: string,
diff --git a/extensions/ql-vscode/src/view/model-editor/MethodRow.tsx b/extensions/ql-vscode/src/view/model-editor/MethodRow.tsx
@@ -5,7 +5,7 @@ import {
   VSCodeProgressRing,
 } from "@vscode/webview-ui-toolkit/react";
 import * as React from "react";
-import { forwardRef, useCallback, useEffect, useRef } from "react";
+import { forwardRef, useCallback, useEffect, useMemo, useRef } from "react";
 import { styled } from "styled-components";
 import { vscode } from "../vscode-api";
 
@@ -68,7 +68,7 @@ export type MethodRowProps = {
   modelingInProgress: boolean;
   viewState: ModelEditorViewState;
   revealedMethodSignature: string | null;
-  onChange: (modeledMethod: ModeledMethod) => void;
+  onChange: (methodSignature: string, modeledMethods: ModeledMethod[]) => void;
 };
 
 export const MethodRow = (props: MethodRowProps) => {
@@ -103,9 +103,20 @@ const ModelableMethodRow = forwardRef<HTMLElement | undefined, MethodRowProps>(
       onChange,
     } = props;
 
-    const modeledMethods = viewState.showMultipleModels
-      ? modeledMethodsProp
-      : modeledMethodsProp.slice(0, 1);
+    const modeledMethods = useMemo(
+      () => modeledMethodsToDisplay(modeledMethodsProp, method, viewState),
+      [modeledMethodsProp, method, viewState],
+    );
+
+    const modeledMethodChangedHandlers = useMemo(
+      () =>
+        modeledMethods.map((_, index) => (modeledMethod: ModeledMethod) => {
+          const newModeledMethods = [...modeledMethods];
+          newModeledMethods[index] = modeledMethod;
+          onChange(method.signature, newModeledMethods);
+        }),
+      [method, modeledMethods, onChange],
+    );
 
     const jumpToMethod = useCallback(
       () => sendJumpToMethodMessage(method),
@@ -153,42 +164,42 @@ const ModelableMethodRow = forwardRef<HTMLElement | undefined, MethodRowProps>(
         {!props.modelingInProgress && (
           <>
             <MultiModelColumn gridColumn={2}>
-              {forEachModeledMethod(modeledMethods, (modeledMethod, index) => (
+              {modeledMethods.map((modeledMethod, index) => (
                 <ModelTypeDropdown
                   key={index}
                   method={method}
                   modeledMethod={modeledMethod}
-                  onChange={onChange}
+                  onChange={modeledMethodChangedHandlers[index]}
                 />
               ))}
             </MultiModelColumn>
             <MultiModelColumn gridColumn={3}>
-              {forEachModeledMethod(modeledMethods, (modeledMethod, index) => (
+              {modeledMethods.map((modeledMethod, index) => (
                 <ModelInputDropdown
                   key={index}
                   method={method}
                   modeledMethod={modeledMethod}
-                  onChange={onChange}
+                  onChange={modeledMethodChangedHandlers[index]}
                 />
               ))}
             </MultiModelColumn>
             <MultiModelColumn gridColumn={4}>
-              {forEachModeledMethod(modeledMethods, (modeledMethod, index) => (
+              {modeledMethods.map((modeledMethod, index) => (
                 <ModelOutputDropdown
                   key={index}
                   method={method}
                   modeledMethod={modeledMethod}
-                  onChange={onChange}
+                  onChange={modeledMethodChangedHandlers[index]}
                 />
               ))}
             </MultiModelColumn>
             <MultiModelColumn gridColumn={5}>
-              {forEachModeledMethod(modeledMethods, (modeledMethod, index) => (
+              {modeledMethods.map((modeledMethod, index) => (
                 <ModelKindDropdown
                   key={index}
                   method={method}
                   modeledMethod={modeledMethod}
-                  onChange={onChange}
+                  onChange={modeledMethodChangedHandlers[index]}
                 />
               ))}
             </MultiModelColumn>
@@ -245,16 +256,31 @@ function sendJumpToMethodMessage(method: Method) {
   });
 }
 
-function forEachModeledMethod(
+function modeledMethodsToDisplay(
   modeledMethods: ModeledMethod[],
-  renderer: (
-    modeledMethod: ModeledMethod | undefined,
-    index: number,
-  ) => JSX.Element,
-): JSX.Element | JSX.Element[] {
+  method: Method,
+  viewState: ModelEditorViewState,
+): ModeledMethod[] {
   if (modeledMethods.length === 0) {
-    return renderer(undefined, 0);
+    return [
+      {
+        type: "none",
+        input: "",
+        output: "",
+        kind: "",
+        provenance: "manual",
+        signature: method.signature,
+        packageName: method.packageName,
+        typeName: method.typeName,
+        methodName: method.methodName,
+        methodParameters: method.methodParameters,
+      },
+    ];
+  }
+
+  if (viewState.showMultipleModels) {
+    return modeledMethods;
   } else {
-    return modeledMethods.map(renderer);
+    return modeledMethods.slice(0, 1);
   }
 }
diff --git a/extensions/ql-vscode/src/view/model-editor/ModelEditor.tsx b/extensions/ql-vscode/src/view/model-editor/ModelEditor.tsx
@@ -180,12 +180,16 @@ export function ModelEditor({
     [methods],
   );
 
-  const onChange = useCallback((model: ModeledMethod) => {
-    vscode.postMessage({
-      t: "setModeledMethod",
-      method: model,
-    });
-  }, []);
+  const onChange = useCallback(
+    (methodSignature: string, modeledMethods: ModeledMethod[]) => {
+      vscode.postMessage({
+        t: "setMultipleModeledMethods",
+        methodSignature,
+        modeledMethods,
+      });
+    },
+    [],
+  );
 
   const onRefreshClick = useCallback(() => {
     vscode.postMessage({
diff --git a/extensions/ql-vscode/src/view/model-editor/ModeledMethodDataGrid.tsx b/extensions/ql-vscode/src/view/model-editor/ModeledMethodDataGrid.tsx
@@ -24,7 +24,7 @@ export type ModeledMethodDataGridProps = {
   viewState: ModelEditorViewState;
   hideModeledMethods: boolean;
   revealedMethodSignature: string | null;
-  onChange: (modeledMethod: ModeledMethod) => void;
+  onChange: (methodSignature: string, modeledMethods: ModeledMethod[]) => void;
 };
 
 export const ModeledMethodDataGrid = ({
diff --git a/extensions/ql-vscode/src/view/model-editor/ModeledMethodsList.tsx b/extensions/ql-vscode/src/view/model-editor/ModeledMethodsList.tsx
@@ -19,7 +19,7 @@ export type ModeledMethodsListProps = {
   revealedMethodSignature: string | null;
   viewState: ModelEditorViewState;
   hideModeledMethods: boolean;
-  onChange: (modeledMethod: ModeledMethod) => void;
+  onChange: (methodSignature: string, modeledMethods: ModeledMethod[]) => void;
   onSaveModelClick: (methodSignatures: string[]) => void;
   onGenerateFromLlmClick: (
     packageName: string,
diff --git a/extensions/ql-vscode/src/view/model-editor/__tests__/MethodRow.spec.tsx b/extensions/ql-vscode/src/view/model-editor/__tests__/MethodRow.spec.tsx
@@ -73,6 +73,33 @@ describe(MethodRow.name, () => {
     expect(screen.queryByLabelText("Loading")).not.toBeInTheDocument();
   });
 
+  it("can change the type when there is no modeled method", async () => {
+    render({ modeledMethods: [] });
+
+    onChange.mockReset();
+
+    await userEvent.selectOptions(
+      screen.getByRole("combobox", { name: "Model type" }),
+      "source",
+    );
+
+    expect(onChange).toHaveBeenCalledTimes(1);
+    expect(onChange).toHaveBeenCalledWith(method.signature, [
+      {
+        type: "source",
+        input: "Argument[0]",
+        output: "ReturnValue",
+        kind: "value",
+        provenance: "manual",
+        signature: method.signature,
+        packageName: method.packageName,
+        typeName: method.typeName,
+        methodName: method.methodName,
+        methodParameters: method.methodParameters,
+      },
+    ]);
+  });
+
   it("can change the kind", async () => {
     render();
 
@@ -86,10 +113,12 @@ describe(MethodRow.name, () => {
     );
 
     expect(onChange).toHaveBeenCalledTimes(1);
-    expect(onChange).toHaveBeenCalledWith({
-      ...modeledMethod,
-      kind: "value",
-    });
+    expect(onChange).toHaveBeenCalledWith(method.signature, [
+      {
+        ...modeledMethod,
+        kind: "value",
+      },
+    ]);
   });
 
   it("has the correct input options", () => {
@@ -181,6 +210,38 @@ describe(MethodRow.name, () => {
     expect(kindInputs[0]).toHaveValue("source");
   });
 
+  it("can update fields when there are multiple models", async () => {
+    render({
+      modeledMethods: [
+        { ...modeledMethod, type: "source" },
+        { ...modeledMethod, type: "sink", kind: "code-injection" },
+        { ...modeledMethod, type: "summary" },
+      ],
+      viewState: {
+        ...viewState,
+        showMultipleModels: true,
+      },
+    });
+
+    onChange.mockReset();
+
+    expect(screen.getAllByRole("combobox", { name: "Kind" })[1]).toHaveValue(
+      "code-injection",
+    );
+
+    await userEvent.selectOptions(
+      screen.getAllByRole("combobox", { name: "Kind" })[1],
+      "sql-injection",
+    );
+
+    expect(onChange).toHaveBeenCalledTimes(1);
+    expect(onChange).toHaveBeenCalledWith(method.signature, [
+      { ...modeledMethod, type: "source" },
+      { ...modeledMethod, type: "sink", kind: "sql-injection" },
+      { ...modeledMethod, type: "summary" },
+    ]);
+  });
+
   it("renders an unmodelable method", () => {
     render({
       methodCanBeModeled: false,
