fix(Table): fix drag when virtualization enabled (#7)

Co-authored-by: kseniyakuzina <[email protected]>

src/components/Table/Table.tsx

@@ -130,7 +130,9 @@ export const Table = React.memo(
         onColumnSizingChange,
         onColumnSizingInfoChange,
     }: TableProps<TData>) => {
-        const draggable = Boolean(React.useContext(SortableListContext));
+        const draggableContext = React.useContext(SortableListContext);
+        const draggable = Boolean(draggableContext);
+        const draggableItemKey = draggableContext?.activeItemKey;
 
         const virtualizationContext = React.useContext(VirtualizationContext);
         const virtual = Boolean(virtualizationContext);
@@ -212,6 +214,13 @@ export const Table = React.memo(
             tableOptions.onColumnSizingInfoChange = onColumnSizingInfoChange;
         }
 
+        const draggableIndex = React.useMemo(() => {
+            if (!draggableItemKey) {
+                return undefined;
+            }
+            return data.findIndex((item) => getRowId(item) === draggableItemKey);
+        }, [data, draggableItemKey, getRowId]);
+
         const table = useReactTable(tableOptions);
 
         const getRowByIndex = React.useCallback<TableContextProviderProps<TData>['getRowByIndex']>(
@@ -266,7 +275,7 @@ export const Table = React.memo(
                 enableNesting={enableNesting}
                 getTableState={table.getState}
             >
-                <table className={b(null, className)}>
+                <table className={b(null, className)} data-draggable-index={draggableIndex}>
                     {withHeader && (
                         <thead className={b('header', headerClassName)}>
                             {headerGroups.map((headerGroup, index) => (
@@ -285,7 +294,7 @@ export const Table = React.memo(
                         </thead>
                     )}
                     <tbody
-                        className={bodyClassName}
+                        className={b('body', bodyClassName)}
                         style={{
                             height: totalSize,
                         }}

src/components/VirtualRow/VirtualRow.tsx

@@ -20,7 +20,7 @@ export const VirtualRow = React.forwardRef(
 
         const virtualRowStyle = React.useMemo(
             () => ({
-                transform: `translateY(${virtualRow.start}px)`,
+                top: virtualRow.start,
                 ...style,
             }),
             [style, virtualRow.start],

src/components/VirtualizationProvider/VirtualizationProvider.tsx

@@ -4,6 +4,7 @@ import {useVirtualizer} from '@tanstack/react-virtual';
 
 import {VirtualizationContainer} from '../VirtualizationContainer';
 import {VirtualizationContext} from '../VirtualizationContext';
+import {getVirtualRangeExtractor} from '../utils/getVirtualRangeExtractor';
 
 export interface VirtualizationProviderProps {
     children?: React.ReactNode;
@@ -30,6 +31,7 @@ export const VirtualizationProvider = ({
         getScrollElement: () => containerRef.current,
         measureElement: (element) => element?.getBoundingClientRect().height,
         overscan: overscanRowCount,
+        rangeExtractor: getVirtualRangeExtractor(containerRef.current),
     });
 
     const virtualItems = rowVirtualizer.getVirtualItems();

src/components/WindowVirtualizationProvider/WindowVirtualizationProvider.tsx

@@ -4,6 +4,7 @@ import {useWindowVirtualizer} from '@tanstack/react-virtual';
 
 import {VirtualizationContainer} from '../VirtualizationContainer';
 import {VirtualizationContext} from '../VirtualizationContext';
+import {getVirtualRangeExtractor} from '../utils/getVirtualRangeExtractor';
 
 export interface WindowVirtualizationProviderProps {
     rowsCount: number;
@@ -20,11 +21,14 @@ export const WindowVirtualizationProvider = ({
     containerClassName,
     estimateRowSize,
 }: WindowVirtualizationProviderProps) => {
+    const containerRef = React.useRef<HTMLDivElement>(null);
+
     const rowVirtualizer = useWindowVirtualizer({
         count: rowsCount,
         estimateSize: estimateRowSize,
         measureElement: (element) => element?.getBoundingClientRect().height,
         overscan: overscanRowCount,
+        rangeExtractor: getVirtualRangeExtractor(containerRef.current),
     });
 
     const virtualItems = rowVirtualizer.getVirtualItems();
@@ -41,7 +45,7 @@ export const WindowVirtualizationProvider = ({
 
     return (
         <VirtualizationContext.Provider value={contextValue}>
-            <VirtualizationContainer className={containerClassName}>
+            <VirtualizationContainer ref={containerRef} className={containerClassName}>
                 {children}
             </VirtualizationContainer>
         </VirtualizationContext.Provider>

src/components/__stories__/Grid2.stories.tsx

@@ -10,6 +10,7 @@ import {GroupingDemo2} from './GroupingDemo2';
 import {GroupingWithSelectionDemo} from './GroupingWithSelectionDemo';
 import {ReorderingDemo} from './ReorderingDemo';
 import {ReorderingTreeDemo} from './ReorderingTreeDemo';
+import {ReorderingWithVirtualizationDemo} from './ReorderingWithVirtualizationDemo';
 import {SortingDemo} from './SortingDemo';
 import {TreeDemo} from './TreeDemo';
 import {TreeWithGroupsDemo} from './TreeWithGroupsDemo';
@@ -185,6 +186,21 @@ WindowVirtualization.args = {
     onSelectedChange: undefined,
 };
 
+const ReorderingWithVirtualizationTemplate: StoryFn<typeof Table<Item>> = (args) => (
+    <ReorderingWithVirtualizationDemo {...args} />
+);
+
+export const ReorderingWithVirtualization: StoryFn<typeof Table<Item>> =
+    ReorderingWithVirtualizationTemplate.bind({});
+
+ReorderingWithVirtualization.args = {
+    data: generateData(300),
+    columns,
+    getRowId: (item: Item) => item.id,
+    onRowClick: undefined,
+    onSelectedChange: undefined,
+};
+
 export const Resizing: StoryFn<typeof Table<Item>> = Template.bind({});
 
 Resizing.args = {

src/components/__stories__/ReorderingWithVirtualizationDemo.tsx

@@ -0,0 +1,55 @@
+import React from 'react';
+
+import type {TableProps} from '../../components/Table';
+import {Table} from '../../components/Table';
+import {withTableReorder} from '../../hocs';
+import type {SortableListDragResult} from '../SortableList';
+import {WindowVirtualizationProvider} from '../WindowVirtualizationProvider';
+
+const GridWithReordering = withTableReorder(Table);
+
+const rowHeight = 20;
+const estimateRowSize = () => rowHeight;
+
+export const ReorderingWithVirtualizationDemo = <ItemType extends unknown>(
+    props: TableProps<ItemType>,
+) => {
+    const {getRowId} = props;
+
+    const [data, setData] = React.useState(props.data);
+
+    const handleDragEnd = React.useCallback(
+        ({draggedItemKey, baseItemKey}: SortableListDragResult) => {
+            setData((data) => {
+                const dataClone = data.slice();
+
+                const index = dataClone.findIndex((item) => getRowId(item) === draggedItemKey);
+                if (index >= 0) {
+                    const dragged = dataClone.splice(index, 1)[0]!;
+
+                    const insertIndex = dataClone.findIndex(
+                        (value) => getRowId(value) === baseItemKey,
+                    );
+                    if (insertIndex >= 0) {
+                        dataClone.splice(insertIndex + 1, 0, dragged);
+                    } else {
+                        dataClone.unshift(dragged);
+                    }
+                }
+
+                return dataClone;
+            });
+        },
+        [getRowId],
+    );
+
+    return (
+        <WindowVirtualizationProvider
+            estimateRowSize={estimateRowSize}
+            overscanRowCount={5}
+            rowsCount={data.length}
+        >
+            <GridWithReordering<ItemType> {...props} data={data} onReorder={handleDragEnd} />
+        </WindowVirtualizationProvider>
+    );
+};

src/components/utils/getVirtualRangeExtractor.ts

@@ -0,0 +1,17 @@
+import {defaultRangeExtractor} from '@tanstack/react-virtual';
+import type {Range} from '@tanstack/react-virtual';
+
+export function getVirtualRangeExtractor(container?: HTMLElement | null) {
+    return function (range: Range) {
+        const table = container?.querySelector('[data-draggable-index]');
+        const draggableItemIndex = Number(table?.getAttribute('data-draggable-index') ?? -1);
+
+        const result = defaultRangeExtractor(range);
+
+        if (draggableItemIndex !== -1 && !result.includes(draggableItemIndex)) {
+            result.push(draggableItemIndex);
+        }
+
+        return result;
+    };
+}