From 3cf4f2d987f7e6ad22cb270773128b06e26c25fb Mon Sep 17 00:00:00 2001
From: Michael Leibman <michael.leibman@databricks.com>
Date: Tue, 8 Oct 2024 16:54:34 -0700
Subject: [PATCH] Only register a single set of event listeners for the entire
 document when tracking whether a keyboard or a mouse is being used.

---
 packages/react/menu/src/Menu.tsx              | 27 ++--------
 .../react/menu/src/useIsUsingKeyboard.tsx     | 49 +++++++++++++++++++
 2 files changed, 54 insertions(+), 22 deletions(-)
 create mode 100644 packages/react/menu/src/useIsUsingKeyboard.tsx

diff --git a/packages/react/menu/src/Menu.tsx b/packages/react/menu/src/Menu.tsx
index d8f340689..57040e6e2 100644
--- a/packages/react/menu/src/Menu.tsx
+++ b/packages/react/menu/src/Menu.tsx
@@ -19,6 +19,7 @@ import { Slot } from '@radix-ui/react-slot';
 import { useCallbackRef } from '@radix-ui/react-use-callback-ref';
 import { hideOthers } from 'aria-hidden';
 import { RemoveScroll } from 'react-remove-scroll';
+import { useIsUsingKeyboard } from './useIsUsingKeyboard';
 
 import type { Scope } from '@radix-ui/react-context';
 
@@ -69,7 +70,6 @@ const [MenuProvider, useMenuContext] = createMenuContext<MenuContextValue>(MENU_
 
 type MenuRootContextValue = {
   onClose(): void;
-  isUsingKeyboardRef: React.RefObject<boolean>;
   dir: Direction;
   modal: boolean;
 };
@@ -88,27 +88,9 @@ const Menu: React.FC<MenuProps> = (props: ScopedProps<MenuProps>) => {
   const { __scopeMenu, open = false, children, dir, onOpenChange, modal = true } = props;
   const popperScope = usePopperScope(__scopeMenu);
   const [content, setContent] = React.useState<MenuContentElement | null>(null);
-  const isUsingKeyboardRef = React.useRef(false);
   const handleOpenChange = useCallbackRef(onOpenChange);
   const direction = useDirection(dir);
 
-  React.useEffect(() => {
-    // Capture phase ensures we set the boolean before any side effects execute
-    // in response to the key or pointer event as they might depend on this value.
-    const handleKeyDown = () => {
-      isUsingKeyboardRef.current = true;
-      document.addEventListener('pointerdown', handlePointer, { capture: true, once: true });
-      document.addEventListener('pointermove', handlePointer, { capture: true, once: true });
-    };
-    const handlePointer = () => (isUsingKeyboardRef.current = false);
-    document.addEventListener('keydown', handleKeyDown, { capture: true });
-    return () => {
-      document.removeEventListener('keydown', handleKeyDown, { capture: true });
-      document.removeEventListener('pointerdown', handlePointer, { capture: true });
-      document.removeEventListener('pointermove', handlePointer, { capture: true });
-    };
-  }, []);
-
   return (
     <PopperPrimitive.Root {...popperScope}>
       <MenuProvider
@@ -121,7 +103,6 @@ const Menu: React.FC<MenuProps> = (props: ScopedProps<MenuProps>) => {
         <MenuRootProvider
           scope={__scopeMenu}
           onClose={React.useCallback(() => handleOpenChange(false), [handleOpenChange])}
-          isUsingKeyboardRef={isUsingKeyboardRef}
           dir={direction}
           modal={modal}
         >
@@ -386,6 +367,7 @@ const MenuContentImpl = React.forwardRef<MenuContentImplElement, MenuContentImpl
     const pointerGraceIntentRef = React.useRef<GraceIntent | null>(null);
     const pointerDirRef = React.useRef<Side>('right');
     const lastPointerXRef = React.useRef(0);
+    const getIsUsingKeyboard = useIsUsingKeyboard();
 
     const ScrollLockWrapper = disableOutsideScroll ? RemoveScroll : React.Fragment;
     const scrollLockWrapperProps = disableOutsideScroll
@@ -490,7 +472,7 @@ const MenuContentImpl = React.forwardRef<MenuContentImplElement, MenuContentImpl
                 onCurrentTabStopIdChange={setCurrentItemId}
                 onEntryFocus={composeEventHandlers(onEntryFocus, (event) => {
                   // only focus first item when using keyboard
-                  if (!rootContext.isUsingKeyboardRef.current) event.preventDefault();
+                  if (!getIsUsingKeyboard()) event.preventDefault();
                 })}
                 preventScrollOnEntryFocus
               >
@@ -1167,6 +1149,7 @@ const MenuSubContent = React.forwardRef<MenuSubContentElement, MenuSubContentPro
     const subContext = useMenuSubContext(SUB_CONTENT_NAME, props.__scopeMenu);
     const ref = React.useRef<MenuSubContentElement>(null);
     const composedRefs = useComposedRefs(forwardedRef, ref);
+    const getIsUsingKeyboard = useIsUsingKeyboard();
     return (
       <Collection.Provider scope={props.__scopeMenu}>
         <Presence present={forceMount || context.open}>
@@ -1183,7 +1166,7 @@ const MenuSubContent = React.forwardRef<MenuSubContentElement, MenuSubContentPro
               trapFocus={false}
               onOpenAutoFocus={(event) => {
                 // when opening a submenu, focus content for keyboard users only
-                if (rootContext.isUsingKeyboardRef.current) ref.current?.focus();
+                if (getIsUsingKeyboard()) ref.current?.focus();
                 event.preventDefault();
               }}
               // The menu might close because of focusing another menu item in the parent menu. We
diff --git a/packages/react/menu/src/useIsUsingKeyboard.tsx b/packages/react/menu/src/useIsUsingKeyboard.tsx
new file mode 100644
index 000000000..37caf2ad0
--- /dev/null
+++ b/packages/react/menu/src/useIsUsingKeyboard.tsx
@@ -0,0 +1,49 @@
+import * as React from 'react';
+
+let isUsingKeyboard = false;
+let isUsingKeyboardSubscriberCount = 0;
+
+function handleKeydown() {
+  isUsingKeyboard = true;
+}
+
+function handlePointer() {
+  isUsingKeyboard = false;
+}
+
+function getIsUsingKeyboard() {
+  return isUsingKeyboard;
+}
+
+function subscribeShared() {
+  if (isUsingKeyboardSubscriberCount++ === 0) {
+    // Capture phase ensures we set the boolean before any side effects execute
+    // in response to the key or pointer event as they might depend on this value.
+    document.addEventListener('keydown', handleKeydown, { capture: true });
+    document.addEventListener('pointerdown', handlePointer, { capture: true });
+    document.addEventListener('pointermove', handlePointer, { capture: true });
+  }
+}
+
+function unsubscribeShared() {
+  if (--isUsingKeyboardSubscriberCount === 0) {
+    document.removeEventListener('keydown', handleKeydown, { capture: true });
+    document.removeEventListener('pointerdown', handlePointer, { capture: true });
+    document.removeEventListener('pointermove', handlePointer, { capture: true });
+  }
+}
+
+/**
+ * Starts tracking whether the user is using a keyboard or a mouse.
+ * This implementation keeps track of how many subscribers it has and makes sure
+ * only one set of listeners is attached to the document to improve performance and
+ * prevent memory leaks.
+ */
+export function useIsUsingKeyboard() {
+  React.useEffect(() => {
+    subscribeShared();
+    return unsubscribeShared;
+  }, []);
+
+  return getIsUsingKeyboard;
+}