diff --git a/packages/material-ui-lab/src/TreeItem/TreeItem.js b/packages/material-ui-lab/src/TreeItem/TreeItem.js index 93a3a9577f58d3..9a9878d319876f 100644 --- a/packages/material-ui-lab/src/TreeItem/TreeItem.js +++ b/packages/material-ui-lab/src/TreeItem/TreeItem.js @@ -81,7 +81,8 @@ const TreeItem = React.forwardRef(function TreeItem(props, ref) { focusPreviousNode, handleFirstChars, handleLeftArrow, - handleNodeMap, + addNodeToNodeMap, + removeNodeFromNodeMap, icons: contextIcons, isExpanded, isFocused, @@ -222,11 +223,20 @@ const TreeItem = React.forwardRef(function TreeItem(props, ref) { }; React.useEffect(() => { - const childIds = React.Children.map(children, child => child.props.nodeId); - if (handleNodeMap) { - handleNodeMap(nodeId, childIds); + const childIds = React.Children.map(children, child => child.props.nodeId) || []; + if (addNodeToNodeMap) { + addNodeToNodeMap(nodeId, childIds); } - }, [children, nodeId, handleNodeMap]); + }, [children, nodeId, addNodeToNodeMap]); + + React.useEffect(() => { + if (removeNodeFromNodeMap) { + return () => { + removeNodeFromNodeMap(nodeId); + }; + } + return undefined; + }, [nodeId, removeNodeFromNodeMap]); React.useEffect(() => { if (handleFirstChars && label) { diff --git a/packages/material-ui-lab/src/TreeView/TreeView.js b/packages/material-ui-lab/src/TreeView/TreeView.js index 52996df0d5b4a3..286224cc78726b 100644 --- a/packages/material-ui-lab/src/TreeView/TreeView.js +++ b/packages/material-ui-lab/src/TreeView/TreeView.js @@ -13,6 +13,16 @@ export const styles = { }, }; +function arrayDiff(arr1, arr2) { + if (arr1.length !== arr2.length) return true; + + for (let i = 0; i < arr1.length; i += 1) { + if (arr1[i] !== arr2[i]) return true; + } + + return false; +} + const defaultExpandedDefault = []; const TreeView = React.forwardRef(function TreeView(props, ref) { @@ -40,18 +50,21 @@ const TreeView = React.forwardRef(function TreeView(props, ref) { const isTabable = id => tabable === id; const isFocused = id => focused === id; + const prevChildIds = React.useRef([]); React.useEffect(() => { - nodeMap.current = {}; - const childIds = React.Children.map(children, child => child.props.nodeId); - nodeMap.current[-1] = { parent: null, children: childIds }; - - (childIds || []).forEach((id, index) => { - if (index === 0) { - firstNode.current = id; - setTabable(id); - } - nodeMap.current[id] = { parent: null }; - }); + const childIds = React.Children.map(children, child => child.props.nodeId) || []; + if (arrayDiff(prevChildIds.current, childIds)) { + nodeMap.current[-1] = { parent: null, children: childIds }; + + childIds.forEach((id, index) => { + if (index === 0) { + firstNode.current = id; + setTabable(id); + } + nodeMap.current[id] = { parent: null }; + }); + prevChildIds.current = childIds; + } }, [children]); const getLastNode = React.useCallback( @@ -248,15 +261,30 @@ const TreeView = React.forwardRef(function TreeView(props, ref) { } }; - const handleNodeMap = (id, childrenIds) => { + const addNodeToNodeMap = (id, childrenIds) => { const currentMap = nodeMap.current[id]; nodeMap.current[id] = { ...currentMap, children: childrenIds, id }; - (childrenIds || []).forEach(childId => { + childrenIds.forEach(childId => { const currentChildMap = nodeMap.current[childId]; nodeMap.current[childId] = { ...currentChildMap, parent: id, id: childId }; }); }; + const removeNodeFromNodeMap = id => { + const map = nodeMap.current[id]; + if (map) { + if (map.parent) { + const parentMap = nodeMap.current[map.parent]; + if (parentMap && parentMap.children) { + const parentChildren = parentMap.children.filter(c => c !== id); + nodeMap.current[map.parent] = { ...parentMap, children: parentChildren }; + } + } + + delete nodeMap.current[id]; + } + }; + const handleFirstChars = (id, firstChar) => { firstCharMap.current[id] = firstChar; }; @@ -272,7 +300,8 @@ const TreeView = React.forwardRef(function TreeView(props, ref) { focusPreviousNode, handleFirstChars, handleLeftArrow, - handleNodeMap, + addNodeToNodeMap, + removeNodeFromNodeMap, icons: { defaultCollapseIcon, defaultExpandIcon, defaultParentIcon, defaultEndIcon }, isExpanded, isFocused, diff --git a/packages/material-ui-lab/src/TreeView/TreeView.test.js b/packages/material-ui-lab/src/TreeView/TreeView.test.js index e298f9c6b74735..12601505c05ad2 100644 --- a/packages/material-ui-lab/src/TreeView/TreeView.test.js +++ b/packages/material-ui-lab/src/TreeView/TreeView.test.js @@ -21,6 +21,38 @@ describe('', () => { after: () => mount.cleanUp(), })); + it('should not error when component state changes', () => { + function MyComponent() { + const [, setState] = React.useState(1); + + return ( + + { + setState(Math.random); + }} + > + + + + ); + } + + const { getByText, getByTestId } = render(); + + fireEvent.click(getByText('one')); + expect(getByTestId('one')).to.be.focused; + fireEvent.keyDown(document.activeElement, { key: 'ArrowDown' }); + expect(getByTestId('two')).to.be.focused; + fireEvent.keyDown(document.activeElement, { key: 'ArrowUp' }); + expect(getByTestId('one')).to.be.focused; + fireEvent.keyDown(document.activeElement, { key: 'ArrowDown' }); + expect(getByTestId('two')).to.be.focused; + }); + describe('onNodeToggle', () => { it('should be called when a parent node is clicked', () => { const handleNodeToggle = spy();