[5574] - add moveBefore and moveAfter to useTreeData (#7689)

* [5574] - add moveBefore and moveAfter to useTreeData

* add docs

* remove onlys

* remove console logs

Author: rob-clayburn

Diff for: packages/@react-stately/data/docs/useTreeData.mdx

@@ -182,6 +182,33 @@ tree.move('Sam', 'Animals', 1);
 tree.move('Sam', null, 1);
 ```
 
+### Move before
+An alias to move
+
+```tsx
+// Move an item within the same parent
+tree.moveBefore('Sam', 'People', 0);
+
+// Move an item to a different parent
+tree.moveBefore('Sam', 'Animals', 1);
+
+// Move an item to the root
+tree.moveBefore('Sam', null, 1);
+```
+
+### Move after
+
+```tsx
+// Move an item within the same parent
+tree.moveAfter('Sam', 'People', 0);
+
+// Move an item to a different parent
+tree.moveAfter('Sam', 'Animals', 1);
+
+// Move an item to the root
+tree.moveAfter('Sam', null, 1);
+```
+
 ### Updating items
 
 ```tsx

Diff for: packages/@react-stately/data/src/useTreeData.ts

@@ -107,6 +107,22 @@ export interface TreeData<T extends object> {
    */
   move(key: Key, toParentKey: Key | null, index: number): void,
 
+  /**
+   * Moves an item before a node within the tree.
+   * @param key - The key of the item to move.
+   * @param toParentKey - The key of the new parent to insert into. `null` for the root.
+   * @param index - The index within the new parent to insert before.
+   */
+  moveBefore(key: Key, toParentKey: Key | null, index: number): void,
+
+  /**
+   * Moves an item after a node within the tree.
+   * @param key - The key of the item to move.
+   * @param toParentKey - The key of the new parent to insert into. `null` for the root.
+   * @param index - The index within the new parent to insert after.
+   */
+  moveAfter(key: Key, toParentKey: Key | null, index: number): void,
+
   /**
    * Updates an item in the tree.
    * @param key - The key of the item to update.
@@ -373,6 +389,50 @@ export function useTreeData<T extends object>(options: TreeOptions<T>): TreeData
         }), newMap);
       });
     },
+    moveBefore(key: Key, toParentKey: Key | null, index: number) {
+      this.move(key, toParentKey, index);
+    },
+    moveAfter(key: Key, toParentKey: Key | null, index: number) {
+      setItems(({items, nodeMap: originalMap}) => {
+        let node = originalMap.get(key);
+        if (!node) {
+          return {items, nodeMap: originalMap};
+        }
+
+        let {items: newItems, nodeMap: newMap} = updateTree(items, key, () => null, originalMap);
+
+        const movedNode = {
+          ...node,
+          parentKey: toParentKey
+        };
+
+        const afterIndex = items.length === index ? index : index + 1;
+        // If parentKey is null, insert into the root.
+        if (toParentKey == null) {
+          newMap.set(movedNode.key, movedNode);
+          return {items: [
+            ...newItems.slice(0, afterIndex),
+            movedNode,
+            ...newItems.slice(afterIndex)
+          ], nodeMap: newMap};
+        }
+
+        // Otherwise, update the parent node and its ancestors.
+        return updateTree(newItems, toParentKey, parentNode => {
+          const c = [
+            ...parentNode.children!.slice(0, afterIndex),
+            movedNode,
+            ...parentNode.children!.slice(afterIndex)
+          ];
+          return {
+            key: parentNode.key,
+            parentKey: parentNode.parentKey,
+            value: parentNode.value,
+            children: c
+          };
+        }, newMap);
+      });
+    },
     update(oldKey: Key, newValue: T) {
       setItems(({items, nodeMap: originalMap}) => updateTree(items, oldKey, oldNode => {
         let node: TreeNode<T> = {

Diff for: packages/@react-stately/data/test/useTreeData.test.js

@@ -676,4 +676,74 @@ describe('useTreeData', function () {
     expect(result.current.items[1].value).toEqual(initialResult.items[2].value);
     expect(result.current.items[2]).toEqual(initialResult.items[1]);
   });
+
+
+  it('should move an item within its same level before the target', function () {
+    const initialItems = [...initial, {name: 'Emily'}, {name: 'Eli'}];
+
+    let {result} = renderHook(() =>
+      useTreeData({initialItems, getChildren, getKey})
+    );
+    act(() => {
+      result.current.moveBefore('Eli', null, 0);
+    });
+    expect(result.current.items[0].key).toEqual('Eli');
+    expect(result.current.items[1].key).toEqual('David');
+    expect(result.current.items[2].key).toEqual('Emily');
+    expect(result.current.items.length).toEqual(initialItems.length);
+  });
+
+  it('should move an item to a different level before the target', function () {
+    const initialItems = [...initial, {name: 'Emily'}, {name: 'Eli'}];
+
+    let {result} = renderHook(() =>
+      useTreeData({initialItems, getChildren, getKey})
+    );
+    act(() => {
+      result.current.moveBefore('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');
+    expect(result.current.items.length).toEqual(2);
+  });
+
+  it('should move an item to a different level after the target', function () {
+    const initialItems = [...initial, {name: 'Emily'}, {name: 'Eli'}];
+    let {result} = renderHook(() =>
+      useTreeData({initialItems, getChildren, getKey})
+    );
+
+    act(() => {
+      result.current.moveAfter('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('Sam');
+    expect(result.current.items[0].children[2].key).toEqual('Eli');
+    expect(result.current.items[1].key).toEqual('Emily');
+    expect(result.current.items.length).toEqual(2);
+  });
+
+  it('should move an item to a different level at the end when the index is greater than the node list length', function () {
+    const initialItems = [...initial, {name: 'Emily'}, {name: 'Eli'}];
+    console.log('initialItems', initialItems[0]);
+    let {result} = renderHook(() =>
+      useTreeData({initialItems, getChildren, getKey})
+    );
+
+    act(() => {
+      result.current.moveAfter('Eli', 'David', 100);
+    });
+    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('Sam');
+    expect(result.current.items[0].children[2].key).toEqual('Jane');
+    expect(result.current.items[0].children[3].key).toEqual('Eli');
+    expect(result.current.items[1].key).toEqual('Emily');
+    expect(result.current.items.length).toEqual(2);
+  });
 });