Fix virtualizer persisted keys with drag and drop

packages/@react-aria/dnd/src/DragManager.ts

@@ -13,7 +13,6 @@
 import {announce} from '@react-aria/live-announcer';
 import {ariaHideOutside} from '@react-aria/overlays';
 import {DragEndEvent, DragItem, DropActivateEvent, DropEnterEvent, DropEvent, DropExitEvent, DropItem, DropOperation, DropTarget as DroppableCollectionTarget, FocusableElement} from '@react-types/shared';
-import {flushSync} from 'react-dom';
 import {getDragModality, getTypes} from './utils';
 import {isVirtualClick, isVirtualPointerEvent} from '@react-aria/utils';
 import type {LocalizedStringFormatter} from '@internationalized/string';
@@ -26,6 +25,7 @@ let subscriptions = new Set<() => void>();
 
 interface DropTarget {
   element: FocusableElement,
+  preventFocusOnDrop?: boolean,
   getDropOperation?: (types: Set<string>, allowedOperations: DropOperation[]) => DropOperation,
   onDropEnter?: (e: DropEnterEvent, dragTarget: DragTarget) => void,
   onDropExit?: (e: DropExitEvent) => void,
@@ -513,19 +513,11 @@ class DragSession {
       });
     }
 
-    // Blur and re-focus the drop target so that the focus ring appears.
-    if (this.currentDropTarget) {
-      // Since we cancel all focus events in drag sessions, refire blur to make sure state gets updated so drag target doesn't think it's still focused
-      // i.e. When you from one list to another during a drag session, we need the blur to fire on the first list after the drag.
-      if (!this.dragTarget.element.contains(this.currentDropTarget.element)) {
-        this.dragTarget.element.dispatchEvent(new FocusEvent('blur'));
-        this.dragTarget.element.dispatchEvent(new FocusEvent('focusout', {bubbles: true}));
-      }
-      // Re-focus the focusedKey upon reorder. This requires a React rerender between blurring and focusing.
-      flushSync(() => {
-        this.currentDropTarget.element.blur();
-      });
-      this.currentDropTarget.element.focus();
+    if (this.currentDropTarget && !this.currentDropTarget.preventFocusOnDrop) {
+      // Re-trigger focus event on active element, since it will not have received it during dragging (see cancelEvent).
+      // This corrects state such as whether focus ring should appear.
+      // useDroppableCollection handles this itself, so this is only for standalone drop zones.
+      document.activeElement.dispatchEvent(new FocusEvent('focusin', {bubbles: true}));
     }
 
     this.setCurrentDropTarget(null);

packages/@react-aria/dnd/src/useDroppableCollection.ts

@@ -57,6 +57,9 @@ interface DroppingState {
   collection: Collection<Node<unknown>>,
   focusedKey: Key,
   selectedKeys: Set<Key>,
+  target: DropTarget,
+  draggingKeys: Set<Key>,
+  isInternal: boolean,
   timeout: ReturnType<typeof setTimeout>
 }
 
@@ -213,26 +216,93 @@ export function useDroppableCollection(props: DroppableCollectionOptions, state:
   });
 
   let droppingState = useRef<DroppingState>(null);
-  let onDrop = useCallback((e: DropEvent, target: DropTarget) => {
+  let updateFocusAfterDrop = useCallback(() => {
     let {state} = localState;
+    if (droppingState.current) {
+      let {
+        target,
+        collection: prevCollection,
+        selectedKeys: prevSelectedKeys,
+        focusedKey: prevFocusedKey,
+        isInternal,
+        draggingKeys
+      } = droppingState.current;
+
+      // If an insert occurs during a drop, we want to immediately select these items to give
+      // feedback to the user that a drop occurred. Only do this if the selection didn't change
+      // since the drop started so we don't override if the user or application did something.
+      if (
+        state.collection.size > prevCollection.size &&
+        state.selectionManager.isSelectionEqual(prevSelectedKeys)
+      ) {
+        let newKeys = new Set<Key>();
+        for (let key of state.collection.getKeys()) {
+          if (!prevCollection.getItem(key)) {
+            newKeys.add(key);
+          }
+        }
 
-    // Focus the collection.
-    state.selectionManager.setFocused(true);
+        state.selectionManager.setSelectedKeys(newKeys);
 
-    // Save some state of the collection/selection before the drop occurs so we can compare later.
-    let focusedKey = state.selectionManager.focusedKey;
+        // If the focused item didn't change since the drop occurred, also focus the first
+        // inserted item. If selection is disabled, then also show the focus ring so there
+        // is some indication that items were added.
+        if (state.selectionManager.focusedKey === prevFocusedKey) {
+          let first = newKeys.keys().next().value;
+          let item = state.collection.getItem(first);
+
+          // If this is a cell, focus the parent row.
+          if (item?.type === 'cell') {
+            first = item.parentKey;
+          }
+
+          state.selectionManager.setFocusedKey(first);
 
-    // If parent key was dragged, we want to use it instead (i.e. focus row instead of cell after dropping)
-    if (globalDndState.draggingKeys.has(state.collection.getItem(focusedKey)?.parentKey)) {
-      focusedKey = state.collection.getItem(focusedKey).parentKey;
-      state.selectionManager.setFocusedKey(focusedKey);
+          if (state.selectionManager.selectionMode === 'none') {
+            setInteractionModality('keyboard');
+          }
+        }
+      } else if (
+        state.selectionManager.focusedKey === prevFocusedKey &&
+        isInternal &&
+        target.type === 'item' &&
+        target.dropPosition !== 'on' &&
+        draggingKeys.has(state.collection.getItem(prevFocusedKey)?.parentKey)
+      ) {
+        // Focus row instead of cell when reordering.
+        state.selectionManager.setFocusedKey(state.collection.getItem(prevFocusedKey).parentKey);
+        setInteractionModality('keyboard');
+      } else if (
+        state.selectionManager.focusedKey === prevFocusedKey &&
+        target.type === 'item' && 
+        target.dropPosition === 'on' && 
+        state.collection.getItem(target.key) != null
+      ) {
+        // If focus didn't move already (e.g. due to an insert), and the user dropped on an item,
+        // focus that item and show the focus ring to give the user feedback that the drop occurred.
+        // Also show the focus ring if the focused key is not selected, e.g. in case of a reorder.
+        state.selectionManager.setFocusedKey(target.key);
+        setInteractionModality('keyboard');
+      } else if (!state.selectionManager.isSelected(state.selectionManager.focusedKey)) {
+        setInteractionModality('keyboard');
+      }
+
+      state.selectionManager.setFocused(true);
     }
+  }, [localState]);
 
+  let onDrop = useCallback((e: DropEvent, target: DropTarget) => {
+    let {state} = localState;
+
+    // Save some state of the collection/selection before the drop occurs so we can compare later.
     droppingState.current = {
       timeout: null,
-      focusedKey,
+      focusedKey: state.selectionManager.focusedKey,
       collection: state.collection,
-      selectedKeys: state.selectionManager.selectedKeys
+      selectedKeys: state.selectionManager.selectedKeys,
+      draggingKeys: globalDndState.draggingKeys,
+      isInternal: isInternalDropOperation(ref),
+      target
     };
 
     let onDropFn = localState.props.onDrop || defaultOnDrop;
@@ -246,26 +316,13 @@ export function useDroppableCollection(props: DroppableCollectionOptions, state:
     });
 
     // Wait for a short time period after the onDrop is called to allow the data to be read asynchronously
-    // and for React to re-render. If an insert occurs during this time, it will be selected/focused below.
-    // If items are not "immediately" inserted by the onDrop handler, the application will need to handle
-    // selecting and focusing those items themselves.
+    // and for React to re-render. If the collection didn't already change during this time (handled below),
+    // update the focused key here.
     droppingState.current.timeout = setTimeout(() => {
-      // If focus didn't move already (e.g. due to an insert), and the user dropped on an item,
-      // focus that item and show the focus ring to give the user feedback that the drop occurred.
-      // Also show the focus ring if the focused key is not selected, e.g. in case of a reorder.
-      let {state} = localState;
-
-      if (target.type === 'item' && target.dropPosition === 'on' && state.collection.getItem(target.key) != null) {
-        state.selectionManager.setFocusedKey(target.key);
-        state.selectionManager.setFocused(true);
-        setInteractionModality('keyboard');
-      } else if (!state.selectionManager.isSelected(focusedKey)) {
-        setInteractionModality('keyboard');
-      }
-
+      updateFocusAfterDrop();
       droppingState.current = null;
     }, 50);
-  }, [localState, defaultOnDrop]);
+  }, [localState, defaultOnDrop, ref, updateFocusAfterDrop]);
 
   // eslint-disable-next-line arrow-body-style
   useEffect(() => {
@@ -277,44 +334,9 @@ export function useDroppableCollection(props: DroppableCollectionOptions, state:
   }, []);
 
   useLayoutEffect(() => {
-    // If an insert occurs during a drop, we want to immediately select these items to give
-    // feedback to the user that a drop occurred. Only do this if the selection didn't change
-    // since the drop started so we don't override if the user or application did something.
-    if (
-      droppingState.current &&
-      state.selectionManager.isFocused &&
-      state.collection.size > droppingState.current.collection.size &&
-      state.selectionManager.isSelectionEqual(droppingState.current.selectedKeys)
-    ) {
-      let newKeys = new Set<Key>();
-      for (let key of state.collection.getKeys()) {
-        if (!droppingState.current.collection.getItem(key)) {
-          newKeys.add(key);
-        }
-      }
-
-      state.selectionManager.setSelectedKeys(newKeys);
-
-      // If the focused item didn't change since the drop occurred, also focus the first
-      // inserted item. If selection is disabled, then also show the focus ring so there
-      // is some indication that items were added.
-      if (state.selectionManager.focusedKey === droppingState.current.focusedKey) {
-        let first = newKeys.keys().next().value;
-        let item = state.collection.getItem(first);
-
-        // If this is a cell, focus the parent row.
-        if (item?.type === 'cell') {
-          first = item.parentKey;
-        }
-
-        state.selectionManager.setFocusedKey(first);
-
-        if (state.selectionManager.selectionMode === 'none') {
-          setInteractionModality('keyboard');
-        }
-      }
-
-      droppingState.current = null;
+    // If the collection changed after a drop, update the focused key.
+    if (droppingState.current && state.collection !== droppingState.current.collection) {
+      updateFocusAfterDrop();
     }
   });
 
@@ -470,6 +492,7 @@ export function useDroppableCollection(props: DroppableCollectionOptions, state:
 
     return DragManager.registerDropTarget({
       element: ref.current,
+      preventFocusOnDrop: true,
       getDropOperation(types, allowedOperations) {
         if (localState.state.target) {
           let {draggingKeys} = globalDndState;

packages/@react-aria/dnd/stories/DraggableCollection.tsx

@@ -58,7 +58,7 @@ function DraggableCollection(props) {
   let state = useListState(props);
   let gridState = useGridState({
     selectionMode: 'multiple',
-    collection: new GridCollection({
+    collection: React.useMemo(() => new GridCollection({
       columnCount: 1,
       items: [...state.collection].map(item => ({
         ...item,
@@ -74,7 +74,7 @@ function DraggableCollection(props) {
           childNodes: []
         }]
       }))
-    })
+    }), [state.collection])
   });
 
   let preview = useRef(null);

packages/@react-aria/dnd/stories/DroppableGrid.tsx

@@ -122,7 +122,7 @@ const DroppableGrid = React.forwardRef(function (props: any, ref) {
     focusMode: 'cell',
     selectedKeys: props.selectedKeys,
     onSelectionChange: props.onSelectionChange,
-    collection: new GridCollection({
+    collection: React.useMemo(() => new GridCollection({
       columnCount: 1,
       items: [...state.collection].map(item => ({
         ...item,
@@ -138,7 +138,7 @@ const DroppableGrid = React.forwardRef(function (props: any, ref) {
           childNodes: []
         }]
       }))
-    })
+    }), [state.collection])
   });
 
   React.useImperativeHandle(ref, () => ({

packages/@react-aria/dnd/stories/Reorderable.tsx

@@ -71,7 +71,7 @@ function ReorderableGrid(props) {
   let keyboardDelegate = new ListKeyboardDelegate(state.collection, new Set(), ref);
   let gridState = useGridState({
     selectionMode: 'multiple',
-    collection: new GridCollection({
+    collection: React.useMemo(() => new GridCollection({
       columnCount: 1,
       items: [...state.collection].map(item => ({
         ...item,
@@ -87,7 +87,7 @@ function ReorderableGrid(props) {
           childNodes: []
         }]
       }))
-    })
+    }), [state.collection])
   });
 
   // Use a random drag type so the items can only be reordered within this list and not dragged elsewhere.

packages/@react-aria/dnd/stories/VirtualizedListBox.tsx

@@ -21,7 +21,7 @@ import Folder from '@spectrum-icons/workflow/Folder';
 import {Item} from '@react-stately/collections';
 import {ListKeyboardDelegate} from '@react-aria/selection';
 import {ListLayout} from '@react-stately/layout';
-import React from 'react';
+import React, {useMemo} from 'react';
 import {useDropIndicator, useDroppableCollection, useDroppableItem} from '..';
 import {useDroppableCollectionState} from '@react-stately/dnd';
 import {useListBox, useOption} from '@react-aria/listbox';
@@ -159,6 +159,8 @@ export const VirtualizedListBox = React.forwardRef(function (props: any, ref) {
     isVirtualized: true
   }, state, domRef);
   let isDropTarget = dropState.isDropTarget({type: 'root'});
+  let focusedKey = dropState.target?.type === 'item' ? dropState.target.key : state.selectionManager.focusedKey;
+  let persistedKeys = useMemo(() => focusedKey != null ? new Set([focusedKey]) : null, [focusedKey]);
 
   return (
     <Context.Provider value={{state, dropState}}>
@@ -170,7 +172,7 @@ export const VirtualizedListBox = React.forwardRef(function (props: any, ref) {
         scrollDirection="vertical"
         layout={layout}
         collection={state.collection}
-        focusedKey={dropState.target?.type === 'item' ? dropState.target.key : state.selectionManager.focusedKey}>
+        persistedKeys={persistedKeys}>
         {(type, item) => (
           <>
             {state.collection.getKeyBefore(item.key) == null &&

packages/@react-aria/dnd/stories/dnd.stories.tsx

@@ -386,7 +386,7 @@ function DraggableCollection(props) {
   let gridState = useGridState({
     ...props,
     selectionMode: 'multiple',
-    collection: new GridCollection({
+    collection: React.useMemo(() => new GridCollection({
       columnCount: 1,
       items: [...state.collection].map(item => ({
         ...item,
@@ -402,7 +402,7 @@ function DraggableCollection(props) {
           childNodes: []
         }]
       }))
-    })
+    }), [state.collection])
   });
 
   let preview = useRef(null);

packages/@react-aria/grid/stories/example.tsx

@@ -11,7 +11,7 @@ export function Grid(props) {
   let gridState = useGridState({
     ...props,
     selectionMode: 'multiple',
-    collection: new GridCollection({
+    collection: React.useMemo(() => new GridCollection({
       columnCount: 1,
       items: [...state.collection].map(item => ({
         type: 'item',
@@ -21,7 +21,7 @@ export function Grid(props) {
           type: 'cell'
         }]
       }))
-    })
+    }), [state.collection])
   });
 
   let ref = React.useRef(undefined);

packages/@react-aria/virtualizer/src/Virtualizer.tsx

@@ -13,7 +13,7 @@
 import {Collection, Key} from '@react-types/shared';
 import {Layout, Rect, ReusableView, useVirtualizerState, VirtualizerState} from '@react-stately/virtualizer';
 import {mergeProps, useLayoutEffect} from '@react-aria/utils';
-import React, {HTMLAttributes, ReactElement, ReactNode, RefObject, useCallback, useMemo, useRef} from 'react';
+import React, {HTMLAttributes, ReactElement, ReactNode, RefObject, useCallback, useRef} from 'react';
 import {ScrollView} from './ScrollView';
 import {VirtualizerItem} from './VirtualizerItem';
 
@@ -29,7 +29,7 @@ interface VirtualizerProps<T extends object, V, O> extends Omit<HTMLAttributes<H
   renderWrapper?: RenderWrapper<T, V>,
   layout: Layout<T, O>,
   collection: Collection<T>,
-  focusedKey?: Key,
+  persistedKeys?: Set<Key> | null,
   sizeToFit?: 'width' | 'height',
   scrollDirection?: 'horizontal' | 'vertical' | 'both',
   isLoading?: boolean,
@@ -49,7 +49,7 @@ function Virtualizer<T extends object, V extends ReactNode, O>(props: Virtualize
     isLoading,
     // eslint-disable-next-line @typescript-eslint/no-unused-vars
     onLoadMore,
-    focusedKey,
+    persistedKeys,
     layoutOptions,
     ...otherProps
   } = props;
@@ -65,7 +65,7 @@ function Virtualizer<T extends object, V extends ReactNode, O>(props: Virtualize
       ref.current.scrollLeft = rect.x;
       ref.current.scrollTop = rect.y;
     },
-    persistedKeys: useMemo(() => focusedKey != null ? new Set([focusedKey]) : new Set(), [focusedKey]),
+    persistedKeys,
     layoutOptions
   });