[TreeView] Update firstCharMap when a TreeItem is removed

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

@@ -519,7 +519,7 @@ describe('<TreeItem />', () => {
                 <button
                   data-testid="button"
                   type="button"
-                  onClick={() => setState(value => !value)}
+                  onClick={() => setState((value) => !value)}
                 >
                   Toggle Hide
                 </button>
@@ -738,6 +738,36 @@ describe('<TreeItem />', () => {
           fireEvent.keyDown(document.activeElement, { key: 'v', shiftKey: true });
           expect(getByTestId('apple')).to.have.focus;
         });
+
+        it('should not throw when an item is removed', () => {
+          function TestComponent() {
+            const [hide, setState] = React.useState(false);
+            return (
+              <React.Fragment>
+                <button type="button" onClick={() => setState(true)}>
+                  Hide
+                </button>
+                <TreeView>
+                  {!hide && <TreeItem nodeId="hide" label="ab" />}
+                  <TreeItem nodeId="keyDown" data-testid="keyDown" />
+                  <TreeItem nodeId="navTo" data-testid="navTo" label="ac" />
+                </TreeView>
+              </React.Fragment>
+            );
+          }
+
+          const { getByText, getByTestId } = render(<TestComponent />);
+          fireEvent.click(getByText('Hide'));
+          const navTreeItem = getByTestId('navTo');
+          expect(navTreeItem).not.to.have.focus;
+
+          expect(() => {
+            fireEvent.keyDown(getByTestId('keyDown'), { key: 'a' });
+          }).not.to.throw();
+
+          expect(navTreeItem).to.have.focus;
+          expect(navTreeItem).to.have.attribute('tabindex', '0');
+        });
       });
 
       describe('asterisk key interaction', () => {

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

@@ -416,9 +416,20 @@ const TreeView = React.forwardRef(function TreeView(props, ref) {
     return nodes;
   }, []);
 
+  const cleanUpFirstCharMap = React.useCallback((nodes) => {
+    const newMap = { ...firstCharMap.current };
+    nodes.forEach((node) => {
+      if (newMap[node]) {
+        delete newMap[node];
+      }
+    });
+    firstCharMap.current = newMap;
+  }, []);
+
   const removeNodeFromNodeMap = React.useCallback(
     (id) => {
       const nodes = getNodesToRemove(id);
+      cleanUpFirstCharMap(nodes);
       const newMap = { ...nodeMap.current };
 
       nodes.forEach((node) => {
@@ -444,7 +455,7 @@ const TreeView = React.forwardRef(function TreeView(props, ref) {
         return oldFocusedNodeId;
       });
     },
-    [getNodesToRemove],
+    [getNodesToRemove, cleanUpFirstCharMap],
   );
 
   const mapFirstChar = (id, firstChar) => {