Allow editing multiple models in model editor

Author: robertbrignull

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

@@ -614,7 +614,7 @@ export type FromModelEditorMessage =
   | StopGeneratingMethodsFromLlmMessage
   | ModelDependencyMessage
   | HideModeledMethodsMessage
-  | SetModeledMethodMessage;
+  | SetMultipleModeledMethodsMessage;
 
 interface RevealInEditorMessage {
   t: "revealInModelEditor";

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,
@@ -309,11 +308,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":

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,

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

@@ -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) => {
@@ -113,6 +113,21 @@ const ModelableMethodRow = forwardRef<HTMLElement | undefined, MethodRowProps>(
       [modeledMethodsProp, viewState],
     );
 
+    const modeledMethodChangedHandlers = useMemo(
+      () =>
+        modeledMethods.map((_, index) => (modeledMethod: ModeledMethod) => {
+          const newModeledMethods = [...modeledMethods];
+          newModeledMethods[index] = modeledMethod;
+          onChange(
+            method.signature,
+            newModeledMethods.filter(
+              (m): m is ModeledMethod => m !== undefined,
+            ),
+          );
+        }),
+      [method, modeledMethods, onChange],
+    );
+
     const jumpToMethod = useCallback(
       () => sendJumpToMethodMessage(method),
       [method],
@@ -164,7 +179,7 @@ const ModelableMethodRow = forwardRef<HTMLElement | undefined, MethodRowProps>(
                   key={index}
                   method={method}
                   modeledMethod={modeledMethod}
-                  onChange={onChange}
+                  onChange={modeledMethodChangedHandlers[index]}
                 />
               ))}
             </MultiModelColumn>
@@ -174,7 +189,7 @@ const ModelableMethodRow = forwardRef<HTMLElement | undefined, MethodRowProps>(
                   key={index}
                   method={method}
                   modeledMethod={modeledMethod}
-                  onChange={onChange}
+                  onChange={modeledMethodChangedHandlers[index]}
                 />
               ))}
             </MultiModelColumn>
@@ -184,7 +199,7 @@ const ModelableMethodRow = forwardRef<HTMLElement | undefined, MethodRowProps>(
                   key={index}
                   method={method}
                   modeledMethod={modeledMethod}
-                  onChange={onChange}
+                  onChange={modeledMethodChangedHandlers[index]}
                 />
               ))}
             </MultiModelColumn>
@@ -194,7 +209,7 @@ const ModelableMethodRow = forwardRef<HTMLElement | undefined, MethodRowProps>(
                   key={index}
                   method={method}
                   modeledMethod={modeledMethod}
-                  onChange={onChange}
+                  onChange={modeledMethodChangedHandlers[index]}
                 />
               ))}
             </MultiModelColumn>

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({

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 = ({

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,