fix(useFloatingWithInteractions): improve restoreFocus, add restoreFocus to anchor element

packages/vkui/src/components/Popover/Popover.tsx

@@ -269,7 +269,7 @@ export const Popover = ({
             mount={!hidden}
             disabled={hidden}
             autoFocus={disableInteractive ? false : autoFocus}
-            restoreFocus={restoreFocus ? onRestoreFocus : false}
+            restoreFocus={restoreFocus ? () => onRestoreFocus(restoreFocus) : false}
             onClose={onEscapeKeyDown}
           >
             {arrow}

packages/vkui/src/components/Popover/Readme.md

@@ -39,6 +39,7 @@ const PopoverWithTriggerHover = () => {
           <Text>Привет</Text>
         </Div>
       }
+      restoreFocus="anchor-element"
     >
       <Button id="tooltip-1" mode="outline">
         Наведи на меня

packages/vkui/src/hooks/useFocusTrap.ts

@@ -1,4 +1,4 @@
-import { type RefObject, useCallback, useRef, useState } from 'react';
+import { type RefObject, useRef, useState } from 'react';
 import { arraysEquals } from '../helpers/array';
 import { FOCUSABLE_ELEMENTS_LIST, Keys, pressedKey } from '../lib/accessibility';
 import {
@@ -10,6 +10,72 @@ import {
 } from '../lib/dom';
 import { useIsomorphicLayoutEffect } from '../lib/useIsomorphicLayoutEffect';
 import { useMutationObserver } from './useMutationObserver';
+import { useStableCallback } from './useStableCallback';
+
+const useRestoreFocus = ({
+  restoreFocus,
+  timeout,
+  mount,
+  ref,
+}: Pick<UseFocusTrapProps, 'restoreFocus' | 'timeout' | 'mount'> & {
+  ref: RefObject<HTMLElement | null>;
+}) => {
+  const restoreFocusRef = useRef(restoreFocus);
+  restoreFocusRef.current = restoreFocus;
+  const [restoreFocusTo, setRestoreFocusTo] = useState<Element | null>(null);
+
+  const restoreFocusImpl = useStableCallback(() => {
+    const shouldRestoreFocus =
+      typeof restoreFocusRef.current === 'function'
+        ? restoreFocusRef.current()
+        : restoreFocusRef.current;
+
+    if (!shouldRestoreFocus) {
+      return;
+    }
+
+    setTimeout(() => {
+      const restoreFocusElement =
+        (isHTMLElement(shouldRestoreFocus) && shouldRestoreFocus) ||
+        (isHTMLElement(restoreFocusTo) && restoreFocusTo) ||
+        null;
+
+      if (restoreFocusElement) {
+        restoreFocusElement.focus();
+        setRestoreFocusTo(null);
+      }
+    }, timeout);
+  });
+
+  useIsomorphicLayoutEffect(
+    function calculateRestoreFocusTo() {
+      if (!ref.current || !restoreFocusRef.current || !mount) {
+        setRestoreFocusTo(null);
+        return;
+      }
+      setRestoreFocusTo(getActiveElementByAnotherElement(ref.current));
+    },
+    [ref, mount],
+  );
+
+  useIsomorphicLayoutEffect(
+    function tryToRestoreFocusOnUnmount() {
+      return () => {
+        restoreFocusImpl();
+      };
+    },
+    [restoreFocusImpl],
+  );
+
+  useIsomorphicLayoutEffect(
+    function tryToRestoreFocusWhenFakeUnmount() {
+      if (!mount) {
+        restoreFocusImpl();
+      }
+    },
+    [mount, restoreFocusImpl],
+  );
+};
 
 const FOCUSABLE_ELEMENTS: string = FOCUSABLE_ELEMENTS_LIST.join();
 
@@ -31,7 +97,7 @@ export type UseFocusTrapProps = {
   /**
    * @default true
    */
-  restoreFocus?: boolean | (() => boolean);
+  restoreFocus?: boolean | (() => boolean | HTMLElement);
   /**
    * @default 0
    */
@@ -60,8 +126,6 @@ export const useFocusTrap = (
 
   const focusableNodesRef = useRef<HTMLElement[]>([]);
 
-  const [restoreFocusTo, setRestoreFocusTo] = useState<Element | null>(null);
-
   const focusNodeByIndex = (nodeIndex: number) => {
     const element = focusableNodesRef.current[nodeIndex];
 
@@ -72,6 +136,13 @@ export const useFocusTrap = (
     }
   };
 
+  useRestoreFocus({
+    restoreFocus,
+    mount,
+    timeout,
+    ref,
+  });
+
   const recalculateFocusableNodesRef = (parentNode: HTMLElement) => {
     // eslint-disable-next-line no-restricted-properties
     const newFocusableElements = parentNode.querySelectorAll<HTMLElement>(FOCUSABLE_ELEMENTS);
@@ -142,48 +213,6 @@ export const useFocusTrap = (
     [autoFocus, timeout, disabled],
   );
 
-  const restoreFocusImpl = useCallback(() => {
-    const shouldRestoreFocus = typeof restoreFocus === 'function' ? restoreFocus() : restoreFocus;
-
-    if (!restoreFocusTo || !isHTMLElement(restoreFocusTo) || !shouldRestoreFocus) {
-      return;
-    }
-
-    setTimeout(() => {
-      if (restoreFocusTo) {
-        restoreFocusTo.focus();
-        setRestoreFocusTo(null);
-      }
-    }, timeout);
-  }, [restoreFocus, restoreFocusTo, timeout]);
-
-  useIsomorphicLayoutEffect(
-    function calculateRestoreFocusTo() {
-      if (!ref.current || !restoreFocus || !mount) {
-        setRestoreFocusTo(null);
-        return;
-      }
-      setRestoreFocusTo(getActiveElementByAnotherElement(ref.current));
-    },
-    [ref, mount, restoreFocus],
-  );
-
-  useIsomorphicLayoutEffect(
-    function tryToRestoreFocusOnUnmount() {
-      return () => restoreFocusImpl();
-    },
-    [restoreFocusImpl],
-  );
-
-  useIsomorphicLayoutEffect(
-    function tryToRestoreFocusWhenFakeUnmount() {
-      if (!mount) {
-        restoreFocusImpl();
-      }
-    },
-    [mount, restoreFocusImpl],
-  );
-
   useIsomorphicLayoutEffect(
     function initializeFocusTrap() {
       if (!ref.current) {

packages/vkui/src/lib/floating/types/component.ts

@@ -4,6 +4,7 @@ import type {
   UseFloatingWithInteractionsProps,
   UseFloatingWithInteractionsReturn,
 } from '../useFloatingWithInteractions';
+import { type RestoreFocusType } from '../useFloatingWithInteractions/types';
 import type { OnPlacementChange } from './common';
 
 /**
@@ -44,7 +45,7 @@ export interface FloatingComponentProps
   /**
    * Нужно ли после закрытия всплывающего элемента возвращать фокус на предыдущий активный элемент.
    */
-  restoreFocus?: boolean;
+  restoreFocus?: RestoreFocusType;
   /**
    * Перебивает zIndex заданный по умолчанию.
    */

packages/vkui/src/lib/floating/useFloatingWithInteractions/types.ts

@@ -11,6 +11,8 @@ export type ManualTriggerType = 'manual';
 
 export type TriggerType = ManualTriggerType | InteractiveTriggerType | InteractiveTriggerType[];
 
+export type RestoreFocusType = boolean | 'latest-active' | 'anchor-element' | HTMLElement;
+
 export type ShownChangeReason =
   | 'click-outside'
   | 'escape-key'
@@ -110,5 +112,5 @@ export interface UseFloatingWithInteractionsReturn<T extends HTMLElement = HTMLE
   middlewareData: UseFloatingData['middlewareData'];
   onClose: (this: void) => void;
   onEscapeKeyDown?: (this: void) => void;
-  onRestoreFocus: (this: void) => boolean;
+  onRestoreFocus: (restoreFocus?: RestoreFocusType) => boolean | HTMLElement;
 }

packages/vkui/src/lib/floating/useFloatingWithInteractions/useFloatingWithInteractions.ts

@@ -229,9 +229,21 @@
     commitShownLocalState(false, 'callback');
   }, [commitShownLocalState]);
 
-  const handleRestoreFocus = React.useCallback(
-    () => (triggerOnFocus ? blockFocusRef.current : true),
-    [triggerOnFocus],
+  const handleRestoreFocus: UseFloatingWithInteractionsReturn['onRestoreFocus'] = React.useCallback(
+    (restoreFocus = true) => {
+      if (!restoreFocus) {
+        return false;
+      }
+      if (restoreFocus === true || restoreFocus === 'latest-active') {
+        return triggerOnFocus ? blockFocusRef.current : true;
+      } else if (restoreFocus === 'anchor-element') {
+        return refs.reference.current as HTMLElement;
+      } else if (restoreFocus instanceof HTMLElement) {
+        return restoreFocus;
+      }
+      return false;
+    },
+    [refs.reference, triggerOnFocus],
   );
 
   const handleEscapeKeyDown = React.useCallback(() => {