fix moving an item to the root (#8225)

Author: reidbarber

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

@@ -370,9 +370,7 @@ 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);
+          addNode(movedNode, newMap);
           return {items: [
             ...newItems.slice(0, index),
             movedNode,
@@ -406,7 +404,7 @@ export function useTreeData<T extends object>(options: TreeOptions<T>): TreeData
           parent = nodeMap.get(toParentKey) ?? null;
         }
         let toIndex = parent?.children ? parent.children.indexOf(node) : items.indexOf(node);
-        return moveItems(prevState, keys, parent, toIndex, updateTree);
+        return moveItems(prevState, keys, parent, toIndex, updateTree, addNode);
       });
     },
     moveAfter(key: Key, keys: Iterable<Key>) {
@@ -423,7 +421,7 @@ export function useTreeData<T extends object>(options: TreeOptions<T>): TreeData
         }
         let toIndex = parent?.children ? parent.children.indexOf(node) : items.indexOf(node);
         toIndex++;
-        return moveItems(prevState, keys, parent, toIndex, updateTree);
+        return moveItems(prevState, keys, parent, toIndex, updateTree, addNode);
       });
     },
     update(oldKey: Key, newValue: T) {
@@ -450,10 +448,11 @@ function moveItems<T extends object>(
   toIndex: number,
   updateTree: (
     items: TreeNode<T>[],
-    key: Key,
+    key: Key | null,
     update: (node: TreeNode<T>) => TreeNode<T> | null,
     originalMap: Map<Key, TreeNode<T>>
-  ) => TreeDataState<T>
+  ) => TreeDataState<T>,
+  addNode: (node: TreeNode<T>, map: Map<Key, TreeNode<T>>) => void
 ): TreeDataState<T> {
   let {items, nodeMap} = state;
 
@@ -515,8 +514,9 @@ function moveItems<T extends object>(
   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));
+    inOrderItems.forEach(movedNode => {
+      addNode(movedNode, newMap); 
+    });
     return {items: [
       ...newItems.slice(0, toIndex),
       ...inOrderItems,

packages/@react-stately/data/test/useTreeData.test.js

@@ -610,6 +610,43 @@ describe('useTreeData', function () {
     expect(result.current.items[1].children[2]).toBe(
       initialResult.items[0].children[2]
     );
+
+    /*
+      Expected tree structure after moving 'Stacy' to root:
+      - Stacy
+      - David
+        |-- John
+        |   -- Suzie
+        |-- Sam
+        |   -- Brad
+        |-- Jane
+    */
+    let stacyNode = result.current.getItem('Stacy');
+    expect(stacyNode.parentKey).toBeNull();
+    expect(stacyNode.children).toHaveLength(0);
+    expect(stacyNode).toBe(result.current.items[0]);
+
+    let davidNode = result.current.getItem('David');
+    expect(davidNode.parentKey).toBeNull();
+    expect(davidNode.children.map(c => c.key)).toEqual(['John', 'Sam', 'Jane']);
+    expect(davidNode).toBe(result.current.items[1]);
+
+    let samNode = result.current.getItem('Sam');
+    expect(samNode.parentKey).toBe('David');
+    expect(samNode.children.map(c => c.key)).toEqual(['Brad']);
+    expect(samNode).toBe(result.current.items[1].children[1]);
+
+    let bradNode = result.current.getItem('Brad');
+    expect(bradNode.parentKey).toBe('Sam');
+    expect(bradNode).toBe(result.current.items[1].children[1].children[0]);
+
+    let johnNode = result.current.getItem('John');
+    expect(johnNode.parentKey).toBe('David');
+    expect(johnNode).toBe(result.current.items[1].children[0]);
+
+    let janeNode = result.current.getItem('Jane');
+    expect(janeNode.parentKey).toBe('David');
+    expect(janeNode).toBe(result.current.items[1].children[2]);
   });
 
   it('should move an item to a new index within its current parent', function () {
@@ -876,6 +913,7 @@ describe('useTreeData', function () {
       result.current.move('Eli', 'David', 1);
     });
     expect(result.current.items[0].key).toEqual('David');
+
     expect(result.current.items[0].children[0].key).toEqual('John');
     expect(result.current.items[0].children[1].key).toEqual('Eli');
     expect(result.current.items[1].key).toEqual('Emily');