fix: tree component focus loss issue

In some case, for example on high CPU load the scroll to active node can
not keep up with the user input and the active element removed because
it is not in the view.

- check and scroll to active element when it is not in the DOM after
  each arrow key event.
- optimize MutationObserver creation

src/components/Tree/Tree.constants.ts

@@ -7,6 +7,7 @@ const DEFAULTS = {
   SELECTABLE_NODES: 'leafOnly' as const,
   IS_REQUIRED: false,
   NODE_ID_PREFIX: 'md-tree-node',
+  SHOULD_NODE_FOCUS_BE_INSET: true,
 };
 
 const STYLE = {

src/components/Tree/Tree.hooks.tsx

@@ -9,13 +9,18 @@ import { NODE_ID_ATTRIBUTE_NAME, NODE_ID_DATA_NAME } from '../TreeNodeBase/TreeN
 /**
  * Handle DOM changes for virtual tree
  *
+ * This hook manage 2 use cases:
+ * 1) When Virtual Tree removes active node, the hook adds a cloned node to not lose focus.
+ * 2) When the active node added back to the DOM, the hook invoke focus on the element and trigger
+ *    key event if it happened on the cloned element.
+ *
  * @param props
  * @internal
  */
 export const useVirtualTreeNavigation = ({
   virtualTreeConnector,
   treeRef,
-  activeNodeId,
+  activeNodeIdRef,
 }: UseVirtualTreeNavigationProps): void => {
   // Handle DOM changes for virtual tree
   useEffect(() => {
@@ -36,12 +41,13 @@ export const useVirtualTreeNavigation = ({
       );
 
     // Filter element by active node id
-    const getByNodeId = (node: HTMLElement) => node.dataset[NODE_ID_DATA_NAME] === activeNodeId;
+    const getByNodeId = (node: HTMLElement) =>
+      node.dataset[NODE_ID_DATA_NAME] === activeNodeIdRef.current;
 
     // Mutation observer to handle the focus change
     const mutationHandler: MutationCallback = (mutationList) => {
       for (const mutation of mutationList) {
-        // Handle removed active node
+        // Active node moved out of view and removed
         const removed = Array.from(mutation.removedNodes).find(getByNodeId) as HTMLElement;
         if (removed) {
           // Clone the active node and make some modifications
@@ -58,11 +64,11 @@ export const useVirtualTreeNavigation = ({
 
           // add focus and keydown event listeners
           clonedNode.addEventListener('focus', () => {
-            virtualTreeConnector.scrollToNode(activeNodeId);
+            virtualTreeConnector.scrollToNode(activeNodeIdRef.current);
           });
           clonedNode.addEventListener('keydown', (evt) => {
             if (TREE_NAVIGATION_KEYS.includes(evt.key) || (evt.key === 'Tab' && !evt.shiftKey)) {
-              virtualTreeConnector.scrollToNode(activeNodeId);
+              virtualTreeConnector.scrollToNode(activeNodeIdRef.current);
               keyDownEvent = new KeyboardEvent('keydown', evt);
               evt.stopPropagation();
               evt.preventDefault();
@@ -71,7 +77,7 @@ export const useVirtualTreeNavigation = ({
           return;
         }
 
-        // Handle adding back the focused node
+        // Active node moved back into the view and it added back to the DOM
         const added = Array.from(mutation.addedNodes).find(getByNodeId) as HTMLElement;
         if (added) {
           cleanUp();
@@ -103,5 +109,5 @@ export const useVirtualTreeNavigation = ({
       observer.disconnect();
       cleanUp();
     };
-  }, [virtualTreeConnector, activeNodeId, treeRef.current]);
+  }, [virtualTreeConnector, treeRef.current]);
 };

src/components/Tree/Tree.test.tsx

@@ -533,7 +533,7 @@ describe('<Tree />', () => {
       expect(scrollToNode).toHaveBeenNthCalledWith(1, '1');
     });
 
-    it('should not call scrollToNode when cloned node exists and user press Shift+Tab', async () => {
+    it('should call scrollToNode when the focus moves back to the tree and the active node is not in the DOM', async () => {
       expect.assertions(2);
       const scrollToNode = jest.fn();
       const { getByText } = await renderTreeAndRemoveNode({
@@ -545,9 +545,33 @@ describe('<Tree />', () => {
 
       // Only adding back Node 1 will remove the cloned node
       expect(getByText('1')).not.toBeNull();
-      expect(scrollToNode).toHaveBeenCalledTimes(0);
+      expect(scrollToNode).toHaveBeenCalledTimes(1);
     });
 
+    it.each`
+      keyPressed
+      ${'ArrowUp'}
+      ${'ArrowDown'}
+      ${'ArrowLeft'}
+      ${'ArrowRight'}
+    `(
+      'should call scrollToNode when active node is not in the DOM and $kepPressed key pressed',
+      async ({ keyPressed }) => {
+        expect.assertions(2);
+        const scrollToNode = jest.fn();
+        const { getByText } = await renderTreeAndRemoveNode({
+          scrollToNode,
+          setNodeOpen: jest.fn(),
+        });
+
+        await userEvent.keyboard(`{${keyPressed}}`);
+
+        // Only adding back Node 1 will remove the cloned node
+        expect(getByText('1')).not.toBeNull();
+        expect(scrollToNode).toHaveBeenCalledTimes(1);
+      }
+    );
+
     it('should remove the cloned node when the real node added back', async () => {
       expect.assertions(5);
       const scrollToNode = jest.fn();

src/components/Tree/Tree.tsx

@@ -14,11 +14,11 @@ import { TreeIdNodeMap, Props, TreeContextValue, TreeNodeId, TreeRefObject } fro
 import './Tree.style.scss';
 import {
   convertNestedTree2MappedTree,
+  isActiveNodeInDOM,
   getFistActiveNode,
   getNextActiveNode,
   getNodeDOMId,
   getTreeRootId,
-  isActiveNodeInDOM,
   migrateTreeState,
   toggleTreeNodeRecord,
   TreeContext,
@@ -35,7 +35,7 @@ const Tree = forwardRef((props: Props, ref: ForwardedRef<TreeRefObject>) => {
     id,
     style,
     children,
-    shouldNodeFocusBeInset,
+    shouldNodeFocusBeInset = DEFAULTS.SHOULD_NODE_FOCUS_BE_INSET,
     treeStructure,
     isRenderedFlat = DEFAULTS.IS_RENDERED_FLAT,
     excludeTreeRoot = DEFAULTS.EXCLUDE_TREE_ROOT,
@@ -60,10 +60,11 @@ const Tree = forwardRef((props: Props, ref: ForwardedRef<TreeRefObject>) => {
     isRequired,
   });
   const [tree, setTree] = useState<TreeIdNodeMap>(convertNestedTree2MappedTree(treeStructure));
-  const [activeNodeId, setActiveNodeId] = useState<TreeNodeId | undefined>(
+  const [activeNode, setActiveNode] = useState<TreeNodeId | undefined>(
     getFistActiveNode(tree, excludeTreeRoot)
   );
   const [isFocusWithin, setIsFocusWithin] = useState(false);
+  const activeNodeIdRef = useRef<TreeNodeId>(activeNode);
 
   const previousTree = usePrevious(tree);
 
@@ -72,7 +73,7 @@ const Tree = forwardRef((props: Props, ref: ForwardedRef<TreeRefObject>) => {
     migrateTreeState(previousTree, newTree);
     setTree(newTree);
     // Find the closest node to the last active node in the new tree
-    let newActiveNodeId = activeNodeId;
+    let newActiveNodeId = activeNode;
     while (newActiveNodeId) {
       if (newTree.has(newActiveNodeId) && !newTree.get(newActiveNodeId).isHidden) break;
       newActiveNodeId = previousTree.get(newActiveNodeId)?.parent;
@@ -86,17 +87,30 @@ const Tree = forwardRef((props: Props, ref: ForwardedRef<TreeRefObject>) => {
 
   const isVirtualTree = virtualTreeConnector !== undefined;
 
+  const scrollToVTreeNode = useCallback(() => {
+    if (isVirtualTree && virtualTreeConnector && !isActiveNodeInDOM(treeRef, activeNode)) {
+      virtualTreeConnector.scrollToNode?.(activeNode);
+    }
+  }, [isVirtualTree, virtualTreeConnector, activeNode]);
+
   // Handle DOM changes for virtual tree
-  useVirtualTreeNavigation({ virtualTreeConnector, treeRef: treeRef, activeNodeId });
+  useVirtualTreeNavigation({ virtualTreeConnector, treeRef: treeRef, activeNodeIdRef });
+
+  const setActiveNodeId = useCallback(
+    (newNodeId) => {
+      activeNodeIdRef.current = newNodeId;
+      setActiveNode(newNodeId);
+      scrollToVTreeNode();
+    },
+    [isVirtualTree, setActiveNode]
+  );
 
   const toggleTreeNode = useCallback(
     async (id: TreeNodeId, isOpen?: boolean): Promise<void> => {
       const newOpenState = isOpen !== undefined ? isOpen : !tree.get(id).isOpen;
 
       if (isVirtualTree) {
-        if (!isActiveNodeInDOM(treeRef, activeNodeId)) {
-          virtualTreeConnector.scrollToNode?.(id);
-        }
+        scrollToVTreeNode();
         await virtualTreeConnector.setNodeOpen?.(id, newOpenState);
       }
 
@@ -115,10 +129,10 @@ const Tree = forwardRef((props: Props, ref: ForwardedRef<TreeRefObject>) => {
 
       const parent = node.parent && tree.get(node.parent);
       const isRoot = parent === undefined;
-
       // tabindex depends on the treeNodeBase params as well
       const nodeProps = {
         id: getNodeDOMId(node.id),
+
         'aria-setsize': isRoot ? 1 : parent.children.length,
         'aria-level': node.level + (excludeTreeRoot ? 0 : 1),
         'aria-posinset': node.index + 1,
@@ -152,15 +166,15 @@ const Tree = forwardRef((props: Props, ref: ForwardedRef<TreeRefObject>) => {
       getNodeDetails,
       isRenderedFlat,
       shouldNodeFocusBeInset,
-      activeNodeId,
+      activeNodeId: activeNode,
       setActiveNodeId,
       toggleTreeNode,
       selectableNodes,
       itemSelection,
-      isFocusWithin: isFocusWithin,
+      isFocusWithin,
     }),
     [
-      activeNodeId,
+      activeNode,
       getNodeAriaProps,
       getNodeDetails,
       isRenderedFlat,
@@ -186,7 +200,10 @@ const Tree = forwardRef((props: Props, ref: ForwardedRef<TreeRefObject>) => {
   );
 
   const { focusWithinProps } = useFocusWithin({
-    onFocusWithinChange: setIsFocusWithin,
+    onFocusWithinChange: (val) => {
+      setIsFocusWithin(val);
+      scrollToVTreeNode();
+    },
   });
 
   const { keyboardProps } = useKeyboard({
@@ -198,26 +215,20 @@ const Tree = forwardRef((props: Props, ref: ForwardedRef<TreeRefObject>) => {
         case 'ArrowRight':
         case 'ArrowLeft': {
           evt.preventDefault();
-          if (activeNodeId) {
-            const next = getNextActiveNode(
-              tree,
-              activeNodeId,
-              key,
-              excludeTreeRoot,
-              toggleTreeNode
-            );
+          if (activeNode) {
+            const next = getNextActiveNode(tree, activeNode, key, excludeTreeRoot, toggleTreeNode);
             setActiveNodeId(next);
           }
           break;
         }
         case 'Space': // Space
           if (
             selectionMode !== 'none' &&
-            activeNodeId &&
-            (selectableNodes === 'any' || tree.get(activeNodeId)?.isLeaf)
+            activeNode &&
+            (selectableNodes === 'any' || tree.get(activeNode)?.isLeaf)
           ) {
             evt.preventDefault();
-            itemSelection.toggle(activeNodeId);
+            itemSelection.toggle(activeNode);
           }
           break;
 

src/components/Tree/Tree.types.ts

@@ -204,9 +204,12 @@ export interface Props
  */
 export interface UseVirtualTreeNavigationProps extends Pick<Props, 'virtualTreeConnector'> {
   /**
-   * The active node id in the tree.
+   * Reference to the active node id in the tree.
+   *
+   * @remarks This prevent to destroy and re-create MutationObserver every time when active node
+   * changes. Also, this solves the problem of missed mutations.
    */
-  activeNodeId: TreeNodeId;
+  activeNodeIdRef: MutableRefObject<TreeNodeId>;
   /**
    * The reference of the tree DOM element.
    */

src/components/Tree/Tree.utils.test.tsx

@@ -21,7 +21,6 @@ import {
 } from './Tree.types';
 import { createTreeNode as tNode } from './test.utils';
 import { renderHook } from '@testing-library/react-hooks';
-import tree from './Tree';
 
 const createSingleLevelTree = () =>
   // prettier-ignore