From a9cfedc1e0c283f9aa22383f75515aaa03ba86f2 Mon Sep 17 00:00:00 2001
From: Giovanni Gonzaga <giovanni@buildo.io>
Date: Fri, 22 Sep 2023 11:15:58 +0200
Subject: [PATCH] allow virtualizing rows when groupBy is active

---
 .../bento-design-system/src/Table/Table.tsx   | 105 ++++++++----------
 .../stories/Components/Table.stories.tsx      |  17 +++
 2 files changed, 65 insertions(+), 57 deletions(-)

diff --git a/packages/bento-design-system/src/Table/Table.tsx b/packages/bento-design-system/src/Table/Table.tsx
index dfa8c38e7..1724269d2 100644
--- a/packages/bento-design-system/src/Table/Table.tsx
+++ b/packages/bento-design-system/src/Table/Table.tsx
@@ -94,6 +94,11 @@ type Props<
 > = {
   columns: C;
   data: ReadonlyArray<RowType<C>>;
+  groupBy?: C extends ReadonlyArray<SimpleColumnType<string, any, any>>
+    ? C[number]["accessor"]
+    : C extends ReadonlyArray<GroupedColumnType<string, any, any>>
+    ? C[number]["columns"][number]["accessor"]
+    : never;
   noResultsTitle?: LocalizedString;
   noResultsDescription?: LocalizedString;
   noResultsFeedbackSize?: FeedbackProps["size"];
@@ -101,18 +106,8 @@ type Props<
   stickyHeaders?: boolean;
   height?: { custom: string | number };
   onRowPress?: (row: Row<RowType<C>>) => void;
-} & (
-  | {
-      groupBy?: C extends ReadonlyArray<SimpleColumnType<string, any, any>>
-        ? C[number]["accessor"]
-        : C extends ReadonlyArray<GroupedColumnType<string, any, any>>
-        ? C[number]["columns"][number]["accessor"]
-        : never;
-      virtualizeRows?: never;
-    }
-  | { groupBy?: never; virtualizeRows?: boolean | { estimateRowHeight: (index: number) => number } }
-) &
-  SortingProps<C>;
+  virtualizeRows?: boolean | { estimateRowHeight: (index: number) => number };
+} & SortingProps<C>;
 
 /**
  * A component that renders a Table, with sorting capabilities
@@ -360,8 +355,9 @@ export function Table<
       .exhaustive();
   }
 
-  const gridTemplateColumns = flatColumns
-    .filter(({ accessor }) => accessor !== groupBy)
+  const nonGroupedColumns = flatColumns.filter(({ accessor }) => accessor !== groupBy);
+
+  const gridTemplateColumns = nonGroupedColumns
     .map(({ gridWidth = "fit-content" }) => gridWidthStyle(gridWidth))
     .join(" ");
 
@@ -385,50 +381,43 @@ export function Table<
     ));
   }
 
-  const renderedRows = virtualizeRows
-    ? columns
-        .map((_, index) => (
-          <div key={`paddingTop${index}`} style={{ marginTop: virtualPaddingTop }} />
-        ))
-        .concat(
-          virtualRows.map((virtualRow) => {
-            const index = virtualRow.index;
-            const row = rows[index];
-            prepareRow(row);
-            return (
-              <RowContainer key={index} row={row} onPress={onRowPress}>
-                {renderCells(row.cells, index, onRowPress !== undefined)}
-              </RowContainer>
-            );
-          })
-        )
-        .concat(
-          columns.map((_, index) => (
-            <div key={`paddingBottom${index}`} style={{ marginBottom: virtualPaddingBottom }} />
-          ))
-        )
-    : rows.flatMap((row, index) => {
-        if (row.isGrouped) {
-          return [
-            <SectionHeader
-              key={row.groupByVal}
-              label={row.groupByVal}
-              numberOfStickyColumns={stickyLeftColumnsIds.length}
-            />,
-            ...row.leafRows.map((row, index) => {
-              prepareRow(row);
-              return renderCells(row.cells, index, false);
-            }),
-          ];
-        } else {
+  const rowsToRender = virtualizeRows
+    ? virtualRows.map((virtualRow) => [rows[virtualRow.index], virtualRow.index] as const)
+    : rows.map((row, index) => [row, index] as const);
+
+  const paddingTopRow = virtualizeRows
+    ? nonGroupedColumns.map((_, index) => (
+        <div key={`paddingTop${index}`} style={{ marginTop: virtualPaddingTop }} />
+      ))
+    : [];
+  const paddingBottomRow = virtualizeRows
+    ? nonGroupedColumns.map((_, index) => (
+        <div key={`paddingBottom${index}`} style={{ marginBottom: virtualPaddingBottom }} />
+      ))
+    : [];
+
+  const renderedRows = rowsToRender.flatMap(([row, index]) => {
+    if (row.isGrouped) {
+      return [
+        <SectionHeader
+          key={row.groupByVal}
+          label={row.groupByVal}
+          numberOfStickyColumns={stickyLeftColumnsIds.length}
+        />,
+        ...row.leafRows.map((row, index) => {
           prepareRow(row);
-          return (
-            <RowContainer key={index} row={row} onPress={onRowPress}>
-              {renderCells(row.cells, index, onRowPress !== undefined)}
-            </RowContainer>
-          );
-        }
-      });
+          return renderCells(row.cells, index, false);
+        }),
+      ];
+    } else {
+      prepareRow(row);
+      return (
+        <RowContainer key={index} row={row} onPress={onRowPress}>
+          {renderCells(row.cells, index, onRowPress !== undefined)}
+        </RowContainer>
+      );
+    }
+  });
 
   return (
     <Box
@@ -469,7 +458,9 @@ export function Table<
           />
         ))
       )}
+      {paddingTopRow}
       {renderedRows}
+      {paddingBottomRow}
     </Box>
   );
 }
diff --git a/packages/bento-design-system/stories/Components/Table.stories.tsx b/packages/bento-design-system/stories/Components/Table.stories.tsx
index 7ac70ebda..6d92a4a56 100644
--- a/packages/bento-design-system/stories/Components/Table.stories.tsx
+++ b/packages/bento-design-system/stories/Components/Table.stories.tsx
@@ -577,3 +577,20 @@ export const VirtualizedRows = {
     data: repeatToLength(exampleData, 1_000),
   },
 } satisfies Story;
+
+export const VirtualizedRowsGrouped = {
+  args: {
+    columns: [
+      ...exampleColumns,
+      tableColumn.text({
+        headerLabel: "Group",
+        accessor: "group",
+      }),
+    ] as const,
+    groupBy: "group",
+    stickyHeaders: true,
+    height: { custom: 340 },
+    virtualizeRows: true,
+    data: repeatToLength(exampleData, 1_000),
+  },
+} satisfies Story;