Refactor focus on mount

packages/components/src/modal/frame.js

@@ -41,9 +41,7 @@ function ModalFrameContent( {
 			}
 		}
 	}
-	const focusOnMountRef = useFocusOnMount(
-		focusOnMount ? 'container' : false
-	);
+	const focusOnMountRef = useFocusOnMount();
 	const constrainedTabbingRef = useConstrainedTabbing();
 	const focusReturnRef = useFocusReturn();
 
@@ -60,9 +58,9 @@ function ModalFrameContent( {
 				className={ classnames( 'components-modal__frame', className ) }
 				style={ style }
 				ref={ mergeRefs( [
-					focusOnMountRef,
 					constrainedTabbingRef,
 					focusReturnRef,
+					focusOnMount ? focusOnMountRef : null,
 				] ) }
 				role={ role }
 				aria-label={ contentLabel }

packages/components/src/popover/index.js

@@ -425,7 +425,7 @@ const Popover = ( {
 		containerRef,
 		focusOnMount ? constrainedTabbingRef : null,
 		focusOnMount ? focusReturnRef : null,
-		focusOnMountRef,
+		focusOnMount ? focusOnMountRef : null,
 	];
 	const mergedRefs = useCallback( mergeRefs( allRefs ), allRefs );
 

packages/compose/README.md

@@ -212,7 +212,7 @@ _Parameters_
 
 _Returns_
 
--   `(Function|Object)`: Element Ref.
+-   `Function`: Ref callback.
 
 <a name="useFocusReturn" href="#useFocusReturn">#</a> **useFocusReturn**
 

packages/compose/src/hooks/use-dialog/index.js

@@ -58,8 +58,8 @@ function useDialog( options ) {
 	return [
 		mergeRefs( [
 			constrainedTabbingRef,
-			focusOnMountRef,
 			focusReturnRef,
+			focusOnMountRef,
 			closeOnEscapeRef,
 		] ),
 		focusOutsideProps,

packages/compose/src/hooks/use-focus-on-mount/index.js

@@ -1,14 +1,14 @@
 /**
  * WordPress dependencies
  */
+import { useCallback, useRef } from '@wordpress/element';
 import { focus } from '@wordpress/dom';
-import { useCallback, useEffect, useRef, useState } from '@wordpress/element';
 
 /**
  * Hook used to focus the first tabbable element on mount.
  *
  * @param {boolean|string} focusOnMount Focus on mount mode.
- * @return {Function|Object} Element Ref.
+ * @return {Function} Ref callback.
  *
  * @example
  * ```js
@@ -25,58 +25,29 @@ import { useCallback, useEffect, useRef, useState } from '@wordpress/element';
  * }
  * ```
  */
-function useFocusOnMount( focusOnMount = 'firstElement' ) {
-	const focusOnMountRef = useRef( focusOnMount );
-	const [ node, setNode ] = useState();
-	useEffect( () => {
-		focusOnMountRef.current = focusOnMount;
-	}, [ focusOnMount ] );
+export default function useFocusOnMount( focusOnMount ) {
+	const didMount = useRef( false );
+	return useCallback( ( node ) => {
+		if ( ! node || didMount.current === true ) {
+			return;
+		}
 
-	// Ideally we should be running the focus behavior in the useCallback directly
-	// Right now we have some issues where the link popover remounts
-	// prevents us from doing that.
-	const ref = useCallback( setNode, [] );
+		didMount.current = true;
 
-	// Focus handling
-	useEffect( () => {
-		if ( ! node ) {
+		if ( node.contains( node.ownerDocument.activeElement ) ) {
 			return;
 		}
-		/*
-		 * Without the setTimeout, the dom node is not being focused. Related:
-		 * https://stackoverflow.com/questions/35522220/react-ref-with-focus-doesnt-work-without-settimeout-my-example
-		 *
-		 * TODO: Treat the cause, not the symptom.
-		 */
-		const focusTimeout = setTimeout( () => {
-			if ( focusOnMountRef.current === false || ! node ) {
-				return;
-			}
-
-			if ( focusOnMountRef.current === 'firstElement' ) {
-				const firstTabbable = focus.tabbable.find( node )[ 0 ];
 
-				if ( firstTabbable ) {
-					firstTabbable.focus();
-				} else {
-					node.focus();
-				}
+		let target = node;
 
-				return;
-			}
+		if ( focusOnMount === 'firstElement' ) {
+			const firstTabbable = focus.tabbable.find( node )[ 0 ];
 
-			if (
-				focusOnMountRef.current === 'container' ||
-				focusOnMountRef.current === true
-			) {
-				node.focus();
+			if ( firstTabbable ) {
+				target = firstTabbable;
 			}
-		}, 0 );
-
-		return () => clearTimeout( focusTimeout );
-	}, [ node ] );
+		}
 
-	return ref;
+		target.focus();
+	}, [] );
 }
-
-export default useFocusOnMount;