Convert MethodRow to display multiple modelings

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

@@ -7,6 +7,9 @@ import { CallClassification, Method } from "../../model-editor/method";
 import { ModeledMethod } from "../../model-editor/modeled-method";
 import { VSCodeDataGrid } from "@vscode/webview-ui-toolkit/react";
 import { GRID_TEMPLATE_COLUMNS } from "../../view/model-editor/ModeledMethodDataGrid";
+import { ModelEditorViewState } from "../../model-editor/shared/view-state";
+import { createMockExtensionPack } from "../../../test/factories/model-editor/extension-pack";
+import { Mode } from "../../model-editor/shared/mode";
 
 export default {
   title: "CodeQL Model Editor/Method Row",
@@ -66,51 +69,66 @@ const modeledMethod: ModeledMethod = {
   methodParameters: "()",
 };
 
+const viewState: ModelEditorViewState = {
+  extensionPack: createMockExtensionPack(),
+  showFlowGeneration: true,
+  showLlmButton: true,
+  showMultipleModels: true,
+  mode: Mode.Application,
+};
+
 export const Unmodeled = Template.bind({});
 Unmodeled.args = {
   method,
-  modeledMethod: undefined,
+  modeledMethods: [],
   methodCanBeModeled: true,
+  viewState,
 };
 
 export const Source = Template.bind({});
 Source.args = {
   method,
-  modeledMethod: { ...modeledMethod, type: "source" },
+  modeledMethods: [{ ...modeledMethod, type: "source" }],
   methodCanBeModeled: true,
+  viewState,
 };
 
 export const Sink = Template.bind({});
 Sink.args = {
   method,
-  modeledMethod: { ...modeledMethod, type: "sink" },
+  modeledMethods: [{ ...modeledMethod, type: "sink" }],
   methodCanBeModeled: true,
+  viewState,
 };
 
 export const Summary = Template.bind({});
 Summary.args = {
   method,
-  modeledMethod: { ...modeledMethod, type: "summary" },
+  modeledMethods: [{ ...modeledMethod, type: "summary" }],
   methodCanBeModeled: true,
+  viewState,
 };
 
 export const Neutral = Template.bind({});
 Neutral.args = {
   method,
-  modeledMethod: { ...modeledMethod, type: "neutral" },
+  modeledMethods: [{ ...modeledMethod, type: "neutral" }],
   methodCanBeModeled: true,
+  viewState,
 };
 
 export const AlreadyModeled = Template.bind({});
 AlreadyModeled.args = {
   method: { ...method, supported: true },
-  modeledMethod: undefined,
+  modeledMethods: [],
+  viewState,
 };
 
 export const ModelingInProgress = Template.bind({});
 ModelingInProgress.args = {
   method,
-  modeledMethod,
+  modeledMethods: [modeledMethod],
   modelingInProgress: true,
   methodCanBeModeled: true,
+  viewState,
 };

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

@@ -23,6 +23,12 @@ import { ModelInputDropdown } from "./ModelInputDropdown";
 import { ModelOutputDropdown } from "./ModelOutputDropdown";
 import { ModelEditorViewState } from "../../model-editor/shared/view-state";
 
+const MultiModelColumn = styled(VSCodeDataGridCell)`
+  display: flex;
+  flex-direction: column;
+  gap: 0.5em;
+`;
+
 const ApiOrMethodRow = styled.div`
   min-height: calc(var(--input-height) * 1px);
   display: flex;
@@ -57,7 +63,7 @@ const DataGridRow = styled(VSCodeDataGridRow)<{ focused?: boolean }>`
 export type MethodRowProps = {
   method: Method;
   methodCanBeModeled: boolean;
-  modeledMethod: ModeledMethod | undefined;
+  modeledMethods: ModeledMethod[];
   methodIsUnsaved: boolean;
   modelingInProgress: boolean;
   viewState: ModelEditorViewState;
@@ -90,22 +96,23 @@ const ModelableMethodRow = forwardRef<HTMLElement | undefined, MethodRowProps>(
   (props, ref) => {
     const {
       method,
-      modeledMethod,
+      modeledMethods: modeledMethodsArg,
       methodIsUnsaved,
       viewState,
       revealedMethodSignature,
       onChange,
     } = props;
 
+    const modeledMethods = viewState.showMultipleModels
+      ? modeledMethodsArg
+      : modeledMethodsArg.slice(0, 1);
+
     const jumpToUsage = useCallback(
       () => sendJumpToUsageMessage(method),
       [method],
     );
 
-    const modelingStatus = getModelingStatus(
-      modeledMethod ? [modeledMethod] : [],
-      methodIsUnsaved,
-    );
+    const modelingStatus = getModelingStatus(modeledMethods, methodIsUnsaved);
 
     return (
       <DataGridRow
@@ -145,34 +152,46 @@ const ModelableMethodRow = forwardRef<HTMLElement | undefined, MethodRowProps>(
         )}
         {!props.modelingInProgress && (
           <>
-            <VSCodeDataGridCell gridColumn={2}>
-              <ModelTypeDropdown
-                method={method}
-                modeledMethod={modeledMethod}
-                onChange={onChange}
-              />
-            </VSCodeDataGridCell>
-            <VSCodeDataGridCell gridColumn={3}>
-              <ModelInputDropdown
-                method={method}
-                modeledMethod={modeledMethod}
-                onChange={onChange}
-              />
-            </VSCodeDataGridCell>
-            <VSCodeDataGridCell gridColumn={4}>
-              <ModelOutputDropdown
-                method={method}
-                modeledMethod={modeledMethod}
-                onChange={onChange}
-              />
-            </VSCodeDataGridCell>
-            <VSCodeDataGridCell gridColumn={5}>
-              <ModelKindDropdown
-                method={method}
-                modeledMethod={modeledMethod}
-                onChange={onChange}
-              />
-            </VSCodeDataGridCell>
+            <MultiModelColumn gridColumn={2}>
+              {forEachModeledMethod(modeledMethods, (modeledMethod) => (
+                <ModelTypeDropdown
+                  key={JSON.stringify(modeledMethod)}
+                  method={method}
+                  modeledMethod={modeledMethod}
+                  onChange={onChange}
+                />
+              ))}
+            </MultiModelColumn>
+            <MultiModelColumn gridColumn={3}>
+              {forEachModeledMethod(modeledMethods, (modeledMethod) => (
+                <ModelInputDropdown
+                  key={JSON.stringify(modeledMethod)}
+                  method={method}
+                  modeledMethod={modeledMethod}
+                  onChange={onChange}
+                />
+              ))}
+            </MultiModelColumn>
+            <MultiModelColumn gridColumn={4}>
+              {forEachModeledMethod(modeledMethods, (modeledMethod) => (
+                <ModelOutputDropdown
+                  key={JSON.stringify(modeledMethod)}
+                  method={method}
+                  modeledMethod={modeledMethod}
+                  onChange={onChange}
+                />
+              ))}
+            </MultiModelColumn>
+            <MultiModelColumn gridColumn={5}>
+              {forEachModeledMethod(modeledMethods, (modeledMethod) => (
+                <ModelKindDropdown
+                  key={JSON.stringify(modeledMethod)}
+                  method={method}
+                  modeledMethod={modeledMethod}
+                  onChange={onChange}
+                />
+              ))}
+            </MultiModelColumn>
           </>
         )}
       </DataGridRow>
@@ -227,3 +246,14 @@ function sendJumpToUsageMessage(method: Method) {
     usage: method.usages[0],
   });
 }
+
+function forEachModeledMethod(
+  modeledMethods: ModeledMethod[],
+  renderer: (modeledMethod: ModeledMethod | undefined) => JSX.Element,
+): JSX.Element | JSX.Element[] {
+  if (modeledMethods.length === 0) {
+    return renderer(undefined);
+  } else {
+    return modeledMethods.map(renderer);
+  }
+}

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

@@ -84,22 +84,25 @@ export const ModeledMethodDataGrid = ({
               Kind
             </VSCodeDataGridCell>
           </VSCodeDataGridRow>
-          {methodsWithModelability.map(({ method, methodCanBeModeled }) => (
-            <MethodRow
-              key={method.signature}
-              method={method}
-              methodCanBeModeled={methodCanBeModeled}
-              modeledMethod={modeledMethods[method.signature]}
-              methodIsUnsaved={modifiedSignatures.has(method.signature)}
-              modelingInProgress={inProgressMethods.hasMethod(
-                packageName,
-                method.signature,
-              )}
-              viewState={viewState}
-              revealedMethodSignature={revealedMethodSignature}
-              onChange={onChange}
-            />
-          ))}
+          {methodsWithModelability.map(({ method, methodCanBeModeled }) => {
+            const modeledMethod = modeledMethods[method.signature];
+            return (
+              <MethodRow
+                key={method.signature}
+                method={method}
+                methodCanBeModeled={methodCanBeModeled}
+                modeledMethods={modeledMethod ? [modeledMethod] : []}
+                methodIsUnsaved={modifiedSignatures.has(method.signature)}
+                modelingInProgress={inProgressMethods.hasMethod(
+                  packageName,
+                  method.signature,
+                )}
+                viewState={viewState}
+                revealedMethodSignature={revealedMethodSignature}
+                onChange={onChange}
+              />
+            );
+          })}
         </>
       )}
       <HiddenMethodsRow

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

@@ -46,7 +46,7 @@ describe(MethodRow.name, () => {
       <MethodRow
         method={method}
         methodCanBeModeled={true}
-        modeledMethod={modeledMethod}
+        modeledMethods={[modeledMethod]}
         methodIsUnsaved={false}
         modelingInProgress={false}
         revealedMethodSignature={null}
@@ -120,7 +120,7 @@ describe(MethodRow.name, () => {
 
   it("shows the modeling status indicator when unmodeled", () => {
     render({
-      modeledMethod: undefined,
+      modeledMethods: [],
     });
 
     expect(screen.getByLabelText("Method not modeled")).toBeInTheDocument();
@@ -137,7 +137,7 @@ describe(MethodRow.name, () => {
   it("renders an unmodelable method", () => {
     render({
       methodCanBeModeled: false,
-      modeledMethod: undefined,
+      modeledMethods: [],
     });
 
     expect(screen.queryByRole("combobox")).not.toBeInTheDocument();