Group modeled methods by library in the data extensions editor

extensions/ql-vscode/src/view/data-extensions-editor/DataExtensionsEditor.tsx

@@ -4,23 +4,18 @@ import {
   ShowProgressMessage,
   ToDataExtensionsEditorMessage,
 } from "../../pure/interface-types";
-import {
-  VSCodeButton,
-  VSCodeDataGrid,
-  VSCodeDataGridCell,
-  VSCodeDataGridRow,
-} from "@vscode/webview-ui-toolkit/react";
+import { VSCodeButton } from "@vscode/webview-ui-toolkit/react";
 import styled from "styled-components";
 import { ExternalApiUsage } from "../../data-extensions-editor/external-api-usage";
 import { ModeledMethod } from "../../data-extensions-editor/modeled-method";
-import { MethodRow } from "./MethodRow";
 import { assertNever } from "../../pure/helpers-pure";
 import { vscode } from "../vscode-api";
 import { calculateModeledPercentage } from "./modeled";
 import { LinkIconButton } from "../variant-analysis/LinkIconButton";
 import { basename } from "../common/path";
 import { ViewTitle } from "../common";
 import { DataExtensionEditorViewState } from "../../data-extensions-editor/shared/view-state";
+import { ModeledMethodsList } from "./ModeledMethodsList";
 
 const DataExtensionsEditorContainer = styled.div`
   margin-top: 1rem;
@@ -42,6 +37,12 @@ const EditorContainer = styled.div`
   margin-top: 1rem;
 `;
 
+const ButtonsContainer = styled.div`
+  display: flex;
+  gap: 0.4em;
+  margin-bottom: 1rem;
+`;
+
 type ProgressBarProps = {
   completion: number;
 };
@@ -217,57 +218,24 @@ export function DataExtensionsEditor({
           </DetailsContainer>
 
           <EditorContainer>
-            <VSCodeButton onClick={onApplyClick}>Apply</VSCodeButton>
-            &nbsp;
-            <VSCodeButton onClick={onGenerateClick}>
-              Download and generate
-            </VSCodeButton>
-            {viewState?.showLlmButton && (
-              <>
-                &nbsp;
-                <VSCodeButton onClick={onGenerateFromLlmClick}>
-                  Generate using LLM
-                </VSCodeButton>
-              </>
-            )}
-            <br />
-            <br />
-            <VSCodeDataGrid>
-              <VSCodeDataGridRow rowType="header">
-                <VSCodeDataGridCell cellType="columnheader" gridColumn={1}>
-                  Library
-                </VSCodeDataGridCell>
-                <VSCodeDataGridCell cellType="columnheader" gridColumn={2}>
-                  Type
-                </VSCodeDataGridCell>
-                <VSCodeDataGridCell cellType="columnheader" gridColumn={3}>
-                  Method
-                </VSCodeDataGridCell>
-                <VSCodeDataGridCell cellType="columnheader" gridColumn={4}>
-                  Usages
-                </VSCodeDataGridCell>
-                <VSCodeDataGridCell cellType="columnheader" gridColumn={5}>
-                  Model type
-                </VSCodeDataGridCell>
-                <VSCodeDataGridCell cellType="columnheader" gridColumn={6}>
-                  Input
-                </VSCodeDataGridCell>
-                <VSCodeDataGridCell cellType="columnheader" gridColumn={7}>
-                  Output
-                </VSCodeDataGridCell>
-                <VSCodeDataGridCell cellType="columnheader" gridColumn={8}>
-                  Kind
-                </VSCodeDataGridCell>
-              </VSCodeDataGridRow>
-              {externalApiUsages.map((externalApiUsage) => (
-                <MethodRow
-                  key={externalApiUsage.signature}
-                  externalApiUsage={externalApiUsage}
-                  modeledMethod={modeledMethods[externalApiUsage.signature]}
-                  onChange={onChange}
-                />
-              ))}
-            </VSCodeDataGrid>
+            <ButtonsContainer>
+              <VSCodeButton onClick={onApplyClick}>Apply</VSCodeButton>
+              <VSCodeButton onClick={onGenerateClick}>
+                Download and generate
+              </VSCodeButton>
+              {viewState?.showLlmButton && (
+                <>
+                  <VSCodeButton onClick={onGenerateFromLlmClick}>
+                    Generate using LLM
+                  </VSCodeButton>
+                </>
+              )}
+            </ButtonsContainer>
+            <ModeledMethodsList
+              externalApiUsages={externalApiUsages}
+              modeledMethods={modeledMethods}
+              onChange={onChange}
+            />
           </EditorContainer>
         </>
       )}

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

@@ -149,17 +149,14 @@ export const MethodRow = ({
   return (
     <VSCodeDataGridRow>
       <VSCodeDataGridCell gridColumn={1}>
-        {externalApiUsage.library}
-      </VSCodeDataGridCell>
-      <VSCodeDataGridCell gridColumn={2}>
         <SupportSpan
           supported={externalApiUsage.supported}
           modeled={modeledMethod}
         >
           {externalApiUsage.packageName}.{externalApiUsage.typeName}
         </SupportSpan>
       </VSCodeDataGridCell>
-      <VSCodeDataGridCell gridColumn={3}>
+      <VSCodeDataGridCell gridColumn={2}>
         <SupportSpan
           supported={externalApiUsage.supported}
           modeled={modeledMethod}
@@ -168,12 +165,12 @@ export const MethodRow = ({
           {externalApiUsage.methodParameters}
         </SupportSpan>
       </VSCodeDataGridCell>
-      <VSCodeDataGridCell gridColumn={4}>
+      <VSCodeDataGridCell gridColumn={3}>
         <UsagesButton onClick={jumpToUsage}>
           {externalApiUsage.usages.length}
         </UsagesButton>
       </VSCodeDataGridCell>
-      <VSCodeDataGridCell gridColumn={5}>
+      <VSCodeDataGridCell gridColumn={4}>
         {(!externalApiUsage.supported ||
           (modeledMethod && modeledMethod?.type !== "none")) && (
           <Dropdown
@@ -188,7 +185,7 @@ export const MethodRow = ({
           </Dropdown>
         )}
       </VSCodeDataGridCell>
-      <VSCodeDataGridCell gridColumn={6}>
+      <VSCodeDataGridCell gridColumn={5}>
         {modeledMethod?.type &&
           ["sink", "summary"].includes(modeledMethod?.type) && (
             <Dropdown value={modeledMethod?.input} onInput={handleInputInput}>
@@ -201,7 +198,7 @@ export const MethodRow = ({
             </Dropdown>
           )}
       </VSCodeDataGridCell>
-      <VSCodeDataGridCell gridColumn={7}>
+      <VSCodeDataGridCell gridColumn={6}>
         {modeledMethod?.type &&
           ["source", "summary"].includes(modeledMethod?.type) && (
             <Dropdown value={modeledMethod?.output} onInput={handleOutputInput}>
@@ -215,7 +212,7 @@ export const MethodRow = ({
             </Dropdown>
           )}
       </VSCodeDataGridCell>
-      <VSCodeDataGridCell gridColumn={8}>
+      <VSCodeDataGridCell gridColumn={7}>
         {predicate?.supportedKinds && (
           <KindInput
             kinds={predicate.supportedKinds}

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

@@ -0,0 +1,60 @@
+import * as React from "react";
+import {
+  VSCodeDataGrid,
+  VSCodeDataGridCell,
+  VSCodeDataGridRow,
+} from "@vscode/webview-ui-toolkit/react";
+import { MethodRow } from "./MethodRow";
+import { ExternalApiUsage } from "../../data-extensions-editor/external-api-usage";
+import { ModeledMethod } from "../../data-extensions-editor/modeled-method";
+
+type Props = {
+  externalApiUsages: ExternalApiUsage[];
+  modeledMethods: Record<string, ModeledMethod>;
+  onChange: (
+    externalApiUsage: ExternalApiUsage,
+    modeledMethod: ModeledMethod,
+  ) => void;
+};
+
+export const ModeledMethodDataGrid = ({
+  externalApiUsages,
+  modeledMethods,
+  onChange,
+}: Props) => {
+  return (
+    <VSCodeDataGrid>
+      <VSCodeDataGridRow rowType="header">
+        <VSCodeDataGridCell cellType="columnheader" gridColumn={1}>
+          Type
+        </VSCodeDataGridCell>
+        <VSCodeDataGridCell cellType="columnheader" gridColumn={2}>
+          Method
+        </VSCodeDataGridCell>
+        <VSCodeDataGridCell cellType="columnheader" gridColumn={3}>
+          Usages
+        </VSCodeDataGridCell>
+        <VSCodeDataGridCell cellType="columnheader" gridColumn={4}>
+          Model type
+        </VSCodeDataGridCell>
+        <VSCodeDataGridCell cellType="columnheader" gridColumn={5}>
+          Input
+        </VSCodeDataGridCell>
+        <VSCodeDataGridCell cellType="columnheader" gridColumn={6}>
+          Output
+        </VSCodeDataGridCell>
+        <VSCodeDataGridCell cellType="columnheader" gridColumn={7}>
+          Kind
+        </VSCodeDataGridCell>
+      </VSCodeDataGridRow>
+      {externalApiUsages.map((externalApiUsage) => (
+        <MethodRow
+          key={externalApiUsage.signature}
+          externalApiUsage={externalApiUsage}
+          modeledMethod={modeledMethods[externalApiUsage.signature]}
+          onChange={onChange}
+        />
+      ))}
+    </VSCodeDataGrid>
+  );
+};

extensions/ql-vscode/src/view/data-extensions-editor/ModeledMethodsList.tsx

@@ -0,0 +1,55 @@
+import * as React from "react";
+import { useMemo } from "react";
+import styled from "styled-components";
+import { ExternalApiUsage } from "../../data-extensions-editor/external-api-usage";
+import { ModeledMethod } from "../../data-extensions-editor/modeled-method";
+import { ModeledMethodDataGrid } from "./ModeledMethodDataGrid";
+
+const LibraryContainer = styled.div`
+  margin-bottom: 1rem;
+`;
+
+type Props = {
+  externalApiUsages: ExternalApiUsage[];
+  modeledMethods: Record<string, ModeledMethod>;
+  onChange: (
+    externalApiUsage: ExternalApiUsage,
+    modeledMethod: ModeledMethod,
+  ) => void;
+};
+
+export const ModeledMethodsList = ({
+  externalApiUsages,
+  modeledMethods,
+  onChange,
+}: Props) => {
+  const groupedByLibrary = useMemo(() => {
+    const groupedByLibrary: Record<string, ExternalApiUsage[]> = {};
+
+    for (const externalApiUsage of externalApiUsages) {
+      groupedByLibrary[externalApiUsage.library] ??= [];
+      groupedByLibrary[externalApiUsage.library].push(externalApiUsage);
+    }
+
+    return groupedByLibrary;
+  }, [externalApiUsages]);
+
+  const sortedLibraryNames = useMemo(() => {
+    return Object.keys(groupedByLibrary).sort();
+  }, [groupedByLibrary]);
+
+  return (
+    <>
+      {sortedLibraryNames.map((libraryName) => (
+        <LibraryContainer key={libraryName}>
+          <h3>{libraryName}</h3>
+          <ModeledMethodDataGrid
+            externalApiUsages={groupedByLibrary[libraryName]}
+            modeledMethods={modeledMethods}
+            onChange={onChange}
+          />
+        </LibraryContainer>
+      ))}
+    </>
+  );
+};