Skip to content

Commit

Permalink
Pure render for tree data, fixes React Canary useMemo/useState Strict…
Browse files Browse the repository at this point in the history
…Mode (#5835)

* Pure render for useTreeData

* adjust some more calls

* fix accumulation of build trees

* better readability

* fix typo

* fix remaining broken tests

* revert canary testing changes

* fix lint

* updating naming and making things a bit easier to read

* make sure a new map is always generated in insert for purity

* make useTreeData always use and return a new Map

this avoids the case where we may inadvertantly modify the previous map directly which could be problematic in a double render

* guard against empy array in remove

---------

Co-authored-by: Daniel Lu <[email protected]>
  • Loading branch information
snowystinger and LFDanLu authored Feb 8, 2024
1 parent 3973a77 commit 7122c20
Show file tree
Hide file tree
Showing 2 changed files with 84 additions and 64 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down
147 changes: 83 additions & 64 deletions packages/@react-stately/data/src/useTreeData.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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. */
Expand Down Expand Up @@ -126,44 +126,47 @@ 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>[], nodeMap: Map<Key, TreeNode<T>>}>(() => buildTree(initialItems, new Map()));
let {items, nodeMap} = 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[] = [], map: Map<Key, TreeNode<T>>, parentKey?: Key | null) {
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), map, node.key).items;
map.set(node.key, node);
return node;
}),
nodeMap: 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, nodeMap: 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> = {
Expand Down Expand Up @@ -196,26 +199,29 @@ 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 === node) {
return newNode;
}

return item;
});
return item;
}),
nodeMap: map
};
}

function addNode(node: TreeNode<T>) {
function addNode(node: TreeNode<T>, map: Map<Key, TreeNode<T>>) {
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<Key, TreeNode<T>>) {
map.delete(node.key);
for (let child of node.children) {
deleteNode(child);
deleteNode(child, map);
}
}

Expand All @@ -224,19 +230,22 @@ export function useTreeData<T extends object>(options: TreeOptions<T>): TreeData
selectedKeys,
setSelectedKeys,
getItem(key: Key) {
return map.get(key);
return nodeMap.get(key);
},
insert(parentKey: Key | null, index: number, ...values: T[]) {
setItems(items => {
let nodes = buildTree(values, parentKey);
setItems(({items, nodeMap: originalMap}) => {
let {items: newNodes, nodeMap: newMap} = buildTree(values, originalMap, parentKey);

// 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),
...newNodes,
...items.slice(index)
],
nodeMap: newMap
};
}

// Otherwise, update the parent node and its ancestors.
Expand All @@ -246,30 +255,30 @@ export function useTreeData<T extends object>(options: TreeOptions<T>): TreeData
value: parentNode.value,
children: [
...parentNode.children.slice(0, index),
...nodes,
...newNodes,
...parentNode.children.slice(index)
]
}));
}), newMap);
});
},
insertBefore(key: Key, ...values: T[]): void {
let node = map.get(key);
let node = nodeMap.get(key);
if (!node) {
return;
}

let parentNode = map.get(node.parentKey);
let parentNode = nodeMap.get(node.parentKey);
let nodes = parentNode ? parentNode.children : items;
let index = nodes.indexOf(node);
this.insert(parentNode?.key, index, ...values);
},
insertAfter(key: Key, ...values: T[]): void {
let node = map.get(key);
let node = nodeMap.get(key);
if (!node) {
return;
}

let parentNode = map.get(node.parentKey);
let parentNode = nodeMap.get(node.parentKey);
let nodes = parentNode ? parentNode.children : items;
let index = nodes.indexOf(node);
this.insert(parentNode?.key, index + 1, ...values);
Expand All @@ -281,7 +290,7 @@ export function useTreeData<T extends object>(options: TreeOptions<T>): TreeData
if (parentKey == null) {
this.insert(null, items.length, ...values);
} else {
let parentNode = map.get(parentKey);
let parentNode = nodeMap.get(parentKey);
if (!parentNode) {
return;
}
Expand All @@ -290,16 +299,24 @@ export function useTreeData<T extends object>(options: TreeOptions<T>): TreeData
}
},
remove(...keys: Key[]) {
if (keys.length === 0) {
return;
}

let newItems = items;
let prevMap = nodeMap;
let newTree;
for (let key of keys) {
newItems = updateTree(newItems, key, () => null);
newTree = updateTree(newItems, key, () => null, prevMap);
prevMap = newTree.nodeMap;
newItems = newTree.items;
}

setItems(newItems);
setItems(newTree);

let selection = new Set(selectedKeys);
for (let key of selectedKeys) {
if (!map.has(key)) {
if (!newTree.nodeMap.has(key)) {
selection.delete(key);
}
}
Expand All @@ -310,13 +327,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, nodeMap: originalMap}) => {
let node = originalMap.get(key);
if (!node) {
return items;
return {items, nodeMap: originalMap};
}

items = updateTree(items, key, () => null);
let {items: newItems, nodeMap: newMap} = updateTree(items, key, () => null, originalMap);


const movedNode = {
...node,
Expand All @@ -325,15 +343,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)
], nodeMap: newMap};
}

// 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,
Expand All @@ -342,21 +360,22 @@ export function useTreeData<T extends object>(options: TreeOptions<T>): TreeData
movedNode,
...parentNode.children.slice(index)
]
}));
}), newMap);
});
},
update(oldKey: Key, newValue: T) {
setItems(items => updateTree(items, oldKey, oldNode => {
setItems(({items, nodeMap: 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), originalMap, node.key);
node.children = tree.items;
return node;
}));
}, originalMap));
}
};
}

1 comment on commit 7122c20

@rspbot
Copy link

@rspbot rspbot commented on 7122c20 Feb 8, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please sign in to comment.