Refactor Virtualizer to improve performance, stability, and complexity (#6451)

Author: devongovett

package.json

@@ -74,6 +74,7 @@
     "@babel/preset-react": "^7.24.1",
     "@babel/preset-typescript": "^7.24.1",
     "@babel/register": "^7.23.7",
+    "@faker-js/faker": "^8.4.1",
     "@octokit/rest": "*",
     "@parcel/bundler-library": "2.11.1-dev.3224",
     "@parcel/optimizer-data-url": "2.0.0-dev.1601",

packages/@react-aria/focus/src/FocusScope.tsx

@@ -65,6 +65,7 @@ interface IFocusContext {
 }
 
 const FocusContext = React.createContext<IFocusContext | null>(null);
+const RESTORE_FOCUS_EVENT = 'react-aria-focus-scope-restore';
 
 let activeScope: ScopeRef = null;
 
@@ -117,12 +118,21 @@ export function FocusScope(props: FocusScopeProps) {
     // Find all rendered nodes between the sentinels and add them to the scope.
     let node = startRef.current?.nextSibling!;
     let nodes: Element[] = [];
+    let stopPropagation = e => e.stopPropagation();
     while (node && node !== endRef.current) {
       nodes.push(node as Element);
+      // Stop custom restore focus event from propagating to parent focus scopes.
+      node.addEventListener(RESTORE_FOCUS_EVENT, stopPropagation);
       node = node.nextSibling as Element;
     }
 
     scopeRef.current = nodes;
+
+    return () => {
+      for (let node of nodes) {
+        node.removeEventListener(RESTORE_FOCUS_EVENT, stopPropagation);
+      }
+    };
   }, [children]);
 
   useActiveScopeTracker(scopeRef, restoreFocus, contain);
@@ -470,7 +480,7 @@ function focusElement(element: FocusableElement | null, scroll = false) {
   }
 }
 
-function focusFirstInScope(scope: Element[], tabbable:boolean = true) {
+function getFirstInScope(scope: Element[], tabbable = true) {
   let sentinel = scope[0].previousElementSibling!;
   let scopeRoot = getScopeRoot(scope);
   let walker = getFocusableTreeWalker(scopeRoot, {tabbable}, scope);
@@ -485,7 +495,11 @@ function focusFirstInScope(scope: Element[], tabbable:boolean = true) {
     nextNode = walker.nextNode();
   }
 
-  focusElement(nextNode as FocusableElement);
+  return nextNode as FocusableElement;
+}
+
+function focusFirstInScope(scope: Element[], tabbable:boolean = true) {
+  focusElement(getFirstInScope(scope, tabbable));
 }
 
 function useAutoFocus(scopeRef: RefObject<Element[]>, autoFocus?: boolean) {
@@ -692,7 +706,7 @@ function useRestoreFocus(scopeRef: RefObject<Element[]>, restoreFocus?: boolean,
             let treeNode = clonedTree.getTreeNode(scopeRef);
             while (treeNode) {
               if (treeNode.nodeToRestore && treeNode.nodeToRestore.isConnected) {
-                focusElement(treeNode.nodeToRestore);
+                restoreFocusToElement(treeNode.nodeToRestore);
                 return;
               }
               treeNode = treeNode.parent;
@@ -703,7 +717,8 @@ function useRestoreFocus(scopeRef: RefObject<Element[]>, restoreFocus?: boolean,
             treeNode = clonedTree.getTreeNode(scopeRef);
             while (treeNode) {
               if (treeNode.scopeRef && treeNode.scopeRef.current && focusScopeTree.getTreeNode(treeNode.scopeRef)) {
-                focusFirstInScope(treeNode.scopeRef.current, true);
+                let node = getFirstInScope(treeNode.scopeRef.current, true);
+                restoreFocusToElement(node);
                 return;
               }
               treeNode = treeNode.parent;
@@ -715,6 +730,15 @@ function useRestoreFocus(scopeRef: RefObject<Element[]>, restoreFocus?: boolean,
   }, [scopeRef, restoreFocus]);
 }
 
+function restoreFocusToElement(node: FocusableElement) {
+  // Dispatch a custom event that parent elements can intercept to customize focus restoration.
+  // For example, virtualized collection components reuse DOM elements, so the original element
+  // might still exist in the DOM but representing a different item.
+  if (node.dispatchEvent(new CustomEvent(RESTORE_FOCUS_EVENT, {bubbles: true, cancelable: true}))) {
+    focusElement(node);
+  }
+}
+
 /**
  * Create a [TreeWalker]{@link https://developer.mozilla.org/en-US/docs/Web/API/TreeWalker}
  * that matches all focusable/tabbable elements.

packages/@react-aria/focus/test/FocusScope.test.js

@@ -19,6 +19,7 @@ import {Provider} from '@react-spectrum/provider';
 import React, {useEffect, useState} from 'react';
 import ReactDOM from 'react-dom';
 import {Example as StorybookExample} from '../stories/FocusScope.stories';
+import {useEvent} from '@react-aria/utils';
 import userEvent from '@testing-library/user-event';
 
 
@@ -764,6 +765,70 @@ describe('FocusScope', function () {
       expect(document.activeElement).toBe(button2);
       expect(input1).not.toBeInTheDocument();
     });
+
+    it('should allow restoration to be overridden with a custom event', async function () {
+      function Test() {
+        let [show, setShow] = React.useState(false);
+        let ref = React.useRef(null);
+        useEvent(ref, 'react-aria-focus-scope-restore', e => {
+          e.preventDefault();
+        });
+
+        return (
+          <div ref={ref}>
+            <button onClick={() => setShow(true)}>Show</button>
+            {show && <FocusScope restoreFocus>
+              <input autoFocus onKeyDown={() => setShow(false)} />
+            </FocusScope>}
+          </div>
+        );
+      }
+
+      let {getByRole} = render(<Test />);
+      let button = getByRole('button');
+      await user.click(button);
+
+      let input = getByRole('textbox');
+      expect(document.activeElement).toBe(input);
+
+      await user.keyboard('{Escape}');
+      act(() => jest.runAllTimers());
+      expect(input).not.toBeInTheDocument();
+      expect(document.activeElement).toBe(document.body);
+    });
+
+    it('should not bubble focus scope restoration event out of nested focus scopes', async function () {
+      function Test() {
+        let [show, setShow] = React.useState(false);
+        let ref = React.useRef(null);
+        useEvent(ref, 'react-aria-focus-scope-restore', e => {
+          e.preventDefault();
+        });
+
+        return (
+          <div ref={ref}>
+            <FocusScope>
+              <button onClick={() => setShow(true)}>Show</button>
+              {show && <FocusScope restoreFocus>
+                <input autoFocus onKeyDown={() => setShow(false)} />
+              </FocusScope>}
+            </FocusScope>
+          </div>
+        );
+      }
+
+      let {getByRole} = render(<Test />);
+      let button = getByRole('button');
+      await user.click(button);
+
+      let input = getByRole('textbox');
+      expect(document.activeElement).toBe(input);
+
+      await user.keyboard('{Escape}');
+      act(() => jest.runAllTimers());
+      expect(input).not.toBeInTheDocument();
+      expect(document.activeElement).toBe(button);
+    });
   });
 
   describe('auto focus', function () {

packages/@react-aria/selection/src/useSelectableCollection.ts

@@ -293,6 +293,7 @@ export function useSelectableCollection(options: AriaSelectableCollectionOptions
   };
 
   // Store the scroll position so we can restore it later.
+  /// TODO: should this happen all the time??
   let scrollPos = useRef({top: 0, left: 0});
   useEvent(scrollRef, 'scroll', isVirtualized ? null : () => {
     scrollPos.current = {
@@ -342,7 +343,7 @@ export function useSelectableCollection(options: AriaSelectableCollectionOptions
       scrollRef.current.scrollLeft = scrollPos.current.left;
     }
 
-    if (!isVirtualized && manager.focusedKey != null) {
+    if (manager.focusedKey != null) {
       // Refocus and scroll the focused item into view if it exists within the scrollable region.
       let element = scrollRef.current.querySelector(`[data-key="${CSS.escape(manager.focusedKey.toString())}"]`) as HTMLElement;
       if (element) {
@@ -400,17 +401,21 @@ export function useSelectableCollection(options: AriaSelectableCollectionOptions
   // eslint-disable-next-line react-hooks/exhaustive-deps
   }, []);
 
-  // If not virtualized, scroll the focused element into view when the focusedKey changes.
-  // When virtualized, Virtualizer handles this internally.
+  // Scroll the focused element into view when the focusedKey changes.
   let lastFocusedKey = useRef(manager.focusedKey);
   useEffect(() => {
-    let modality = getInteractionModality();
-    if (manager.isFocused && manager.focusedKey != null && scrollRef?.current) {
-      let element = scrollRef.current.querySelector(`[data-key="${CSS.escape(manager.focusedKey.toString())}"]`) as HTMLElement;
-      if (element && (modality === 'keyboard' || autoFocusRef.current)) {
-        if (!isVirtualized) {
-          scrollIntoView(scrollRef.current, element);
-        }
+    if (manager.isFocused && manager.focusedKey != null && manager.focusedKey !== lastFocusedKey.current && scrollRef?.current) {
+      let modality = getInteractionModality();
+      let element = ref.current.querySelector(`[data-key="${CSS.escape(manager.focusedKey.toString())}"]`) as HTMLElement;
+      if (!element) {
+        // If item element wasn't found, return early (don't update autoFocusRef and lastFocusedKey).
+        // The collection may initially be empty (e.g. virtualizer), so wait until the element exists.
+        return;
+      }
+
+      if (modality === 'keyboard' || autoFocusRef.current) {
+        scrollIntoView(scrollRef.current, element);
+
         // Avoid scroll in iOS VO, since it may cause overlay to close (i.e. RAC submenu)
         if (modality !== 'virtual') {
           scrollIntoViewport(element, {containingElement: ref.current});
@@ -425,7 +430,13 @@ export function useSelectableCollection(options: AriaSelectableCollectionOptions
 
     lastFocusedKey.current = manager.focusedKey;
     autoFocusRef.current = false;
-  }, [isVirtualized, scrollRef, manager.focusedKey, manager.isFocused, ref]);
+  });
+
+  // Intercept FocusScope restoration since virtualized collections can reuse DOM nodes.
+  useEvent(ref, 'react-aria-focus-scope-restore', e => {
+    e.preventDefault();
+    manager.setFocused(true);
+  });
 
   let handlers = {
     onKeyDown,

packages/@react-aria/utils/src/platform.ts

@@ -26,40 +26,54 @@ function testPlatform(re: RegExp) {
     : false;
 }
 
-export function isMac() {
-  return testPlatform(/^Mac/i);
+function cached(fn: () => boolean) {
+  if (process.env.NODE_ENV === 'test') {
+    return fn;
+  }
+  
+  let res: boolean | null = null;
+  return () => {
+    if (res == null) {
+      res = fn();
+    }
+    return res;
+  };
 }
 
-export function isIPhone() {
+export const isMac = cached(function () {
+  return testPlatform(/^Mac/i);
+});
+
+export const isIPhone = cached(function () {
   return testPlatform(/^iPhone/i);
-}
+});
 
-export function isIPad() {
+export const isIPad = cached(function () {
   return testPlatform(/^iPad/i) ||
     // iPadOS 13 lies and says it's a Mac, but we can distinguish by detecting touch support.
     (isMac() && navigator.maxTouchPoints > 1);
-}
+});
 
-export function isIOS() {
+export const isIOS = cached(function () {
   return isIPhone() || isIPad();
-}
+});
 
-export function isAppleDevice() {
+export const isAppleDevice = cached(function () {
   return isMac() || isIOS();
-}
+});
 
-export function isWebKit() {
+export const isWebKit = cached(function () {
   return testUserAgent(/AppleWebKit/i) && !isChrome();
-}
+});
 
-export function isChrome() {
+export const isChrome = cached(function () {
   return testUserAgent(/Chrome/i);
-}
+});
 
-export function isAndroid() {
+export const isAndroid = cached(function () {
   return testUserAgent(/Android/i);
-}
+});
 
-export function isFirefox() {
+export const isFirefox = cached(function () {
   return testUserAgent(/Firefox/i);
-}
+});

packages/@react-aria/utils/src/useEvent.ts

@@ -15,7 +15,7 @@ import {useEffectEvent} from './useEffectEvent';
 
 export function useEvent<K extends keyof GlobalEventHandlersEventMap>(
   ref: RefObject<EventTarget>,
-  event: K,
+  event: K | (string & {}),
   handler?: (this: Document, ev: GlobalEventHandlersEventMap[K]) => any,
   options?: boolean | AddEventListenerOptions
 ) {

packages/@react-aria/virtualizer/src/ScrollView.tsx

@@ -155,7 +155,7 @@ function ScrollView(props: ScrollViewProps, ref: RefObject<HTMLDivElement>) {
     updateSize();
   }, [updateSize]);
   let raf = useRef<ReturnType<typeof requestAnimationFrame> | null>();
-  let onResize = () => {
+  let onResize = useCallback(() => {
     if (isOldReact) {
       raf.current ??= requestAnimationFrame(() => {
         updateSize();
@@ -164,7 +164,7 @@ function ScrollView(props: ScrollViewProps, ref: RefObject<HTMLDivElement>) {
     } else {
       updateSize();
     }
-  };
+  }, [updateSize]);
   useResizeObserver({ref, onResize});
   useEffect(() => {
     return () => {