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();