Simple hook using React Context API to give control of the flow of focus.

use-focus.js

import React from 'react'

const FocusContext = React.createContext()

export function FocusProvider({ children }) {
  const focusElement = React.useRef(null)

  const cacheFocusElement = element => (focusElement.current = element)

  const toggleFocus = () => {
    if (focusElement.current) {
      focusElement.current.focus()
    }
  }

  return (
    <FocusContext.Provider
      value={{
        cacheFocusElement,
        toggleFocus,
      }}
    >
      {children}
    </FocusContext.Provider>
  )
}

export default function useFocus() {
  const { cacheFocusElement, toggleFocus } = React.useContext(FocusContext)
  return { cacheFocusElement, toggleFocus }
}

Simple usage

Wrap outer component with FocusProvider

import ReactDOM from 'react-dom'
import Demo from './demo'
import { FocusProvider } from './use-focus'

ReactDOM.render(
  <FocusProvider>
     <Demo />
  </FocusProvider>,
  document.getElementById('root')
)

Now, where ever you want to control focus:

import useFocus from './use-focus'

export default function Demo() {
  const [isModalShowing, toggleModal] = React.useState(false)
  const { cacheFocusElement, toggleFocus } = useFocus()
  
  // Don't trigger on initial render. You can move this to another hook, useUpdatedEffect.
  const initialRender = React.useRef(false)
  React.useEffect(() => {
    if (!initialRender.current) {
      initialRender.current = true
      return
    }
    if (!isModalShowing) {
      // Now that the modal is closed, toggleFocus will toggle focus back on the cacheFocusElement.
      toggleFocus()
    }
  }, [isModalShowing])

  return (
    <>
      <Modal>
         // ...
       </Modal>
      <ButtonThatTogglesModal 
        onClick={e => { 
           toggleModal()
           // Cache the current button so we can place focus on it later.
           cacheFocusElement(e.target)
        }}
      />
    </>
  )
}