Refs on Portal are not set when running useLayoutEffect

[jacobp100]

🐛 Bug report

Current Behavior

() => {
  const ref = useRef();

  useLayoutEffect(() => {
    assert(ref.current != null)
  })

  return <Portal ref={ref} />
}

Expected behavior

Ref is set

Reproducible example

https://codesandbox.io/s/amazing-elgamal-ihrulq

Suggested solution(s)

I'm not fully following why ownerDocument is used - but if we can avoid it, this implementation below will have the ref set on time

const Portal = React.forwardRef(
  ({ container = document.body, children }, ref) => {
    const [target] = React.useState(() => document.createElement('div'));
    React.useImperativeHandle(ref, () => target, []);

    React.useLayoutEffect(() => {
      container.appendChild(target);
      return () => container.removeChild(target);
    }, [container, target]);

    return createPortal(children, container);
  }
);

Additional context

Your environment

Software	Name(s)	Version
Reach Package	17
React	16, 17, 18
Browser	All
Assistive tech	Nil
Node	Nil
npm/yarn	All
Operating System	All