Fix useTreeData subtree move validation

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

@@ -378,6 +378,14 @@ export function useTreeData<T extends object>(options: TreeOptions<T>): TreeData
     },
     move(key: Key, toParentKey: Key | null, index: number) {
       setItems(({items, nodeMap: originalMap}) => {
+        let current = toParentKey;
+        while (current != null) {
+          if (current === key) {
+            throw new Error('Cannot move an item to be a child of itself.');
+          }
+          current = originalMap.get(current)?.parentKey ?? null;
+        }
+
         let node = originalMap.get(key);
         if (!node) {
           return {items, nodeMap: originalMap};
@@ -491,11 +499,11 @@ function moveItems<T extends object>(
 
   let parent = toParent;
   let removeKeys = new Set(keys);
-  while (parent?.parentKey != null) {
+  while (parent != 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;
+    parent = parent.parentKey != null ? (nodeMap.get(parent.parentKey) ?? null) : null;
   }
 
   let originalToIndex = toIndex;

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

@@ -903,7 +903,7 @@ describe('useTreeData', function () {
     expect(result.current.items[0].children[0].children.length).toEqual(3);
   });
 
-  describe('moveBefore error', function () {
+  describe('move errors', function () {
     const consoleError = console.error;
     beforeEach(() => {
       console.error = jest.fn();
@@ -913,15 +913,50 @@ describe('useTreeData', function () {
       console.error = consoleError;
     });
 
-    it('cannot move an item within itself', function () {
+    let reactMajor = parseInt(React.version, 10);
+    let expectMoveError = (result, action) => {
+      if (reactMajor >= 18) {
+        expect(() => act(action)).toThrow('Cannot move an item to be a child of itself.');
+      } else {
+        act(action);
+        expect(result.error?.message).toContain('Cannot move an item to be a child of itself.');
+      }
+    };
+
+    it('cannot move an item before a child inside itself', function () {
       const initialItems = [...initial, {name: 'Emily'}, {name: 'Eli'}];
+      let {result} = renderHook(() => useTreeData({initialItems, getChildren, getKey}));
+      expectMoveError(result, () => result.current.moveBefore('Suzie', ['John', 'Sam', 'Eli']));
+    });
 
+    it('cannot move a root-level item before any of its children', function () {
+      const initialItems = [...initial, {name: 'Emily'}, {name: 'Eli'}];
       let {result} = renderHook(() => useTreeData({initialItems, getChildren, getKey}));
-      try {
-        act(() => result.current.moveBefore('Suzie', ['John', 'Sam', 'Eli']));
-      } catch (e) {
-        expect(e.toString()).toContain('Cannot move an item to be a child of itself.');
-      }
+      expectMoveError(result, () => result.current.moveBefore('Suzie', ['David']));
+    });
+
+    it('cannot move an item after a child inside itself', function () {
+      const initialItems = [...initial, {name: 'Emily'}, {name: 'Eli'}];
+      let {result} = renderHook(() => useTreeData({initialItems, getChildren, getKey}));
+      expectMoveError(result, () => result.current.moveAfter('Suzie', ['John']));
+    });
+
+    it('cannot move a root-level item after any of its children', function () {
+      const initialItems = [...initial, {name: 'Emily'}, {name: 'Eli'}];
+      let {result} = renderHook(() => useTreeData({initialItems, getChildren, getKey}));
+      expectMoveError(result, () => result.current.moveAfter('Suzie', ['David']));
+    });
+
+    it('cannot move an item into itself', function () {
+      const initialItems = [...initial, {name: 'Emily'}, {name: 'Eli'}];
+      let {result} = renderHook(() => useTreeData({initialItems, getChildren, getKey}));
+      expectMoveError(result, () => result.current.move('John', 'Suzie', 0));
+    });
+
+    it('cannot move a root-level item into itself', function () {
+      const initialItems = [...initial, {name: 'Emily'}, {name: 'Eli'}];
+      let {result} = renderHook(() => useTreeData({initialItems, getChildren, getKey}));
+      expectMoveError(result, () => result.current.move('David', 'Suzie', 0));
     });
   });