Repository: mui/material-ui (TypeScript) — Bug Fix

Issue:
[TreeView] Locks up preventing mouse & up/down arrows from working
<!-- Provide a general summary of the issue in the Title above -->

TreeView will often lock up, preventing mouse up/down arrows from working. This happens after it has re-rendered, even if the tree data hasn't changed.

<!--
  Thank you very much for contributing to Material-UI by creating an issue! ❤️
  To avoid duplicate issues we ask you to check off the following list.
-->

- [x] The issue is present in the latest release.
- [x] I have searched the [issues](https://github.com/mui-org/material-ui/issues) of this repository and believe that this is not a duplicate.

## Current Behavior 😯

Keyboard controls stop working (sometimes will start working again)
The lockup prevents up / down, left / right still typically work

Console error:
```
Uncaught TypeError: Cannot read property '0' of undefined
    at getNextNode (VM2322 TreeView.js:129)
    at focusNextNode (VM2322 TreeView.js:180)
    at handleKeyDown (VM2337 TreeItem.js:199)
```

## Expected Behavior 🤔

No errors, keys continues to work

## Steps to Reproduce 🕹

screencast: https://imgur.com/vFcXNxS
https://codesandbox.io/s/material-ui-treeview-bug-td9o7

Steps:
1. Open console
2. Use keyboard to move around the tree (rapidly press... up, down, left, right, etc. )
3. It will eventually lock up logging the error I've pasted above

I also suspect this happens more frequently with a more complex label.

## Context 🔦

I use TreeItem onFocus to update the URL, this causes re-render of TreeView, which causes errors. It doesn't crash the app but the keyboard does stop responding.

## Your Environment 🌎

Reproducible in codesandbox


Discussion:
Exact key press sequence to cause error
1. click first item
2. down, right, up, down
3. down will not respond anymore
4. press left arrow to make it work again
Thanks for the report, @eps1lon was experiencing something similar. I can pick this up and look into it at the same time. It wouldn’t surprise me if they are caused by the same thing.
Thanks for looking into this Josh, let me know if you want me to test or troubleshoot.
@oisinlavery I don't seem to be able to reproduce this. Can you give me details of your environment? Browser, OS, etc.
I accidentally updated the linked codesandbox example instead of forking it.

Here is a new version which shows the bug:
https://codesandbox.io/s/material-ui-treeview-bug-td9o7

screencast: https://imgur.com/vFcXNxS

I'm on macOS (10.14.6) + latest Chrome (77.0.3865.90)
Thanks for the update, looking behind the scenes it seems children aren’t being picked up correctly, I’m still investigating why.
@joshwooding I've come across similar issues to this while investigating #17705, have you made any progress?
@flurmbo Sadly real life came up but I can look a bit more into it soon. The problem reported by @eps1lon was different, so I need to take a closer look.
@joshwooding, no problem. When I investigated @eps1lon's issue I found that when trying to make the TreeView a controlled component that updating the props of the TreeView would force the useEffect hooks dependent on `props.children` to rerun, which would overwrite the `nodeMap`. Then when trying to use the keyboard controls I would get similar errors.

Resolving patch:
diff --git a/packages/material-ui-lab/src/TreeItem/TreeItem.js b/packages/material-ui-lab/src/TreeItem/TreeItem.js
--- 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
--- 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,


Test patch:
diff --git a/packages/material-ui-lab/src/TreeView/TreeView.test.js b/packages/material-ui-lab/src/TreeView/TreeView.test.js
--- 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('<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();