Float.tsx

import React, { createContext, forwardRef, useContext, useEffect, useMemo } from "react"
import { createPortal } from "react-dom"
import { mergeRefs } from "app/helpers/mergeRefs"
import {
  Placement,
  Strategy,
  offset as offsetModifier,
  shift as shiftModifier,
  size as sizeModifier,
  flip as flipModifier,
  useFloating,
  getOverflowAncestors,
} from "@floating-ui/react-dom"

import type { Options } from "@floating-ui/core/src/middleware/offset"

const modalRoot = document.getElementById("modal-root")!

/**
 * Float is a compound component for a flexible popover
 * It is used like this:
 *
 *  <Float>
 *    <Float.Reference>
 *      ...
 *    </Float.Reference>
 *    <Float.Floating>
 *      ...
 *    </Float.Floating>
 *  </Float>
 */
export interface FloatProps {
  placement?: Placement
  strategy?: Strategy
  offset?: Options
  shift?: boolean
  flip?: boolean
  matchSize?: boolean
  children?: React.ReactNode
  portalRoot?: HTMLElement
}
function Float({
  placement,
  strategy,
  offset,
  shift = true,
  flip = true,
  matchSize = false,
  children,
  portalRoot,
}: FloatProps) {
  const {
    floating,
    placement: actualPlacement,
    reference,
    strategy: actualStrategy,
    update,
    x,
    y,
    // destructure the refs as `refs` is not a stable reference
    refs: { floating: floatingRef, reference: referenceRef },
  } = useFloating({
    placement: placement,
    strategy: strategy,
    middleware: [
      ...(offset ? [offsetModifier(offset)] : []),
      ...(shift ? [shiftModifier()] : []),
      ...(flip ? [flipModifier()] : []),
      ...(matchSize
        ? [
            sizeModifier({
              apply({ reference }) {
                if (!floatingRef.current) return
                Object.assign(floatingRef.current.style, {
                  width: `${reference.width}px`,
                })
              },
            }),
          ]
        : []),
    ],
  })

  const value = useMemo(
    () => ({
      placement: actualPlacement,
      strategy: actualStrategy,
      floating,
      floatingRef,
      portalRoot,
      reference,
      referenceRef,
      update,
      x,
      y,
    }),
    [
      actualPlacement,
      actualStrategy,
      floating,
      floatingRef,
      portalRoot,
      reference,
      referenceRef,
      update,
      x,
      y,
    ],
  )

  return <FloatContext.Provider value={value}>{children}</FloatContext.Provider>
}

const Reference = forwardRef<HTMLDivElement, React.HTMLAttributes<HTMLDivElement>>(
  function FloatingButton({ children, ...props }, ref) {
    const { reference } = useFloatContext()

    return (
      <div ref={mergeRefs([reference, ref])} {...props}>
        {children}
      </div>
    )
  },
)

const Floating = forwardRef<HTMLDivElement, React.HTMLAttributes<HTMLDivElement>>(function Modal(
  { children, ...props },
  ref,
) {
  const { floating, strategy, x, y, update, referenceRef, floatingRef, portalRoot } =
    useFloatContext()

  // We trigger updateOnParentChange here instead of in the parent <Float> component
  // because the useEffect there runs before floatingRef is attached to the DOM.
  // When used here, it works as intended.
  useUpdateOnParentChange(referenceRef, floatingRef, update)

  // A non modal root can be passed in so it can be rendered inside
  // a modal and not trigger a click outside close event.
  const root = portalRoot || modalRoot

  return createPortal(
    <div
      ref={mergeRefs([ref, floating])}
      {...props}
      style={{
        ...props.style,
        position: strategy,
        top: y ?? "",
        left: x ?? "",
      }}
    >
      {children}
    </div>,
    root,
  )
})

type UseFloatingReturn = ReturnType<typeof useFloating>
// middlewareData and refs are not stable, and cannot be used in useMemo without causing infinite re-renders
type StableFloatingProps = Omit<UseFloatingReturn, "middlewareData" | "refs">
// We want to pass the refs, but we have to individually because their container is not a stable reference
// https://github.com/floating-ui/floating-ui/issues/1532
interface FloatContextInterface extends StableFloatingProps {
  referenceRef: UseFloatingReturn["refs"]["reference"]
  floatingRef: UseFloatingReturn["refs"]["floating"]
  portalRoot?: HTMLElement | null
}
const FloatContext = createContext<FloatContextInterface | null>(null)
function useFloatContext() {
  const context = useContext(FloatContext)
  if (!context) {
    throw new Error("Float compound components cannot be rendered outside the Float component")
  }
  return context
}

function useUpdateOnParentChange(
  referenceRef: FloatContextInterface["referenceRef"],
  floatingRef: FloatContextInterface["floatingRef"],
  update: FloatContextInterface["update"],
) {
  useEffect(() => {
    if (!referenceRef.current || !floatingRef.current) {
      return
    }

    const parents = [
      ...getOverflowAncestors(referenceRef.current as Node),
      ...getOverflowAncestors(floatingRef.current as Node),
    ]

    parents.forEach((parent) => {
      parent.addEventListener("scroll", update)
      parent.addEventListener("resize", update)
    })

    return () => {
      parents.forEach((parent) => {
        parent.removeEventListener("scroll", update)
        parent.removeEventListener("resize", update)
      })
    }
  }, [referenceRef, floatingRef, update])
}

Float.Reference = Reference
Float.Floating = Floating

export default Float