nextui/packages/components/table/src/virtualized-table.tsx

import type {UseTableProps} from "./use-table";

import {useCallback, useLayoutEffect, useRef, useState} from "react";
import {forwardRef} from "@heroui/system";
import {useVirtualizer} from "@tanstack/react-virtual";

import {useTable} from "./use-table";
import TableRowGroup from "./table-row-group";
import TableHeaderRow from "./table-header-row";
import TableColumnHeader from "./table-column-header";
import TableSelectAllCheckbox from "./table-select-all-checkbox";
import VirtualizedTableBody from "./virtualized-table-body";

export interface TableProps<T = object>
  extends Omit<UseTableProps<T>, "isSelectable" | "isMultiSelectable"> {
  isVirtualized?: boolean;
  rowHeight?: number;
  maxTableHeight?: number;
}

const VirtualizedTable = forwardRef<"table", TableProps>((props, ref) => {
  const {
    BaseComponent,
    Component,
    collection,
    values,
    topContent,
    topContentPlacement,
    bottomContentPlacement,
    bottomContent,
    getBaseProps,
    getWrapperProps,
    getTableProps,
  } = useTable({
    ...props,
    ref,
  });

  const {rowHeight = 40, maxTableHeight = 600} = props;

  const Wrapper = useCallback(
    ({children}: {children: JSX.Element}) => {
      return (
        <BaseComponent
          {...getWrapperProps()}
          ref={parentRef}
          /* Display must be block to maintain the scroll "progress" */
          style={{height: maxTableHeight, display: "block"}}
        >
          {children}
        </BaseComponent>
      );
    },
    [getWrapperProps, maxTableHeight],
  );

  const items = [...collection.body.childNodes];

  const count = items.length;

  const parentRef = useRef(null);

  const [headerHeight, setHeaderHeight] = useState(0);

  const headerRef = useRef<HTMLTableSectionElement>(null);

  useLayoutEffect(() => {
    if (headerRef.current) {
      setHeaderHeight(headerRef.current.getBoundingClientRect().height);
    }
  }, [headerRef]);

  const rowVirtualizer = useVirtualizer({
    count,
    getScrollElement: () => parentRef.current,
    estimateSize: () => rowHeight,
    overscan: 5,
  });

  const tableProps = getTableProps();

  return (
    <div {...getBaseProps()}>
      {/* We need to add p-1 to make the shadow-sm visible */}
      {topContentPlacement === "outside" && topContent}
      <Wrapper>
        <>
          {topContentPlacement === "inside" && topContent}
          <Component
            {...tableProps}
            style={{
              height: `calc(${rowVirtualizer.getTotalSize() + headerHeight}px)`,
              ...tableProps.style,
            }}
          >
            <TableRowGroup ref={headerRef} classNames={values.classNames} slots={values.slots}>
              {collection.headerRows.map((headerRow) => (
                <TableHeaderRow
                  key={headerRow?.key}
                  classNames={values.classNames}
                  node={headerRow}
                  slots={values.slots}
                  state={values.state}
                >
                  {[...headerRow.childNodes].map((column) =>
                    column?.props?.isSelectionCell ? (
                      <TableSelectAllCheckbox
                        key={column?.key}
                        checkboxesProps={values.checkboxesProps}
                        classNames={values.classNames}
                        color={values.color}
                        disableAnimation={values.disableAnimation}
                        node={column}
                        selectionMode={values.selectionMode}
                        slots={values.slots}
                        state={values.state}
                      />
                    ) : (
                      <TableColumnHeader
                        key={column?.key}
                        classNames={values.classNames}
                        node={column}
                        slots={values.slots}
                        state={values.state}
                      />
                    ),
                  )}
                </TableHeaderRow>
              ))}
            </TableRowGroup>
            <VirtualizedTableBody
              checkboxesProps={values.checkboxesProps}
              classNames={values.classNames}
              collection={values.collection}
              color={values.color}
              disableAnimation={values.disableAnimation}
              isSelectable={values.isSelectable}
              rowVirtualizer={rowVirtualizer}
              selectionMode={values.selectionMode}
              slots={values.slots}
              state={values.state}
            />
          </Component>
          {bottomContentPlacement === "inside" && bottomContent}
        </>
      </Wrapper>
      {bottomContentPlacement === "outside" && bottomContent}
    </div>
  );
});

VirtualizedTable.displayName = "HeroUI.VirtualizedTable";

export default VirtualizedTable;