Merge pull request #2910 from github/robertbrignull/multiple-models-method-row

robertbrignull · web-flow · commit 27c4bd83492a · 2023-10-09T16:17:17.000+01:00
Add ability for MethodRow to render multiple modelings of the same method
diff --git a/extensions/ql-vscode/src/model-editor/methods-usage/methods-usage-data-provider.ts b/extensions/ql-vscode/src/model-editor/methods-usage/methods-usage-data-provider.ts
@@ -102,7 +102,10 @@ export class MethodsUsageDataProvider
     const modeledMethod = this.modeledMethods[method.signature];
     const modifiedMethod = this.modifiedMethodSignatures.has(method.signature);
 
-    const status = getModelingStatus(modeledMethod, modifiedMethod);
+    const status = getModelingStatus(
+      modeledMethod ? [modeledMethod] : [],
+      modifiedMethod,
+    );
     switch (status) {
       case "unmodeled":
         return new ThemeIcon("error", new ThemeColor("errorForeground"));
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(
-  modeledMethod: ModeledMethod | undefined,
+  modeledMethods: ModeledMethod[],
   methodIsUnsaved: boolean,
 ): ModelingStatus {
-  if (modeledMethod) {
+  if (modeledMethods.length > 0) {
     if (methodIsUnsaved) {
       return "unsaved";
-    } else if (modeledMethod.type !== "none") {
+    } else if (modeledMethods.some((m) => m.type !== "none")) {
       return "saved";
     }
   }
diff --git a/extensions/ql-vscode/src/stories/model-editor/MethodRow.stories.tsx b/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,78 @@ 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,
+};
+
+export const MultipleModelings = Template.bind({});
+MultipleModelings.args = {
+  method,
+  modeledMethods: [
+    { ...modeledMethod, type: "source" },
+    { ...modeledMethod, type: "sink" },
+    { ...modeledMethod },
+  ],
+  methodCanBeModeled: true,
+  viewState,
 };
diff --git a/extensions/ql-vscode/src/view/method-modeling/MethodModelingView.tsx b/extensions/ql-vscode/src/view/method-modeling/MethodModelingView.tsx
@@ -30,7 +30,8 @@ export function MethodModelingView({ initialViewState }: Props): JSX.Element {
   const [isMethodModified, setIsMethodModified] = useState<boolean>(false);
 
   const modelingStatus = useMemo(
-    () => getModelingStatus(modeledMethod, isMethodModified),
+    () =>
+      getModelingStatus(modeledMethod ? [modeledMethod] : [], isMethodModified),
     [modeledMethod, isMethodModified],
   );
 
diff --git a/extensions/ql-vscode/src/view/model-editor/LibraryRow.tsx b/extensions/ql-vscode/src/view/model-editor/LibraryRow.tsx
@@ -232,7 +232,7 @@ export const LibraryRow = ({
             modeledMethods={modeledMethods}
             modifiedSignatures={modifiedSignatures}
             inProgressMethods={inProgressMethods}
-            mode={viewState.mode}
+            viewState={viewState}
             hideModeledMethods={hideModeledMethods}
             revealedMethodSignature={revealedMethodSignature}
             onChange={onChange}
diff --git a/extensions/ql-vscode/src/view/model-editor/MethodRow.tsx b/extensions/ql-vscode/src/view/model-editor/MethodRow.tsx
@@ -21,8 +21,16 @@ import { MethodName } from "./MethodName";
 import { ModelTypeDropdown } from "./ModelTypeDropdown";
 import { ModelInputDropdown } from "./ModelInputDropdown";
 import { ModelOutputDropdown } from "./ModelOutputDropdown";
+import { ModelEditorViewState } from "../../model-editor/shared/view-state";
 
-const ApiOrMethodCell = styled(VSCodeDataGridCell)`
+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;
   flex-direction: row;
   align-items: center;
@@ -55,10 +63,10 @@ const DataGridRow = styled(VSCodeDataGridRow)<{ focused?: boolean }>`
 export type MethodRowProps = {
   method: Method;
   methodCanBeModeled: boolean;
-  modeledMethod: ModeledMethod | undefined;
+  modeledMethods: ModeledMethod[];
   methodIsUnsaved: boolean;
   modelingInProgress: boolean;
-  mode: Mode;
+  viewState: ModelEditorViewState;
   revealedMethodSignature: string | null;
   onChange: (modeledMethod: ModeledMethod) => void;
 };
@@ -88,38 +96,44 @@ const ModelableMethodRow = forwardRef<HTMLElement | undefined, MethodRowProps>(
   (props, ref) => {
     const {
       method,
-      modeledMethod,
+      modeledMethods: modeledMethodsProp,
       methodIsUnsaved,
-      mode,
+      viewState,
       revealedMethodSignature,
       onChange,
     } = props;
 
+    const modeledMethods = viewState.showMultipleModels
+      ? modeledMethodsProp
+      : modeledMethodsProp.slice(0, 1);
+
     const jumpToUsage = useCallback(
       () => sendJumpToUsageMessage(method),
       [method],
     );
 
-    const modelingStatus = getModelingStatus(modeledMethod, methodIsUnsaved);
+    const modelingStatus = getModelingStatus(modeledMethods, methodIsUnsaved);
 
     return (
       <DataGridRow
         data-testid="modelable-method-row"
         ref={ref}
         focused={revealedMethodSignature === method.signature}
       >
-        <ApiOrMethodCell gridColumn={1}>
-          <ModelingStatusIndicator status={modelingStatus} />
-          <MethodClassifications method={method} />
-          <MethodName {...props.method} />
-          {mode === Mode.Application && (
-            <UsagesButton onClick={jumpToUsage}>
-              {method.usages.length}
-            </UsagesButton>
-          )}
-          <ViewLink onClick={jumpToUsage}>View</ViewLink>
-          {props.modelingInProgress && <ProgressRing />}
-        </ApiOrMethodCell>
+        <VSCodeDataGridCell gridColumn={1}>
+          <ApiOrMethodRow>
+            <ModelingStatusIndicator status={modelingStatus} />
+            <MethodClassifications method={method} />
+            <MethodName {...props.method} />
+            {viewState.mode === Mode.Application && (
+              <UsagesButton onClick={jumpToUsage}>
+                {method.usages.length}
+              </UsagesButton>
+            )}
+            <ViewLink onClick={jumpToUsage}>View</ViewLink>
+            {props.modelingInProgress && <ProgressRing />}
+          </ApiOrMethodRow>
+        </VSCodeDataGridCell>
         {props.modelingInProgress && (
           <>
             <VSCodeDataGridCell gridColumn={2}>
@@ -138,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, index) => (
+                <ModelTypeDropdown
+                  key={index}
+                  method={method}
+                  modeledMethod={modeledMethod}
+                  onChange={onChange}
+                />
+              ))}
+            </MultiModelColumn>
+            <MultiModelColumn gridColumn={3}>
+              {forEachModeledMethod(modeledMethods, (modeledMethod, index) => (
+                <ModelInputDropdown
+                  key={index}
+                  method={method}
+                  modeledMethod={modeledMethod}
+                  onChange={onChange}
+                />
+              ))}
+            </MultiModelColumn>
+            <MultiModelColumn gridColumn={4}>
+              {forEachModeledMethod(modeledMethods, (modeledMethod, index) => (
+                <ModelOutputDropdown
+                  key={index}
+                  method={method}
+                  modeledMethod={modeledMethod}
+                  onChange={onChange}
+                />
+              ))}
+            </MultiModelColumn>
+            <MultiModelColumn gridColumn={5}>
+              {forEachModeledMethod(modeledMethods, (modeledMethod, index) => (
+                <ModelKindDropdown
+                  key={index}
+                  method={method}
+                  modeledMethod={modeledMethod}
+                  onChange={onChange}
+                />
+              ))}
+            </MultiModelColumn>
           </>
         )}
       </DataGridRow>
@@ -178,7 +204,7 @@ const UnmodelableMethodRow = forwardRef<
   HTMLElement | undefined,
   MethodRowProps
 >((props, ref) => {
-  const { method, mode, revealedMethodSignature } = props;
+  const { method, viewState, revealedMethodSignature } = props;
 
   const jumpToUsage = useCallback(
     () => sendJumpToUsageMessage(method),
@@ -191,17 +217,19 @@ const UnmodelableMethodRow = forwardRef<
       ref={ref}
       focused={revealedMethodSignature === method.signature}
     >
-      <ApiOrMethodCell gridColumn={1}>
-        <ModelingStatusIndicator status="saved" />
-        <MethodName {...props.method} />
-        {mode === Mode.Application && (
-          <UsagesButton onClick={jumpToUsage}>
-            {method.usages.length}
-          </UsagesButton>
-        )}
-        <ViewLink onClick={jumpToUsage}>View</ViewLink>
-        <MethodClassifications method={method} />
-      </ApiOrMethodCell>
+      <VSCodeDataGridCell gridColumn={1}>
+        <ApiOrMethodRow>
+          <ModelingStatusIndicator status="saved" />
+          <MethodName {...props.method} />
+          {viewState.mode === Mode.Application && (
+            <UsagesButton onClick={jumpToUsage}>
+              {method.usages.length}
+            </UsagesButton>
+          )}
+          <ViewLink onClick={jumpToUsage}>View</ViewLink>
+          <MethodClassifications method={method} />
+        </ApiOrMethodRow>
+      </VSCodeDataGridCell>
       <VSCodeDataGridCell gridColumn="span 4">
         Method already modeled
       </VSCodeDataGridCell>
@@ -218,3 +246,17 @@ function sendJumpToUsageMessage(method: Method) {
     usage: method.usages[0],
   });
 }
+
+function forEachModeledMethod(
+  modeledMethods: ModeledMethod[],
+  renderer: (
+    modeledMethod: ModeledMethod | undefined,
+    index: number,
+  ) => JSX.Element,
+): JSX.Element | JSX.Element[] {
+  if (modeledMethods.length === 0) {
+    return renderer(undefined, 0);
+  } else {
+    return modeledMethods.map(renderer);
+  }
+}
diff --git a/extensions/ql-vscode/src/view/model-editor/ModeledMethodDataGrid.tsx b/extensions/ql-vscode/src/view/model-editor/ModeledMethodDataGrid.tsx
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
diff --git a/extensions/ql-vscode/src/view/model-editor/__tests__/ModeledMethodDataGrid.spec.tsx b/extensions/ql-vscode/src/view/model-editor/__tests__/ModeledMethodDataGrid.spec.tsx
