feat(RAC): Tree drag and drop (#7692)

* [5574] - add moveBefore and moveAfter to useTreeData

* add docs

* remove onlys

* remove console logs

* [wip] add intial tree drag and drop

* useTreeData - add getDescendantKeys method which is used to determine if a parent node can be dropped into its children

* revert packlog json change

* feat: Add moveBefore and moveAfter to useTreeData

* updates + story

* remove RSP TreeView dnd for now

* cleanup

* fix story

* lint

* lint

* lint

* ts

* fix listMapData destructure

* allow expanding during dragging

* lint

* review comments

* fixes

* lint

* update story to use useTreeData

* handle dropPosition === 'on' case

* set shouldSelectOnPressUp on item if dragging enabled

* update useTreeData move/moveBefore/moveAfter implementations

* fix stories

* typescript

* add drag button

* fix nested item drop button in story

* prevent dropping items onto themselves or their descendants

* add tests and fix broken ones

* remove dead code

* fix lost items when moving to root

* fix tests

---------

Co-authored-by: GitHub <[email protected]>
Co-authored-by: Reid Barber <[email protected]>
Co-authored-by: Robert Snow <[email protected]>

Author: 3 people

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

@@ -36,7 +36,6 @@ export interface DropOptions {
   /**
    * Handler that is called after a valid drag is held over the drop target for a period of time.
    * This typically opens the item so that the user can drop within it.
-   * @private
    */
   onDropActivate?: (e: DropActivateEvent) => void,
   /** Handler that is called when a valid drag exits the drop target. */

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

@@ -30,6 +30,7 @@ import {
   DropTargetDelegate,
   Key,
   KeyboardDelegate,
+  KeyboardEvents,
   Node,
   RefObject
 } from '@react-types/shared';
@@ -46,7 +47,8 @@ export interface DroppableCollectionOptions extends DroppableCollectionProps {
   /** A delegate object that implements behavior for keyboard focus movement. */
   keyboardDelegate: KeyboardDelegate,
   /** A delegate object that provides drop targets for pointer coordinates within the collection. */
-  dropTargetDelegate: DropTargetDelegate
+  dropTargetDelegate: DropTargetDelegate,
+  onKeyDown?: KeyboardEvents['onKeyDown']
 }
 
 export interface DroppableCollectionResult {
@@ -201,7 +203,7 @@ export function useDroppableCollection(props: DroppableCollectionOptions, state:
       autoScroll.stop();
     },
     onDropActivate(e) {
-      if (state.target?.type === 'item' && state.target?.dropPosition === 'on' && typeof props.onDropActivate === 'function') {
+      if (state.target?.type === 'item' && typeof props.onDropActivate === 'function') {
         props.onDropActivate({
           type: 'dropactivate',
           x: e.x, // todo
@@ -741,6 +743,7 @@ export function useDroppableCollection(props: DroppableCollectionOptions, state:
             break;
           }
         }
+        localState.props.onKeyDown?.(e as any);
       }
     });
   }, [localState, ref, onDrop, direction]);

packages/@react-stately/data/src/useTreeData.ts

@@ -107,6 +107,20 @@ export interface TreeData<T extends object> {
    */
   move(key: Key, toParentKey: Key | null, index: number): void,
 
+  /**
+   * Moves one or more items before a given key.
+   * @param key - The key of the item to move the items before.
+   * @param keys - The keys of the items to move.
+   */
+  moveBefore(key: Key, keys: Iterable<Key>): void,
+
+  /**
+   * Moves one or more items after a given key.
+   * @param key - The key of the item to move the items after.
+   * @param keys - The keys of the items to move.
+   */
+  moveAfter(key: Key, keys: Iterable<Key>): void,
+
   /**
    * Updates an item in the tree.
    * @param key - The key of the item to update.
@@ -115,6 +129,11 @@ export interface TreeData<T extends object> {
   update(key: Key, newValue: T): void
 }
 
+interface TreeDataState<T extends object> {
+    items: TreeNode<T>[],
+    nodeMap: Map<Key, TreeNode<T>>
+}
+
 /**
  * Manages state for an immutable tree data structure, and provides convenience methods to
  * update the data over time.
@@ -128,7 +147,7 @@ export function useTreeData<T extends object>(options: TreeOptions<T>): TreeData
   } = options;
 
   // We only want to compute this on initial render.
-  let [tree, setItems] = useState<{items: TreeNode<T>[], nodeMap: Map<Key, TreeNode<T>>}>(() => buildTree(initialItems, new Map()));
+  let [tree, setItems] = useState<TreeDataState<T>>(() => buildTree(initialItems, new Map()));
   let {items, nodeMap} = tree;
 
   let [selectedKeys, setSelectedKeys] = useState(new Set<Key>(initialSelectedKeys || []));
@@ -141,7 +160,7 @@ export function useTreeData<T extends object>(options: TreeOptions<T>): TreeData
       items: initialItems.map(item => {
         let node: TreeNode<T> = {
           key: getKey(item),
-          parentKey: parentKey,
+          parentKey: parentKey ?? null,
           value: item,
           children: null
         };
@@ -154,9 +173,9 @@ export function useTreeData<T extends object>(options: TreeOptions<T>): TreeData
     };
   }
 
-  function updateTree(items: TreeNode<T>[], key: Key, update: (node: TreeNode<T>) => TreeNode<T> | null, originalMap: Map<Key, TreeNode<T>>) {
-    let node = originalMap.get(key);
-    if (!node) {
+  function updateTree(items: TreeNode<T>[], key: Key | null, update: (node: TreeNode<T>) => TreeNode<T> | null, originalMap: Map<Key, TreeNode<T>>) {
+    let node = key == null ? null : originalMap.get(key);
+    if (node == null) {
       return {items, nodeMap: originalMap};
     }
     let map = new Map<Key, TreeNode<T>>(originalMap);
@@ -233,7 +252,6 @@ export function useTreeData<T extends object>(options: TreeOptions<T>): TreeData
       }
     }
   }
-
   return {
     items,
     selectedKeys,
@@ -352,6 +370,8 @@ export function useTreeData<T extends object>(options: TreeOptions<T>): TreeData
 
         // If parentKey is null, insert into the root.
         if (toParentKey == null) {
+          // safe to reuse the original map since no node was actually removed, so we just need to update the one moved node
+          newMap = new Map(originalMap);
           newMap.set(movedNode.key, movedNode);
           return {items: [
             ...newItems.slice(0, index),
@@ -373,6 +393,39 @@ export function useTreeData<T extends object>(options: TreeOptions<T>): TreeData
         }), newMap);
       });
     },
+    moveBefore(key: Key, keys: Iterable<Key>) {
+      setItems((prevState) => {
+        let {items, nodeMap} = prevState;
+        let node = nodeMap.get(key);
+        if (!node) {
+          return prevState;
+        }
+        let toParentKey = node.parentKey ?? null;
+        let parent: null | TreeNode<T> = null;
+        if (toParentKey != null) {
+          parent = nodeMap.get(toParentKey) ?? null;
+        }
+        let toIndex = parent?.children ? parent.children.indexOf(node) : items.indexOf(node);
+        return moveItems(prevState, keys, parent, toIndex, updateTree);
+      });
+    },
+    moveAfter(key: Key, keys: Iterable<Key>) {
+      setItems((prevState) => {
+        let {items, nodeMap} = prevState;
+        let node = nodeMap.get(key);
+        if (!node) {
+          return prevState;
+        }
+        let toParentKey = node.parentKey ?? null;
+        let parent: null | TreeNode<T> = null;
+        if (toParentKey != null) {
+          parent = nodeMap.get(toParentKey) ?? null;
+        }
+        let toIndex = parent?.children ? parent.children.indexOf(node) : items.indexOf(node);
+        toIndex++;
+        return moveItems(prevState, keys, parent, toIndex, updateTree);
+      });
+    },
     update(oldKey: Key, newValue: T) {
       setItems(({items, nodeMap: originalMap}) => updateTree(items, oldKey, oldNode => {
         let node: TreeNode<T> = {
@@ -389,3 +442,97 @@ export function useTreeData<T extends object>(options: TreeOptions<T>): TreeData
     }
   };
 }
+
+function moveItems<T extends object>(
+  state: TreeDataState<T>,
+  keys: Iterable<Key>,
+  toParent: TreeNode<T> | null,
+  toIndex: number,
+  updateTree: (
+    items: TreeNode<T>[],
+    key: Key,
+    update: (node: TreeNode<T>) => TreeNode<T> | null,
+    originalMap: Map<Key, TreeNode<T>>
+  ) => TreeDataState<T>
+): TreeDataState<T> {
+  let {items, nodeMap} = state;
+
+  let parent = toParent;
+  let removeKeys = new Set(keys);
+  while (parent?.parentKey != null) {
+    if (removeKeys.has(parent.key)) {
+      throw new Error('Cannot move an item to be a child of itself.');
+    }
+    parent = nodeMap.get(parent.parentKey!) ?? null;
+  }
+
+  let originalToIndex = toIndex;
+
+  let keyArray = Array.isArray(keys) ? keys : [...keys];
+  // depth first search to put keys in order
+  let inOrderKeys: Map<Key, number> = new Map();
+  let removedItems: Array<TreeNode<T>> = [];
+  let newItems = items;
+  let newMap = nodeMap;
+  let i = 0;
+
+  function traversal(node, {inorder, postorder}) {
+    inorder?.(node);
+    if (node != null) {
+      for (let child of node.children ?? []) {
+        traversal(child, {inorder, postorder});
+        postorder?.(child);
+      }
+    }
+  }
+
+  function inorder(child) {
+    // in-order so we add items as we encounter them in the tree, then we can insert them in expected order later
+    if (keyArray.includes(child.key)) {
+      inOrderKeys.set(child.key, i++);
+    }
+  }
+
+  function postorder(child) {
+    // remove items and update the tree from the leaves and work upwards toward the root, this way
+    // we don't copy child node references from parents inadvertently
+    if (keyArray.includes(child.key)) {
+      removedItems.push({...newMap.get(child.key)!, parentKey: toParent?.key ?? null});
+      let {items: nextItems, nodeMap: nextMap} = updateTree(newItems, child.key, () => null, newMap);
+      newItems = nextItems;
+      newMap = nextMap;
+    }
+    // decrement the index if the child being removed is in the target parent and before the target index
+    if (child.parentKey === toParent?.key
+      && keyArray.includes(child.key)
+      && (toParent?.children ? toParent.children.indexOf(child) : items.indexOf(child)) < originalToIndex) {
+      toIndex--;
+    }
+  }
+
+  traversal({children: items}, {inorder, postorder});
+
+  let inOrderItems = removedItems.sort((a, b) => inOrderKeys.get(a.key)! > inOrderKeys.get(b.key)! ? 1 : -1);
+  // If parentKey is null, insert into the root.
+  if (!toParent || toParent.key == null) {
+    newMap = new Map(nodeMap);
+    inOrderItems.forEach(movedNode => newMap.set(movedNode.key, movedNode));
+    return {items: [
+      ...newItems.slice(0, toIndex),
+      ...inOrderItems,
+      ...newItems.slice(toIndex)
+    ], nodeMap: newMap};
+  }
+
+  // Otherwise, update the parent node and its ancestors.
+  return updateTree(newItems, toParent.key, parentNode => ({
+    key: parentNode.key,
+    parentKey: parentNode.parentKey,
+    value: parentNode.value,
+    children: [
+      ...parentNode.children!.slice(0, toIndex),
+      ...inOrderItems,
+      ...parentNode.children!.slice(toIndex)
+    ]
+  }), newMap);
+}