From 3ccb2ef58bf852872c7903178bb33e700c3aa639 Mon Sep 17 00:00:00 2001
From: Mitchell Austin <mr.fye@oneandthesame.net>
Date: Mon, 24 Jul 2023 17:40:50 -0700
Subject: [PATCH] Return focus more from focus return hook (#52710)

---
 packages/compose/README.md                      |  2 +-
 .../compose/src/hooks/use-focus-return/index.js | 17 +++++++++++------
 2 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/packages/compose/README.md b/packages/compose/README.md
index 0a3fed45bce8f..62ebdef6d798e 100644
--- a/packages/compose/README.md
+++ b/packages/compose/README.md
@@ -323,7 +323,7 @@ _Returns_
 
 ### useFocusReturn
 
-When opening modals/sidebars/dialogs, the focus must move to the opened area and return to the previously focused element when closed. The current hook implements the returning behavior.
+Adds the unmount behavior of returning focus to the element which had it previously as is expected for roles like menus or dialogs.
 
 _Usage_
 
diff --git a/packages/compose/src/hooks/use-focus-return/index.js b/packages/compose/src/hooks/use-focus-return/index.js
index 66751b7028d32..2cd93b279cd31 100644
--- a/packages/compose/src/hooks/use-focus-return/index.js
+++ b/packages/compose/src/hooks/use-focus-return/index.js
@@ -3,11 +3,12 @@
  */
 import { useRef, useEffect, useCallback } from '@wordpress/element';
 
+/** @type {Element|null} */
+let origin = null;
+
 /**
- * When opening modals/sidebars/dialogs, the focus
- * must move to the opened area and return to the
- * previously focused element when closed.
- * The current hook implements the returning behavior.
+ * Adds the unmount behavior of returning focus to the element which had it
+ * previously as is expected for roles like menus or dialogs.
  *
  * @param {() => void} [onFocusReturn] Overrides the default return behavior.
  * @return {import('react').RefCallback<HTMLElement>} Element Ref.
@@ -54,6 +55,7 @@ function useFocusReturn( onFocusReturn ) {
 			);
 
 			if ( ref.current?.isConnected && ! isFocused ) {
+				origin ??= focusedBeforeMount.current;
 				return;
 			}
 
@@ -64,10 +66,13 @@ function useFocusReturn( onFocusReturn ) {
 			if ( onFocusReturnRef.current ) {
 				onFocusReturnRef.current();
 			} else {
-				/** @type {null | HTMLElement} */ (
-					focusedBeforeMount.current
+				/** @type {null|HTMLElement} */ (
+					! focusedBeforeMount.current.isConnected
+						? origin
+						: focusedBeforeMount.current
 				)?.focus();
 			}
+			origin = null;
 		}
 	}, [] );
 }