Refactor

packages/tabs/src/TabPanel/TabPanel.tsx

@@ -1,11 +1,16 @@
 import React, { useMemo } from 'react';
 
-import { useTabDescendantsContext, useTabPanelDescendant } from '../context';
+import { useDescendant } from '@leafygreen-ui/descendants';
+
+import {
+  TabPanelDescendantsContext,
+  useTabDescendantsContext,
+} from '../context';
 
 import { TabPanelProps } from './TabPanel.types';
 
 const TabPanel = ({ child, selectedIndex }: TabPanelProps) => {
-  const { id, index, ref } = useTabPanelDescendant();
+  const { id, index, ref } = useDescendant(TabPanelDescendantsContext);
   const { tabDescendants } = useTabDescendantsContext();
 
   const relatedTab = useMemo(() => {

packages/tabs/src/TabTitle/TabTitle.tsx

@@ -1,20 +1,14 @@
-import React, {
-  RefObject,
-  useCallback,
-  useEffect,
-  useMemo,
-  useRef,
-} from 'react';
+import React, { RefObject, useCallback, useMemo, useRef } from 'react';
 
 import Box, { ExtendableBox } from '@leafygreen-ui/box';
+import { useDescendant } from '@leafygreen-ui/descendants';
 import { cx } from '@leafygreen-ui/emotion';
 import { getNodeTextContent, Theme } from '@leafygreen-ui/lib';
 import { BaseFontSize } from '@leafygreen-ui/tokens';
 import { useUpdatedBaseFontSize } from '@leafygreen-ui/typography';
 
 import {
-  useTabDescendant,
-  useTabDescendantsContext,
+  TabDescendantsContext,
   useTabPanelDescendantsContext,
 } from '../context';
 
@@ -37,31 +31,12 @@ const TabTitle: ExtendableBox<BaseTabTitleProps, 'button'> = ({
 }: BaseTabTitleProps) => {
   const baseFontSize: BaseFontSize = useUpdatedBaseFontSize();
   const titleRef = useRef<HTMLAnchorElement | HTMLButtonElement>(null);
-  const { index, ref, id } = useTabDescendant();
-  const { tabDescendants } = useTabDescendantsContext();
+  const { index, ref, id } = useDescendant(TabDescendantsContext);
   const { tabPanelDescendants } = useTabPanelDescendantsContext();
 
   const theme = darkMode ? Theme.Dark : Theme.Light;
   const selected = index === selectedIndex;
 
-  useEffect(() => {
-    // if tab is disabled or not selected, return early
-    // otherwise, focus may need to be manually moved
-    if (disabled || !selected || !titleRef.current) return;
-
-    // if focus is not on tab descendants, return early
-    // otherwise, focus needs to be manually moved
-    const activeEl = document.activeElement;
-    const tabList = tabDescendants.map(
-      descendant => descendant.element.parentElement,
-    );
-    const isFocusOnTabDescendants =
-      activeEl instanceof HTMLElement && tabList.indexOf(activeEl) !== -1;
-    if (!isFocusOnTabDescendants) return;
-
-    titleRef.current.focus();
-  }, [disabled, selected, tabDescendants, titleRef]);
-
   const relatedTabPanel = useMemo(() => {
     return tabPanelDescendants.find(
       tabPanelDescendant => tabPanelDescendant.index === index,

packages/tabs/src/Tabs/Tabs.tsx

@@ -1,7 +1,11 @@
-import React, { useCallback, useMemo, useState } from 'react';
+import React, { useCallback, useEffect, useState } from 'react';
 import PropTypes from 'prop-types';
 
 import { validateAriaLabelProps } from '@leafygreen-ui/a11y';
+import {
+  DescendantsProvider,
+  useInitDescendants,
+} from '@leafygreen-ui/descendants';
 import { cx } from '@leafygreen-ui/emotion';
 import { useIdAllocator } from '@leafygreen-ui/hooks';
 import LeafyGreenProvider, {
@@ -12,10 +16,7 @@ import { BaseFontSize } from '@leafygreen-ui/tokens';
 import { useUpdatedBaseFontSize } from '@leafygreen-ui/typography';
 
 import { LGIDS_TABS } from '../constants';
-import {
-  TabDescendantsProvider,
-  TabPanelDescendantsProvider,
-} from '../context';
+import { TabDescendantsContext, TabPanelDescendantsContext } from '../context';
 import TabPanel from '../TabPanel';
 import TabTitle from '../TabTitle';
 
@@ -66,23 +67,27 @@ const Tabs = (props: AccessibleTabsProps) => {
 
   const baseFontSize: BaseFontSize = useUpdatedBaseFontSize(baseFontSizeProp);
   const { theme, darkMode } = useDarkMode(darkModeProp);
-  const accessibilityProps = {
-    ['aria-label']: ariaLabel,
-    ['aria-labelledby']: ariaLabelledby,
-  };
-
   const id = useIdAllocator({ prefix: rest.id || 'tabs' });
 
+  const { descendants: tabDescendants, dispatch: tabDispatch } =
+    useInitDescendants<HTMLDivElement>();
+  const { descendants: tabPanelDescendants, dispatch: tabPanelDispatch } =
+    useInitDescendants<HTMLDivElement>();
+
   const [uncontrolledSelected, setUncontrolledSelected] = useState(0);
   const [selected, setSelected] = [
-    controlledSelected ? controlledSelected : uncontrolledSelected,
-    controlledSelected ? setControlledSelected : setUncontrolledSelected,
+    typeof controlledSelected === 'undefined'
+      ? uncontrolledSelected
+      : controlledSelected,
+    typeof controlledSelected === 'undefined'
+      ? setUncontrolledSelected
+      : setControlledSelected,
   ];
 
-  const childrenArray = useMemo(
-    () => React.Children.toArray(children) as Array<React.ReactElement>,
-    [children],
-  );
+  const accessibilityProps = {
+    ['aria-label']: ariaLabel,
+    ['aria-labelledby']: ariaLabelledby,
+  };
 
   const handleClickTab = useCallback(
     (e: React.SyntheticEvent<Element, MouseEvent>, index: number) => {
@@ -91,33 +96,55 @@ const Tabs = (props: AccessibleTabsProps) => {
     [setSelected],
   );
 
-  const getEnabledIndexes: () => [Array<number>, number] = useCallback(() => {
-    const enabledIndexes = childrenArray
-      .filter(child => !child.props.disabled)
-      .map(child => childrenArray.indexOf(child));
-
-    return [enabledIndexes, enabledIndexes.indexOf(selected!)];
-  }, [childrenArray, selected]);
+  const tabTitleElements = tabDescendants.map(
+    descendant => descendant.element.parentNode as HTMLElement,
+  );
+  const getActiveAndEnabledIndices: () => {
+    activeIndex: number;
+    enabledIndices: Array<number>;
+  } = useCallback(() => {
+    const enabledIndices = tabTitleElements
+      .filter(tabTitleEl => !tabTitleEl.hasAttribute('disabled'))
+      .map(tabTitleEl => tabTitleElements.indexOf(tabTitleEl));
+
+    return {
+      activeIndex: enabledIndices.indexOf(selected),
+      enabledIndices,
+    };
+  }, [selected, tabTitleElements]);
 
   const handleKeyDown = useCallback(
     (e: KeyboardEvent) => {
-      if (!(e.metaKey || e.ctrlKey)) {
-        if (e.key === keyMap.ArrowRight) {
-          const [enabledIndexes, current] = getEnabledIndexes();
-          setSelected?.(enabledIndexes[(current + 1) % enabledIndexes.length]);
-        } else if (e.key === keyMap.ArrowLeft) {
-          const [enabledIndexes, current] = getEnabledIndexes();
-          setSelected?.(
-            enabledIndexes[
-              (current - 1 + enabledIndexes.length) % enabledIndexes.length
-            ],
-          );
-        }
-      }
+      if (e.metaKey || e.ctrlKey) return;
+
+      if (e.key !== keyMap.ArrowRight && e.key !== keyMap.ArrowLeft) return;
+
+      const { activeIndex, enabledIndices } = getActiveAndEnabledIndices();
+      const numberOfEnabledTabs = enabledIndices.length;
+      const indexToUpdateTo =
+        enabledIndices[
+        (e.key === keyMap.ArrowRight
+          ? activeIndex + 1
+          : activeIndex - 1 + numberOfEnabledTabs) % numberOfEnabledTabs
+        ];
+      setSelected?.(indexToUpdateTo);
+      tabTitleElements[indexToUpdateTo].focus();
     },
-    [getEnabledIndexes, setSelected],
+    [getActiveAndEnabledIndices, setSelected, tabTitleElements],
   );
 
+  useEffect(() => {
+    if (
+      typeof controlledSelected === 'undefined' ||
+      tabTitleElements.length === 0 ||
+      tabTitleElements[controlledSelected] === undefined
+    ) {
+      return;
+    }
+
+    tabTitleElements[controlledSelected].focus();
+  }, [controlledSelected, tabTitleElements]);
+
   const renderedTabs = React.Children.map(children, child => {
     if (!isComponentType(child, 'Tab')) {
       return child;
@@ -157,8 +184,16 @@ const Tabs = (props: AccessibleTabsProps) => {
 
   return (
     <LeafyGreenProvider baseFontSize={baseFontSize === 16 ? 16 : 14}>
-      <TabDescendantsProvider>
-        <TabPanelDescendantsProvider>
+      <DescendantsProvider
+        context={TabDescendantsContext}
+        descendants={tabDescendants}
+        dispatch={tabDispatch}
+      >
+        <DescendantsProvider
+          context={TabPanelDescendantsContext}
+          descendants={tabPanelDescendants}
+          dispatch={tabPanelDispatch}
+        >
           <div {...rest} className={className} data-lgid={dataLgId}>
             <div className={tabContainerStyle} id={id}>
               <div
@@ -186,8 +221,8 @@ const Tabs = (props: AccessibleTabsProps) => {
               {renderedTabPanels}
             </div>
           </div>
-        </TabPanelDescendantsProvider>
-      </TabDescendantsProvider>
+        </DescendantsProvider>
+      </DescendantsProvider>
     </LeafyGreenProvider>
   );
 };

packages/tabs/src/context/TabDescendantsContext.ts

@@ -0,0 +1,17 @@
+import {
+  createDescendantsContext,
+  useDescendantsContext,
+} from '@leafygreen-ui/descendants';
+
+export const TabDescendantsContext = createDescendantsContext<HTMLDivElement>(
+  'TabDescendantsContext',
+);
+
+/**
+ * Access list of tab descendants
+ */
+export function useTabDescendantsContext() {
+  const { descendants } = useDescendantsContext(TabDescendantsContext);
+
+  return { tabDescendants: descendants };
+}

packages/tabs/src/context/TabPanelDescendantsContext.ts

@@ -0,0 +1,16 @@
+import {
+  createDescendantsContext,
+  useDescendantsContext,
+} from '@leafygreen-ui/descendants';
+
+export const TabPanelDescendantsContext =
+  createDescendantsContext<HTMLDivElement>('TabPanelsDescendantsContext');
+
+/**
+ * Access list of tab panel descendants
+ */
+export function useTabPanelDescendantsContext() {
+  const { descendants } = useDescendantsContext(TabPanelDescendantsContext);
+
+  return { tabPanelDescendants: descendants };
+}