Pure render for tree data, fixes React Canary useMemo/useState StrictMode

packages/@react-spectrum/listbox/stories/ListBox.stories.tsx

@@ -940,6 +940,7 @@ export function FocusExample(args = {}) {
     getKey: (item) => item.name,
     getChildren: (item:{name:string, children?:{name:string, children?:{name:string}[]}[]}) => item.children
   });
+
   let [dialog, setDialog] = useState(null);
   let ref = useRef(null);
   return (

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

@@ -11,7 +11,7 @@
  */
 
 import {Key} from '@react-types/shared';
-import {useMemo, useState} from 'react';
+import {useState} from 'react';
 
 export interface TreeOptions<T extends object> {
   /** Initial root items in the tree. */
@@ -126,44 +126,51 @@ export function useTreeData<T extends object>(options: TreeOptions<T>): TreeData
     getKey = (item: any) => item.id || item.key,
     getChildren = (item: any) => item.children
   } = options;
-  let map = useMemo(() => new Map<Key, TreeNode<T>>(), []);
 
   // We only want to compute this on initial render.
-  // eslint-disable-next-line react-hooks/exhaustive-deps
-  let initialNodes = useMemo(() => buildTree(initialItems), []);
-  let [items, setItems] = useState(initialNodes);
+  let [tree, setItems] = useState<{items: TreeNode<T>[], mappy: Map<Key, TreeNode<T>>}>(() => buildTree(initialItems));
+  let {items, mappy: map} = tree;
+
   let [selectedKeys, setSelectedKeys] = useState(new Set<Key>(initialSelectedKeys || []));
 
-  function buildTree(initialItems: T[] = [], parentKey?: Key | null) {
-    return initialItems.map(item => {
-      let node: TreeNode<T> = {
-        key: getKey(item),
-        parentKey: parentKey,
-        value: item,
-        children: null
-      };
+  function buildTree(initialItems: T[] = [], parentKey?: Key | null, acc?: Map<Key, TreeNode<T>>) {
+    let map = acc;
+    if (!acc) {
+      map = new Map<Key, TreeNode<T>>();
+    }
+    return {
+      items: initialItems.map(item => {
+        let node: TreeNode<T> = {
+          key: getKey(item),
+          parentKey: parentKey,
+          value: item,
+          children: null
+        };
 
-      node.children = buildTree(getChildren(item), node.key);
-      map.set(node.key, node);
-      return node;
-    });
+        node.children = buildTree(getChildren(item), node.key, map).items;
+        map.set(node.key, node);
+        return node;
+      }),
+      mappy: map
+    };
   }
 
-  function updateTree(items: TreeNode<T>[], key: Key, update: (node: TreeNode<T>) => TreeNode<T>) {
-    let node = map.get(key);
+  function updateTree(items: TreeNode<T>[], key: Key, update: (node: TreeNode<T>) => TreeNode<T>, originalMap: Map<Key, TreeNode<T>>) {
+    let node = originalMap.get(key);
     if (!node) {
-      return items;
+      return {items, mappy: originalMap};
     }
+    let map = new Map<Key, TreeNode<T>>(originalMap);
 
     // Create a new node. If null, then delete the node, otherwise replace.
     let newNode = update(node);
     if (newNode == null) {
-      deleteNode(node);
+      deleteNode(node, map);
     } else {
-      addNode(newNode);
+      addNode(newNode, map);
     }
 
-    // Walk up the tree and update each parent to refer to the new chilren.
+    // Walk up the tree and update each parent to refer to the new children.
     while (node.parentKey) {
       let nextParent = map.get(node.parentKey);
       let copy: TreeNode<T> = {
@@ -196,26 +203,30 @@ export function useTreeData<T extends object>(options: TreeOptions<T>): TreeData
       items = items.filter(c => c !== node);
     }
 
-    return items.map(item => {
-      if (item === node) {
-        return newNode;
-      }
+    return {
+      items: items.map(item => {
+        // if (item.key === node.key) {
+        if (item === node) {
+          return newNode;
+        }
 
-      return item;
-    });
+        return item;
+      }),
+      mappy: map
+    };
   }
 
-  function addNode(node: TreeNode<T>) {
+  function addNode(node: TreeNode<T>, map) {
     map.set(node.key, node);
     for (let child of node.children) {
-      addNode(child);
+      addNode(child, map);
     }
   }
 
-  function deleteNode(node: TreeNode<T>) {
+  function deleteNode(node: TreeNode<T>, map) {
     map.delete(node.key);
     for (let child of node.children) {
-      deleteNode(child);
+      deleteNode(child, map);
     }
   }
 
@@ -227,16 +238,19 @@ export function useTreeData<T extends object>(options: TreeOptions<T>): TreeData
       return map.get(key);
     },
     insert(parentKey: Key | null, index: number, ...values: T[]) {
-      setItems(items => {
-        let nodes = buildTree(values, parentKey);
+      setItems(({items, mappy: originalMap}) => {
+        let {items: nodes, mappy: newMap} = buildTree(values, parentKey, originalMap);
 
         // If parentKey is null, insert into the root.
         if (parentKey == null) {
-          return [
-            ...items.slice(0, index),
-            ...nodes,
-            ...items.slice(index)
-          ];
+          return {
+            items: [
+              ...items.slice(0, index),
+              ...nodes,
+              ...items.slice(index)
+            ],
+            mappy: newMap
+          };
         }
 
         // Otherwise, update the parent node and its ancestors.
@@ -249,7 +263,7 @@ export function useTreeData<T extends object>(options: TreeOptions<T>): TreeData
             ...nodes,
             ...parentNode.children.slice(index)
           ]
-        }));
+        }), newMap);
       });
     },
     insertBefore(key: Key, ...values: T[]): void {
@@ -291,15 +305,19 @@ export function useTreeData<T extends object>(options: TreeOptions<T>): TreeData
     },
     remove(...keys: Key[]) {
       let newItems = items;
+      let prevMap = map;
+      let tree;
       for (let key of keys) {
-        newItems = updateTree(newItems, key, () => null);
+        tree = updateTree(newItems, key, () => null, prevMap);
+        prevMap = tree.mappy;
+        newItems = tree.items;
       }
 
-      setItems(newItems);
+      setItems(tree);
 
       let selection = new Set(selectedKeys);
       for (let key of selectedKeys) {
-        if (!map.has(key)) {
+        if (!tree.mappy.has(key)) {
           selection.delete(key);
         }
       }
@@ -310,13 +328,14 @@ export function useTreeData<T extends object>(options: TreeOptions<T>): TreeData
       this.remove(...selectedKeys);
     },
     move(key: Key, toParentKey: Key | null, index: number) {
-      setItems(items => {
-        let node = map.get(key);
+      setItems(({items, mappy: originalMap}) => {
+        let node = originalMap.get(key);
         if (!node) {
-          return items;
+          return {items, mappy: originalMap};
         }
 
-        items = updateTree(items, key, () => null);
+        let tree = updateTree(items, key, () => null, originalMap);
+        let {items: newItems, mappy: map} = tree;
 
         const movedNode = {
           ...node,
@@ -325,15 +344,15 @@ export function useTreeData<T extends object>(options: TreeOptions<T>): TreeData
 
         // If parentKey is null, insert into the root.
         if (toParentKey == null) {
-          return [
-            ...items.slice(0, index),
+          return {items: [
+            ...newItems.slice(0, index),
             movedNode,
-            ...items.slice(index)
-          ];
+            ...newItems.slice(index)
+          ], mappy: map};
         }
 
         // Otherwise, update the parent node and its ancestors.
-        return updateTree(items, toParentKey, parentNode => ({
+        return updateTree(newItems, toParentKey, parentNode => ({
           key: parentNode.key,
           parentKey: parentNode.parentKey,
           value: parentNode.value,
@@ -342,21 +361,22 @@ export function useTreeData<T extends object>(options: TreeOptions<T>): TreeData
             movedNode,
             ...parentNode.children.slice(index)
           ]
-        }));
+        }), map);
       });
     },
     update(oldKey: Key, newValue: T) {
-      setItems(items => updateTree(items, oldKey, oldNode => {
+      setItems(({items, mappy: originalMap}) => updateTree(items, oldKey, oldNode => {
         let node: TreeNode<T> = {
           key: oldNode.key,
           parentKey: oldNode.parentKey,
           value: newValue,
           children: null
         };
 
-        node.children = buildTree(getChildren(newValue), node.key);
+        let tree = buildTree(getChildren(newValue), node.key, originalMap);
+        node.children = tree.items;
         return node;
-      }));
+      }, originalMap));
     }
   };
 }