useRovingTabIndex.tsx

import React, {
  useEffect,
  createContext,
  FunctionComponent,
  useRef,
  MutableRefObject,
  useCallback,
  useContext,
  useReducer,
  Reducer,
  Dispatch,
  useMemo,
  useLayoutEffect,
} from 'react';
import { findNodeHandle } from 'react-native';
import { assertNever } from 'assert-never';
import produce from 'immer';

interface RovingTabIndexState {
  current: HTMLElement | null;
  focusables: HTMLElement[];
  lastAction: RovingTabIndexAction | null;
}

type RovingTabIndexAction =
  | {
      type: 'add';
      element: HTMLElement;
      firstFocus?: boolean;
    }
  | {
      type: 'remove';
      element: HTMLElement;
    }
  | {
      type: 'setCurrent';
      element: HTMLElement;
    }
  | {
      type: 'horizontal';
      left: boolean;
    };

type RovingTabIndexReducer = Reducer<RovingTabIndexState, RovingTabIndexAction>;

const rovingTabIndexInitialState: RovingTabIndexState = {
  current: null,
  focusables: [],
  lastAction: null,
};

const documentPositionComparator = (a: HTMLElement, b: HTMLElement) => {
  if (a === b) return 0;
  const position = a.compareDocumentPosition(b);
  if (
    // eslint-disable-next-line no-bitwise
    position & Node.DOCUMENT_POSITION_FOLLOWING ||
    // eslint-disable-next-line no-bitwise
    position & Node.DOCUMENT_POSITION_CONTAINED_BY
  )
    return -1;
  if (
    // eslint-disable-next-line no-bitwise
    position & Node.DOCUMENT_POSITION_PRECEDING ||
    // eslint-disable-next-line no-bitwise
    position & Node.DOCUMENT_POSITION_CONTAINS
  ) {
    return 1;
  }
  return 0;
};

const rovingTabIndexReducer: RovingTabIndexReducer = (state, action) => {
  // console.log(action);
  state = { ...state, lastAction: action };

  const getSortedFocusables = (focusables: HTMLElement[]) => {
    // Shallow copy to not change draftState.focusables.
    return [...focusables].sort(documentPositionComparator);
  };

  switch (action.type) {
    case 'add':
      // https://github.com/immerjs/immer/issues/376
      return produce(state, (draftState: RovingTabIndexState) => {
        draftState.focusables.push(action.element);
        if (action.firstFocus) {
          draftState.current = action.element;
        }
      });

    case 'remove':
      return produce(state, (draftState: RovingTabIndexState) => {
        const index = draftState.focusables.indexOf(action.element);
        draftState.focusables.splice(index, 1);
        const currentHasBeenRemoved = draftState.current === action.element;
        if (currentHasBeenRemoved) {
          const sortedFocusables = getSortedFocusables(draftState.focusables);
          // Ensure the closest focusable will be current, if exists.
          draftState.current =
            // Removed in the middle should have the same position.
            sortedFocusables[index] ||
            // If last has been removed, get previous.
            sortedFocusables[sortedFocusables.length - 1] ||
            // OK, everything has been removed, set null.
            null;
        }
      });

    case 'setCurrent':
      return produce(state, (draftState: RovingTabIndexState) => {
        draftState.current = action.element;
      });

    case 'horizontal':
      return produce(state, (draftState: RovingTabIndexState) => {
        if (draftState.current == null) return;
        const sortedFocusables = getSortedFocusables(draftState.focusables);
        const index = sortedFocusables.indexOf(draftState.current);
        const next = sortedFocusables[index + (action.left ? -1 : 1)];
        if (next == null) return;
        draftState.current = next;
        next.focus();
      });
    default:
      return assertNever(action);
  }
};

interface RovingTabIndexProviderContext {
  dispatch: Dispatch<RovingTabIndexAction>;
  state: RovingTabIndexState;
}

const RovingTabIndexProviderContext = createContext<
  RovingTabIndexProviderContext
>({
  dispatch: () => {},
  state: rovingTabIndexInitialState,
});

export const RovingTabIndexProvider: FunctionComponent = ({ children }) => {
  const [state, dispatch] = useReducer<RovingTabIndexReducer>(
    rovingTabIndexReducer,
    rovingTabIndexInitialState,
  );

  const value = useMemo(() => {
    return { dispatch, state };
  }, [state]);

  const { current, lastAction } = state;
  useLayoutEffect(() => {
    if (current && lastAction && lastAction.type === 'horizontal') {
      current.focus();
    }
  }, [current, lastAction]);

  return (
    <RovingTabIndexProviderContext.Provider value={value}>
      {children}
    </RovingTabIndexProviderContext.Provider>
  );
};

// TODO: It should be RN something.
type ReactNativeSomethingRef = any;

type TabIndex = -1 | 0;

export const useRovingTabIndex = ({
  firstFocus,
}: {
  firstFocus?: boolean;
}): [MutableRefObject<ReactNativeSomethingRef>, TabIndex] => {
  const componentRef = useRef<ReactNativeSomethingRef | null>(null);
  const elementRef = useRef<HTMLElement | null>(null);
  const { dispatch, state } = useContext(RovingTabIndexProviderContext);

  const setCurrent = useCallback(() => {
    const { current } = elementRef;
    if (current == null) return;
    dispatch({ type: 'setCurrent', element: current });
  }, [dispatch]);

  const handleKeyDown = useCallback(
    (event: KeyboardEvent) => {
      const left = event.key === 'ArrowLeft';
      if (left || event.key === 'ArrowRight') {
        dispatch({ type: 'horizontal', left });
        // return;
      }
    },
    [dispatch],
  );

  const currentTabIndex = state.current === elementRef.current ? 0 : -1;

  // TODO: https://github.com/facebook/react/pull/15022
  useEffect(() => {
    const { current } = componentRef;
    if (current == null) return;
    // We have to bypass React Native Web.
    const element = findNodeHandle(current) as HTMLElement | null;
    if (element == null) return;
    elementRef.current = element;
    // Use useSubscription hook
    // https://github.com/facebook/react/pull/15022
    element.addEventListener('focus', setCurrent);
    element.addEventListener('mousedown', setCurrent);
    element.addEventListener('keydown', handleKeyDown);
    dispatch({ type: 'add', element, firstFocus });
    return () => {
      element.removeEventListener('focus', setCurrent);
      element.removeEventListener('mousedown', setCurrent);
      element.removeEventListener('keydown', handleKeyDown);
      dispatch({ type: 'remove', element });
    };
  }, [dispatch, firstFocus, setCurrent, handleKeyDown]);

  return [componentRef, currentTabIndex];
};