[TreeView] Change when node map is built

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) {

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,

packages/material-ui-lab/src/TreeView/TreeView.test.js

@@ -21,6 +21,38 @@ describe('<TreeView />', () => {
     after: () => mount.cleanUp(),
   }));
 
+  it('should not error when component state changes', () => {
+    function MyComponent() {
+      const [, setState] = React.useState(1);
+
+      return (
+        <TreeView>
+          <TreeItem
+            nodeId="one"
+            label="one"
+            data-testid="one"
+            onFocus={() => {
+              setState(Math.random);
+            }}
+          >
+            <TreeItem nodeId="two" label="two" data-testid="two" />
+          </TreeItem>
+        </TreeView>
+      );
+    }
+
+    const { getByText, getByTestId } = render(<MyComponent />);
+
+    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();