From 148adb33392cf6010f1a578e18914a37bea59bdc Mon Sep 17 00:00:00 2001
From: Jon Koops <jonkoops@gmail.com>
Date: Thu, 22 Jun 2023 12:56:01 +0200
Subject: [PATCH] feat(FocusTrap): convert to function component

---
 .../src/helpers/FocusTrap/FocusTrap.tsx       | 144 ++++++------------
 .../src/helpers/useUnmountEffect.ts           |  18 +++
 2 files changed, 68 insertions(+), 94 deletions(-)
 create mode 100644 packages/react-core/src/helpers/useUnmountEffect.ts

diff --git a/packages/react-core/src/helpers/FocusTrap/FocusTrap.tsx b/packages/react-core/src/helpers/FocusTrap/FocusTrap.tsx
index 17b333e0aa8..977bd8809d2 100644
--- a/packages/react-core/src/helpers/FocusTrap/FocusTrap.tsx
+++ b/packages/react-core/src/helpers/FocusTrap/FocusTrap.tsx
@@ -1,107 +1,63 @@
-import * as React from 'react';
-import { createFocusTrap, Options as FocusTrapOptions, FocusTrap as IFocusTrap } from 'focus-trap';
+import { createFocusTrap, FocusTrap as FocusTrapInstance, Options as FocusTrapOptions } from 'focus-trap';
+import React, { ComponentPropsWithRef, forwardRef, useEffect, useImperativeHandle, useRef } from 'react';
+import { useUnmountEffect } from '../useUnmountEffect';
 
-interface FocusTrapProps extends Omit<React.HTMLProps<HTMLDivElement>, 'ref'> {
-  children: React.ReactNode;
-  className?: string;
+export interface FocusTrapProps extends ComponentPropsWithRef<'div'> {
   active?: boolean;
   paused?: boolean;
   focusTrapOptions?: FocusTrapOptions;
   /** Prevent from scrolling to the previously focused element on deactivation */
   preventScrollOnDeactivate?: boolean;
-  /** @hide Forwarded ref */
-  innerRef?: React.Ref<HTMLDivElement>;
 }
 
-class FocusTrapBase extends React.Component<FocusTrapProps> {
-  static displayName = 'FocusTrap';
-  previouslyFocusedElement: HTMLElement;
-  focusTrap: IFocusTrap;
-  divRef = (this.props.innerRef as React.RefObject<HTMLDivElement>) || React.createRef<HTMLDivElement>();
-
-  static defaultProps = {
-    active: true,
-    paused: false,
-    focusTrapOptions: {},
-    preventScrollOnDeactivate: false
-  };
-
-  constructor(props: FocusTrapProps) {
-    super(props);
-
-    if (typeof document !== 'undefined') {
-      this.previouslyFocusedElement = document.activeElement as HTMLElement;
-    }
-  }
-
-  componentDidMount() {
-    // We need to hijack the returnFocusOnDeactivate option,
-    // because React can move focus into the element before we arrived at
-    // this lifecycle hook (e.g. with autoFocus inputs). So the component
-    // captures the previouslyFocusedElement in componentWillMount,
-    // then (optionally) returns focus to it in componentWillUnmount.
-    this.focusTrap = createFocusTrap(this.divRef.current, {
-      ...this.props.focusTrapOptions,
+export const FocusTrap = forwardRef<HTMLDivElement, FocusTrapProps>(function FocusTrap(
+  { active = true, paused = false, focusTrapOptions = {}, preventScrollOnDeactivate = false, ...props },
+  forwardedRef
+) {
+  // Fall back to internal ref if no forwarded ref is provided.
+  const ref = useRef<HTMLDivElement>(null);
+  useImperativeHandle(forwardedRef, () => ref.current!);
+
+  // Create focus trap instance after rendering DOM.
+  const focusTrapRef = useRef<FocusTrapInstance | null>(null);
+  useEffect(() => {
+    const focusTrap = createFocusTrap(ref.current!, {
+      ...focusTrapOptions,
       returnFocusOnDeactivate: false
     });
-    if (this.props.active) {
-      this.focusTrap.activate();
-    }
-    if (this.props.paused) {
-      this.focusTrap.pause();
-    }
-  }
-
-  componentDidUpdate(prevProps: FocusTrapProps) {
-    if (prevProps.active && !this.props.active) {
-      this.deactivate();
-    } else if (!prevProps.active && this.props.active) {
-      this.focusTrap.activate();
-    }
-
-    if (prevProps.paused && !this.props.paused) {
-      this.focusTrap.unpause();
-    } else if (!prevProps.paused && this.props.paused) {
-      this.focusTrap.pause();
+    focusTrapRef.current = focusTrap;
+
+    // Deactivate focus trap on cleanup.
+    return () => {
+      focusTrap.deactivate();
+    };
+  }, []);
+
+  // Handle activation status based on 'active' prop.
+  useEffect(() => {
+    const focusTrap = focusTrapRef.current;
+    active ? focusTrap?.activate() : focusTrap?.deactivate();
+  }, [active]);
+
+  // Handle pause status based on 'pause' prop.
+  useEffect(() => {
+    const focusTrap = focusTrapRef.current;
+    paused ? focusTrap?.pause() : focusTrap?.unpause();
+  }, [paused]);
+
+  // Store the currently active element to restore focus to later.
+  const previousElementRef = useRef(typeof document !== 'undefined' ? document.activeElement : null);
+
+  // Restore focus to the previously active element on unmount.
+  useUnmountEffect(() => {
+    if (focusTrapOptions.returnFocusOnDeactivate !== false && previousElementRef.current instanceof HTMLElement) {
+      previousElementRef.current.focus({
+        preventScroll: preventScrollOnDeactivate
+      });
     }
-  }
-
-  componentWillUnmount() {
-    this.deactivate();
-  }
+  });
 
-  deactivate() {
-    this.focusTrap.deactivate();
-    if (
-      this.props.focusTrapOptions.returnFocusOnDeactivate !== false &&
-      this.previouslyFocusedElement &&
-      this.previouslyFocusedElement.focus
-    ) {
-      this.previouslyFocusedElement.focus({ preventScroll: this.props.preventScrollOnDeactivate });
-    }
-  }
-
-  render() {
-    const {
-      children,
-      className,
-      /* eslint-disable @typescript-eslint/no-unused-vars */
-      focusTrapOptions,
-      active,
-      paused,
-      preventScrollOnDeactivate,
-      innerRef,
-      /* eslint-enable @typescript-eslint/no-unused-vars */
-      ...rest
-    } = this.props;
-    return (
-      <div ref={this.divRef} className={className} {...rest}>
-        {children}
-      </div>
-    );
-  }
-}
+  return <div ref={ref} {...props} />;
+});
 
-export const FocusTrap = React.forwardRef((props: FocusTrapProps, ref: React.Ref<any>) => (
-  <FocusTrapBase innerRef={ref} {...props} />
-));
+FocusTrap.displayName = 'FocusTrap';
diff --git a/packages/react-core/src/helpers/useUnmountEffect.ts b/packages/react-core/src/helpers/useUnmountEffect.ts
new file mode 100644
index 00000000000..d776db5d2e6
--- /dev/null
+++ b/packages/react-core/src/helpers/useUnmountEffect.ts
@@ -0,0 +1,18 @@
+import { EffectCallback, useEffect, useRef } from 'react';
+
+/**
+ * A `useEffect`-like hook that only triggers when a component unmounts. Does not require a dependency list, as the effect callback will always be kept up to date.
+ */
+export function useUnmountEffect(effect: EffectCallback) {
+  // Always use the latest effect callback so that it can reference the latest props and state.
+  const effectRef = useRef(effect);
+  effectRef.current = effect;
+
+  // Trigger the effect callback when the component unmounts.
+  useEffect(
+    () => () => {
+      effectRef.current();
+    },
+    []
+  );
+}