From da59e031ae91802a2162ab0b16558f462da953ed Mon Sep 17 00:00:00 2001
From: "e.muhamethanov" <eldarmu.1m@gmail.com>
Date: Mon, 21 Oct 2024 14:21:39 +0300
Subject: [PATCH 1/2] fix(useFloatingWithInteractions): improve restoreFocus,
 add restoreFocus to anchor element

---
 .../src/components/FocusTrap/FocusTrap.tsx    | 121 +++++++++++-------
 .../vkui/src/components/Popover/Popover.tsx   |   2 +-
 .../vkui/src/components/Popover/Readme.md     |   1 +
 .../vkui/src/lib/floating/types/component.ts  |   3 +-
 .../useFloatingWithInteractions/types.ts      |   4 +-
 .../useFloatingWithInteractions.ts            |  18 ++-
 6 files changed, 97 insertions(+), 52 deletions(-)

diff --git a/packages/vkui/src/components/FocusTrap/FocusTrap.tsx b/packages/vkui/src/components/FocusTrap/FocusTrap.tsx
index dad2e5be8c..1c25b9eb18 100644
--- a/packages/vkui/src/components/FocusTrap/FocusTrap.tsx
+++ b/packages/vkui/src/components/FocusTrap/FocusTrap.tsx
@@ -1,9 +1,10 @@
 'use client';
 
-import { type AllHTMLAttributes, useCallback, useRef, useState } from 'react';
+import { type AllHTMLAttributes, type MutableRefObject, useRef, useState } from 'react';
 import { arraysEquals } from '../../helpers/array';
 import { useExternRef } from '../../hooks/useExternRef';
 import { useMutationObserver } from '../../hooks/useMutationObserver';
+import { useStableCallback } from '../../hooks/useStableCallback';
 import { FOCUSABLE_ELEMENTS_LIST, Keys, pressedKey } from '../../lib/accessibility';
 import {
   contains,
@@ -21,7 +22,7 @@ export interface FocusTrapProps<T extends HTMLElement = HTMLElement>
     HasRootRef<T>,
     HasComponent {
   autoFocus?: boolean | 'root';
-  restoreFocus?: boolean | (() => boolean);
+  restoreFocus?: boolean | (() => boolean | HTMLElement);
   mount?: boolean;
   timeout?: number;
   onClose?: () => void;
@@ -31,6 +32,71 @@ export interface FocusTrapProps<T extends HTMLElement = HTMLElement>
   disabled?: boolean;
 }
 
+const useRestoreFocus = ({
+  restoreFocus,
+  timeout,
+  mount,
+  ref,
+}: Pick<FocusTrapProps, 'restoreFocus' | 'timeout' | 'mount'> & {
+  ref: MutableRefObject<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],
+  );
+};
+
 /**
  * @see https://vkcom.github.io/VKUI/#/FocusTrap
  */
@@ -51,8 +117,6 @@ export const FocusTrap = <T extends HTMLElement = HTMLElement>({
 
   const focusableNodesRef = useRef<HTMLElement[]>([]);
 
-  const [restoreFocusTo, setRestoreFocusTo] = useState<Element | null>(null);
-
   const focusNodeByIndex = (nodeIndex: number) => {
     const element = focusableNodesRef.current[nodeIndex];
 
@@ -63,6 +127,13 @@ export const FocusTrap = <T extends HTMLElement = HTMLElement>({
     }
   };
 
+  useRestoreFocus({
+    restoreFocus,
+    mount,
+    timeout,
+    ref,
+  });
+
   const recalculateFocusableNodesRef = (parentNode: HTMLElement) => {
     // eslint-disable-next-line no-restricted-properties
     const newFocusableElements = parentNode.querySelectorAll<HTMLElement>(FOCUSABLE_ELEMENTS);
@@ -133,48 +204,6 @@ export const FocusTrap = <T extends HTMLElement = HTMLElement>({
     [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(() => {
     if (!ref.current) {
       return;
diff --git a/packages/vkui/src/components/Popover/Popover.tsx b/packages/vkui/src/components/Popover/Popover.tsx
index 0d163461df..6a59534fff 100644
--- a/packages/vkui/src/components/Popover/Popover.tsx
+++ b/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}
diff --git a/packages/vkui/src/components/Popover/Readme.md b/packages/vkui/src/components/Popover/Readme.md
index d8aa7ee32d..1668c3b32b 100644
--- a/packages/vkui/src/components/Popover/Readme.md
+++ b/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">
         Наведи на меня
diff --git a/packages/vkui/src/lib/floating/types/component.ts b/packages/vkui/src/lib/floating/types/component.ts
index e8a47d7180..0c38e73083 100644
--- a/packages/vkui/src/lib/floating/types/component.ts
+++ b/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 заданный по умолчанию.
    */
diff --git a/packages/vkui/src/lib/floating/useFloatingWithInteractions/types.ts b/packages/vkui/src/lib/floating/useFloatingWithInteractions/types.ts
index 439c19e2ec..0f8db71db4 100644
--- a/packages/vkui/src/lib/floating/useFloatingWithInteractions/types.ts
+++ b/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;
 }
diff --git a/packages/vkui/src/lib/floating/useFloatingWithInteractions/useFloatingWithInteractions.ts b/packages/vkui/src/lib/floating/useFloatingWithInteractions/useFloatingWithInteractions.ts
index 22e4780b42..f2de4cd319 100644
--- a/packages/vkui/src/lib/floating/useFloatingWithInteractions/useFloatingWithInteractions.ts
+++ b/packages/vkui/src/lib/floating/useFloatingWithInteractions/useFloatingWithInteractions.ts
@@ -229,9 +229,21 @@ export const useFloatingWithInteractions = <T extends HTMLElement = HTMLElement>
     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(() => {

From 5b9d4ea47853bcef4f820ae43ab9c63f4316dcf5 Mon Sep 17 00:00:00 2001
From: "e.muhamethanov" <eldarmu.1m@gmail.com>
Date: Tue, 5 Nov 2024 13:17:53 +0300
Subject: [PATCH 2/2] fix(useFocusTrap): resolve conflicts

---
 packages/vkui/src/hooks/useFocusTrap.ts | 121 +++++++++++++++---------
 1 file changed, 75 insertions(+), 46 deletions(-)

diff --git a/packages/vkui/src/hooks/useFocusTrap.ts b/packages/vkui/src/hooks/useFocusTrap.ts
index 0f3422ba77..a79c1684e7 100644
--- a/packages/vkui/src/hooks/useFocusTrap.ts
+++ b/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) {