utils: add descendants hook (prelim)

packages/utils/src/index.tsx

@@ -1,11 +1,16 @@
 /* eslint-disable no-unused-vars */
 
 import {
-  useRef,
-  useMemo,
-  useEffect,
+  cloneElement,
+  createContext,
   isValidElement,
-  cloneElement
+  useCallback,
+  useContext,
+  useEffect,
+  useLayoutEffect,
+  useMemo,
+  useRef,
+  useState
 } from "react";
 import {
   As,
@@ -81,6 +86,14 @@ export function assignRef<T = any>(ref: AssignableRef<T>, value: any) {
   }
 }
 
+export function canUseDOM() {
+  return (
+    typeof window !== "undefined" &&
+    typeof window.document !== "undefined" &&
+    typeof window.document.createElement !== "undefined"
+  );
+}
+
 export function cloneValidElement<P>(
   element: React.ReactElement<P> | React.ReactNode,
   props?: Partial<P> & React.Attributes,
@@ -139,6 +152,29 @@ export function useForkedRef<T = any>(...refs: AssignableRef<T>[]) {
   }, refs);
 }
 
+/**
+ * React currently throws a warning when using useLayoutEffect on the server.
+ * To get around it, we can conditionally useEffect on the server (no-op) and
+ * useLayoutEffect in the browser. We occasionally need useLayoutEffect to ensure
+ * we don't get a render flash for certain operations, but we may also need
+ * affected components to render on the server. One example is when setting a
+ * component's descendants to retrieve their index values. The index value may be
+ * needed to determine whether a descendant is active, but with useEffect in the
+ * browser there will be an initial frame where the active descendant is not set.
+ *
+ * Important to note that using this hook as an escape hatch will break the
+ * eslint dependency warnings, so use sparingly only when needed and pay close
+ * attention to the dependency array!
+ *
+ * https://github.com/reduxjs/react-redux/blob/master/src/utils/useIsomorphicLayoutEffect.js
+ *
+ * @param effect
+ * @param deps
+ */
+export const useIsomorphicLayoutEffect = canUseDOM()
+  ? useLayoutEffect
+  : useEffect;
+
 /**
  * Returns the previous value of a reference after a component update.
  *
@@ -203,6 +239,7 @@ export function forwardRefWithAs<T extends As, P>(
   return React.forwardRef(Comp as any) as ComponentWithAs<T, P>;
 }
 
+// Export types
 export {
   As,
   AssignableRef,
@@ -211,3 +248,191 @@ export {
   PropsFromAs,
   PropsWithAs
 };
+
+////////////////////////////////////////////////////////////////////////////////
+
+type DescendantElement<T = HTMLElement> = T extends HTMLElement
+  ? T
+  : HTMLElement;
+
+type Descendant<T> = {
+  element: DescendantElement<T>;
+  key?: string | number | null;
+  disabled?: boolean;
+};
+
+interface IDescendantContext<T> {
+  descendants: Descendant<T>[];
+  focusNodes: DescendantElement<T>[];
+  registerDescendant(descendant: Descendant<T>): void;
+  unregisterDescendant(element: Descendant<T>["element"]): void;
+}
+
+const DescendantContext = createContext<IDescendantContext<any>>(
+  {} as IDescendantContext<any>
+);
+
+function useDescendantContext<T>() {
+  return useContext(DescendantContext as React.Context<IDescendantContext<T>>);
+}
+
+////////////////////////////////////////////////////////////////////////////////
+// TODO: Move to @reach/descendants once fully tested and implemented
+
+/**
+ * This hook registers our descendant by passing it into an array. We can then
+ * search that array by to find its index when registering it in the component.
+ * We use this for focus management, keyboard navigation, and typeahead
+ * functionality for some components.
+ *
+ * The hook accepts the element node and (optionally) a key. The key is useful
+ * if multiple descendants have identical text values and we need to
+ * differentiate siblings for some reason.
+ *
+ * Our main goals with this are:
+ *   1) maximum composability,
+ *   2) minimal API friction
+ *   3) SSR compatibility*
+ *   4) concurrent safe
+ *   5) index always up-to-date with the tree despite changes
+ *   6) works with memoization of any component in the tree (hopefully)
+ *
+ * * As for SSR, the good news is that we don't actually need the index on the
+ * server for most use-cases, as we are only using it to determine the order of
+ * composed descendants for keyboard navigation. However, in the few cases where
+ * this is not the case, we can require an explicit index from the app.
+ */
+export function useDescendant<T>(
+  { element, key, disabled }: Descendant<T>,
+  indexProp?: number
+) {
+  let [, forceUpdate] = useState();
+  let {
+    registerDescendant,
+    unregisterDescendant,
+    descendants
+  } = useDescendantContext<T>();
+
+  // Prevent any flashing
+  useIsomorphicLayoutEffect(() => {
+    if (!element) forceUpdate({});
+    registerDescendant({ element, key, disabled });
+    return () => unregisterDescendant(element);
+  }, [element, key, disabled]);
+
+  return (
+    indexProp ?? descendants.findIndex(({ element: _el }) => _el === element)
+  );
+}
+
+export function useDescendants<T>() {
+  return useState<Descendant<T>[]>([]);
+}
+
+export function DescendantProvider<T>({
+  children,
+  descendants,
+  setDescendants
+}: {
+  children: React.ReactNode;
+  descendants: Descendant<T>[];
+  setDescendants: React.Dispatch<React.SetStateAction<Descendant<T>[]>>;
+}) {
+  let registerDescendant = React.useCallback(
+    ({ disabled, element, key: providedKey }: Descendant<T>) => {
+      if (!element) {
+        return;
+      }
+
+      setDescendants(items => {
+        if (items.find(({ element: _el }) => _el === element) == null) {
+          let key = providedKey ?? element.textContent;
+
+          /*
+           * When registering a descendant, we need to make sure we insert in
+           * into the array in the same order that it appears in the DOM. So as
+           * new descendants are added or maybe some are removed, we always know
+           * that the array is up-to-date and correct.
+           *
+           * So here we look at our registered descendants and see if the new
+           * element we are adding appears earlier than an existing descendant's
+           * DOM node via `node.compareDocumentPosition`. If it does, we insert
+           * the new element at this index. Because `registerDescendant` will be
+           * called in an effect every time the descendants state value changes,
+           * we should be sure that this index is accurate when descendent
+           * elements come or go from our component.
+           */
+          let index = items.findIndex(({ element: existingElement }) => {
+            if (!existingElement || !element) {
+              return false;
+            }
+            /*
+             * Does this element's DOM node appear before another item in the
+             * array in our DOM tree? If so, return true to grab the index at
+             * this point in the array so we know where to insert the new
+             * element.
+             */
+            return Boolean(
+              existingElement.compareDocumentPosition(element) &
+                Node.DOCUMENT_POSITION_PRECEDING
+            );
+          });
+
+          let newItem = { disabled, element, key };
+
+          // If an index is not found we will push the element to the end.
+          if (index === -1) {
+            return [...items, newItem];
+          }
+          return [...items.slice(0, index), newItem, ...items.slice(index)];
+        }
+        return items;
+      });
+    },
+    /*
+     * setDescendants is a state setter initialized by the useDescendants hook.
+     * We can safely ignore the lint warning here because it will not change
+     * between renders.
+     */
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+    []
+  );
+
+  let unregisterDescendant = useCallback(
+    (element: Descendant<T>["element"]) => {
+      if (!element) {
+        return;
+      }
+
+      setDescendants(items =>
+        items.filter(({ element: _el }) => element !== _el)
+      );
+    },
+    /*
+     * setDescendants is a state setter initialized by the useDescendants hook.
+     * We can safely ignore the lint warning here because it will not change
+     * between renders.
+     */
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+    []
+  );
+
+  let focusNodes = descendants
+    .filter(({ disabled }) => !disabled)
+    .map(({ element }) => element);
+
+  const value: IDescendantContext<T> = useMemo(() => {
+    return {
+      descendants,
+      focusNodes,
+      registerDescendant,
+      unregisterDescendant
+    };
+  }, [descendants, focusNodes, registerDescendant, unregisterDescendant]);
+
+  return (
+    <DescendantContext.Provider value={value}>
+      {children}
+    </DescendantContext.Provider>
+  );
+}