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

Add ability for MethodRow to render multiple modelings of the same method
2023-10-09 16:17:17 +01:00
parent 623df4c2ee edd2aa5e8f
commit 27c4bd8349
9 changed files with 240 additions and 95 deletions
--- 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"));
--- 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";
    }
  }
--- 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,
 };
--- 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],
  );

--- 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}
--- 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,19 +96,23 @@ 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
@@ -108,18 +120,20 @@ const ModelableMethodRow = forwardRef<HTMLElement | undefined, MethodRowProps>(
        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);
+  }
+}
--- a/extensions/ql-vscode/src/view/model-editor/ModeledMethodDataGrid.tsx
+++ b/extensions/ql-vscode/src/view/model-editor/ModeledMethodDataGrid.tsx
@@ -8,10 +8,10 @@ import { MethodRow } from "./MethodRow";
 import { Method } from "../../model-editor/method";
 import { ModeledMethod } from "../../model-editor/modeled-method";
 import { useMemo } from "react";
-import { Mode } from "../../model-editor/shared/mode";
 import { sortMethods } from "../../model-editor/shared/sorting";
 import { InProgressMethods } from "../../model-editor/shared/in-progress-methods";
 import { HiddenMethodsRow } from "./HiddenMethodsRow";
+import { ModelEditorViewState } from "../../model-editor/shared/view-state";

 export const GRID_TEMPLATE_COLUMNS = "0.5fr 0.125fr 0.125fr 0.125fr 0.125fr";

@@ -21,7 +21,7 @@ export type ModeledMethodDataGridProps = {
  modeledMethods: Record<string, ModeledMethod>;
  modifiedSignatures: Set<string>;
  inProgressMethods: InProgressMethods;
-  mode: Mode;
+  viewState: ModelEditorViewState;
  hideModeledMethods: boolean;
  revealedMethodSignature: string | null;
  onChange: (modeledMethod: ModeledMethod) => void;
@@ -33,7 +33,7 @@ export const ModeledMethodDataGrid = ({
  modeledMethods,
  modifiedSignatures,
  inProgressMethods,
-  mode,
+  viewState,
  hideModeledMethods,
  revealedMethodSignature,
  onChange,
@@ -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,
-              )}
-              mode={mode}
-              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
--- a/extensions/ql-vscode/src/view/model-editor/tests/MethodRow.spec.tsx
+++ b/extensions/ql-vscode/src/view/model-editor/tests/MethodRow.spec.tsx
@@ -9,6 +9,8 @@ import { Mode } from "../../../model-editor/shared/mode";
 import { MethodRow, MethodRowProps } from "../MethodRow";
 import { ModeledMethod } from "../../../model-editor/modeled-method";
 import userEvent from "@testing-library/user-event";
+import { ModelEditorViewState } from "../../../model-editor/shared/view-state";
+import { createMockExtensionPack } from "../../../../test/factories/model-editor/extension-pack";

 describe(MethodRow.name, () => {
  const method = createMethod({
@@ -31,16 +33,24 @@ describe(MethodRow.name, () => {
  };
  const onChange = jest.fn();

+  const viewState: ModelEditorViewState = {
+    mode: Mode.Application,
+    showFlowGeneration: false,
+    showLlmButton: false,
+    showMultipleModels: false,
+    extensionPack: createMockExtensionPack(),
+  };
+
  const render = (props: Partial<MethodRowProps> = {}) =>
    reactRender(
      <MethodRow
        method={method}
        methodCanBeModeled={true}
-        modeledMethod={modeledMethod}
+        modeledMethods={[modeledMethod]}
        methodIsUnsaved={false}
        modelingInProgress={false}
        revealedMethodSignature={null}
-        mode={Mode.Application}
+        viewState={viewState}
        onChange={onChange}
        {...props}
      />,
@@ -54,6 +64,14 @@ describe(MethodRow.name, () => {
    expect(screen.queryByLabelText("Loading")).not.toBeInTheDocument();
  });

+  it("renders when there is no modeled method", () => {
+    render({ modeledMethods: [] });
+
+    expect(screen.queryAllByRole("combobox")).toHaveLength(4);
+    expect(screen.getByLabelText("Method not modeled")).toBeInTheDocument();
+    expect(screen.queryByLabelText("Loading")).not.toBeInTheDocument();
+  });
+
  it("can change the kind", async () => {
    render();

@@ -110,7 +128,7 @@ describe(MethodRow.name, () => {

  it("shows the modeling status indicator when unmodeled", () => {
    render({
-      modeledMethod: undefined,
+      modeledMethods: [],
    });

    expect(screen.getByLabelText("Method not modeled")).toBeInTheDocument();
@@ -124,10 +142,48 @@ describe(MethodRow.name, () => {
    expect(screen.getByLabelText("Loading")).toBeInTheDocument();
  });

+  it("can render multiple models", () => {
+    render({
+      modeledMethods: [
+        { ...modeledMethod, type: "source" },
+        { ...modeledMethod, type: "sink" },
+        { ...modeledMethod, type: "summary" },
+      ],
+      viewState: {
+        ...viewState,
+        showMultipleModels: true,
+      },
+    });
+
+    const kindInputs = screen.getAllByRole("combobox", { name: "Model type" });
+    expect(kindInputs).toHaveLength(3);
+    expect(kindInputs[0]).toHaveValue("source");
+    expect(kindInputs[1]).toHaveValue("sink");
+    expect(kindInputs[2]).toHaveValue("summary");
+  });
+
+  it("renders only first model when showMultipleModels feature flag is disabled", () => {
+    render({
+      modeledMethods: [
+        { ...modeledMethod, type: "source" },
+        { ...modeledMethod, type: "sink" },
+        { ...modeledMethod, type: "summary" },
+      ],
+      viewState: {
+        ...viewState,
+        showMultipleModels: false,
+      },
+    });
+
+    const kindInputs = screen.getAllByRole("combobox", { name: "Model type" });
+    expect(kindInputs.length).toBe(1);
+    expect(kindInputs[0]).toHaveValue("source");
+  });
+
  it("renders an unmodelable method", () => {
    render({
      methodCanBeModeled: false,
-      modeledMethod: undefined,
+      modeledMethods: [],
    });

    expect(screen.queryByRole("combobox")).not.toBeInTheDocument();
--- a/extensions/ql-vscode/src/view/model-editor/tests/ModeledMethodDataGrid.spec.tsx
+++ b/extensions/ql-vscode/src/view/model-editor/tests/ModeledMethodDataGrid.spec.tsx
@@ -7,6 +7,8 @@ import {
  ModeledMethodDataGrid,
  ModeledMethodDataGridProps,
 } from "../ModeledMethodDataGrid";
+import { ModelEditorViewState } from "../../../model-editor/shared/view-state";
+import { createMockExtensionPack } from "../../../../test/factories/model-editor/extension-pack";

 describe(ModeledMethodDataGrid.name, () => {
  const method1 = createMethod({
@@ -41,6 +43,14 @@ describe(ModeledMethodDataGrid.name, () => {
  });
  const onChange = jest.fn();

+  const viewState: ModelEditorViewState = {
+    mode: Mode.Application,
+    showFlowGeneration: false,
+    showLlmButton: false,
+    showMultipleModels: false,
+    extensionPack: createMockExtensionPack(),
+  };
+
  const render = (props: Partial<ModeledMethodDataGridProps> = {}) =>
    reactRender(
      <ModeledMethodDataGrid
@@ -58,7 +68,7 @@ describe(ModeledMethodDataGrid.name, () => {
        }}
        modifiedSignatures={new Set([method1.signature])}
        inProgressMethods={new InProgressMethods()}
-        mode={Mode.Application}
+        viewState={viewState}
        hideModeledMethods={false}
        revealedMethodSignature={null}
        onChange={onChange}