From b21a82d4aa2c519b33e7e31124f9f01eae4c2c30 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ella=20van=C2=A0Durpe?= <wp@iseulde.com>
Date: Mon, 9 Nov 2020 16:18:20 +0200
Subject: [PATCH] Focus return: make document.activeElement relative to ref
 (#26814)

* Try without activeElement

* Prepend the focus history with the active element on mount
---
 .../higher-order/with-focus-return/context.js | 45 +++++++++----------
 .../higher-order/with-focus-return/index.js   | 24 +++++-----
 .../with-focus-return/test/index.js           | 30 ++++++++-----
 3 files changed, 50 insertions(+), 49 deletions(-)

diff --git a/packages/components/src/higher-order/with-focus-return/context.js b/packages/components/src/higher-order/with-focus-return/context.js
index f65834f4886927..77def70b3fb89a 100644
--- a/packages/components/src/higher-order/with-focus-return/context.js
+++ b/packages/components/src/higher-order/with-focus-return/context.js
@@ -6,7 +6,7 @@ import { uniq } from 'lodash';
 /**
  * WordPress dependencies
  */
-import { Component, createContext } from '@wordpress/element';
+import { createContext, useEffect, useRef, useState } from '@wordpress/element';
 
 const { Provider, Consumer } = createContext( {
 	focusHistory: [],
@@ -23,20 +23,19 @@ Consumer.displayName = 'FocusReturnConsumer';
  */
 const MAX_STACK_LENGTH = 100;
 
-class FocusReturnProvider extends Component {
-	constructor() {
-		super( ...arguments );
+function FocusReturnProvider( { children, className } ) {
+	const ref = useRef();
+	const [ focusHistory, setFocusHistory ] = useState( [] );
 
-		this.onFocus = this.onFocus.bind( this );
-
-		this.state = {
-			focusHistory: [],
-		};
-	}
-
-	onFocus( event ) {
-		const { focusHistory } = this.state;
+	// Prepend the focus history with the active element on mount.
+	useEffect( () => {
+		setFocusHistory( [
+			ref.current.ownerDocument.activeElement,
+			...focusHistory,
+		] );
+	}, [] );
 
+	function onFocus( event ) {
 		// Push the focused element to the history stack, keeping only unique
 		// members but preferring the _last_ occurrence of any duplicates.
 		// Lodash's `uniq` behavior favors the first occurrence, so the array
@@ -51,20 +50,16 @@ class FocusReturnProvider extends Component {
 				.reverse()
 		).reverse();
 
-		this.setState( { focusHistory: nextFocusHistory } );
+		setFocusHistory( nextFocusHistory );
 	}
 
-	render() {
-		const { children, className } = this.props;
-
-		return (
-			<Provider value={ this.state }>
-				<div onFocus={ this.onFocus } className={ className }>
-					{ children }
-				</div>
-			</Provider>
-		);
-	}
+	return (
+		<Provider value={ focusHistory }>
+			<div ref={ ref } onFocus={ onFocus } className={ className }>
+				{ children }
+			</div>
+		</Provider>
+	);
 }
 
 export default FocusReturnProvider;
diff --git a/packages/components/src/higher-order/with-focus-return/index.js b/packages/components/src/higher-order/with-focus-return/index.js
index 6b12b693899a2e..08fa652cb1176f 100644
--- a/packages/components/src/higher-order/with-focus-return/index.js
+++ b/packages/components/src/higher-order/with-focus-return/index.js
@@ -55,7 +55,6 @@ function withFocusReturn( options ) {
 				super( ...arguments );
 
 				this.ownFocusedElements = new Set();
-				this.activeElementOnMount = document.activeElement;
 				this.setIsFocusedFalse = () => ( this.isFocused = false );
 				this.setIsFocusedTrue = ( event ) => {
 					this.ownFocusedElements.add( event.target );
@@ -64,11 +63,7 @@ function withFocusReturn( options ) {
 			}
 
 			componentWillUnmount() {
-				const {
-					activeElementOnMount,
-					isFocused,
-					ownFocusedElements,
-				} = this;
+				const { isFocused, ownFocusedElements } = this;
 
 				if ( ! isFocused ) {
 					return;
@@ -83,15 +78,13 @@ function withFocusReturn( options ) {
 					return;
 				}
 
-				const stack = [
-					...without(
-						this.props.focus.focusHistory,
-						...ownFocusedElements
-					),
-					activeElementOnMount,
-				];
+				const stack = without(
+					this.props.focusHistory,
+					...ownFocusedElements
+				);
 
 				let candidate;
+
 				while ( ( candidate = stack.pop() ) ) {
 					if ( document.body.contains( candidate ) ) {
 						candidate.focus();
@@ -115,7 +108,10 @@ function withFocusReturn( options ) {
 		return ( props ) => (
 			<Consumer>
 				{ ( context ) => (
-					<FocusReturn childProps={ props } focus={ context } />
+					<FocusReturn
+						childProps={ props }
+						focusHistory={ context }
+					/>
 				) }
 			</Consumer>
 		);
diff --git a/packages/components/src/higher-order/with-focus-return/test/index.js b/packages/components/src/higher-order/with-focus-return/test/index.js
index 67502237260b32..6bef3fb8776d94 100644
--- a/packages/components/src/higher-order/with-focus-return/test/index.js
+++ b/packages/components/src/higher-order/with-focus-return/test/index.js
@@ -70,11 +70,16 @@ describe( 'withFocusReturn()', () => {
 		} );
 
 		it( 'should not switch focus back to the bound focus element', () => {
-			const { unmount } = render( <Composite />, {
-				container: document.body.appendChild(
-					document.createElement( 'div' )
-				),
-			} );
+			const { unmount } = render(
+				<Provider>
+					<Composite />
+				</Provider>,
+				{
+					container: document.body.appendChild(
+						document.createElement( 'div' )
+					),
+				}
+			);
 
 			// Change activeElement.
 			switchFocusTo.focus();
@@ -86,11 +91,16 @@ describe( 'withFocusReturn()', () => {
 		} );
 
 		it( 'should switch focus back when unmounted while having focus', () => {
-			const { container, unmount } = render( <Composite />, {
-				container: document.body.appendChild(
-					document.createElement( 'div' )
-				),
-			} );
+			const { container, unmount } = render(
+				<Provider>
+					<Composite />
+				</Provider>,
+				{
+					container: document.body.appendChild(
+						document.createElement( 'div' )
+					),
+				}
+			);
 
 			const textarea = container.querySelector( 'textarea' );
 			textarea.focus();