LG-4444: OverlayContext and useOverlay (#2454)

* OverlayContext, useOverlayContext, and OverlayProvider

* useOverlay to register, remove, and track isTopMostOverlay

* Export overlay context related components and hooks and include provider in LG provider

* README and changeset

* Fix type and lint

* Update registerOverlay logic and use more explicit generic type names

.changeset/strange-geckos-jam.md

@@ -0,0 +1,5 @@
+---
+'@leafygreen-ui/leafygreen-provider': minor
+---
+
+[LG-4444](https://jira.mongodb.org/browse/LG-4444): Adds `OverlayContext` to track overlays that stack on z-axis. `OverlayProvider` is included in `LeafyGreenProvider`. Adds `useOverlay` hook to register/remove overlays into `OverlayContext`. 

packages/leafygreen-provider/README.md

@@ -111,6 +111,8 @@ function InlineCode({ children, className }: InlineCodeProps) {
 
 This hook is meant for internal use. It allows components to read the value of the dark mode prop from the LeafyGreen provider and overwrite the value locally if necessary.
 
+### Example
+
 ```js
 import { useDarkMode } from '@leafygreen-ui/leafygreen-provider';
 
@@ -119,3 +121,21 @@ function Example({ children, darkMode: darkModeProp, variant }) {
   return <div className={badgeVariants[theme][variant]}>{children}</div>;
 }
 ```
+
+## useOverlay
+
+This hook is used to register/remove a component to the `OverlayContext`.
+
+### Example
+
+```js
+import { useOverlay } from '@leafygreen-ui/leafygreen-provider';
+
+const Popover = forwardRef<HTMLDivElement, PopoverProps>(
+  ({ children }: PopoverProps, fwdRef) => {
+    const { id, isTopMostOverlay, ref } = useOverlay(fwdRef);
+
+    return <div ref={ref}>{children}</div>
+  },
+);
+```

packages/leafygreen-provider/src/LeafyGreenContext.tsx

@@ -4,6 +4,7 @@ import PropTypes from 'prop-types';
 import { DarkModeProps } from '@leafygreen-ui/lib';
 
 import DarkModeProvider, { useDarkModeContext } from './DarkModeContext';
+import { OverlayProvider } from './OverlayContext';
 import PortalContextProvider, {
   PortalContextValues,
   usePopoverPortalContainer,
@@ -55,7 +56,7 @@ function LeafyGreenProvider({
             contextDarkMode={darkModeState}
             setDarkMode={setDarkMode}
           >
-            {children}
+            <OverlayProvider>{children}</OverlayProvider>
           </DarkModeProvider>
         </TypographyProvider>
       </PortalContextProvider>

packages/leafygreen-provider/src/OverlayContext/OverlayContext.spec.tsx

@@ -0,0 +1,99 @@
+import React, { createRef, useState } from 'react';
+import { render } from '@testing-library/react';
+import userEvent from '@testing-library/user-event';
+
+import { renderHook } from '@leafygreen-ui/testing-lib';
+
+import OverlayProvider, { useOverlayContext } from './OverlayContext';
+
+function TestComponent() {
+  const { getOverlays, registerOverlay, removeOverlay, topMostOverlay } =
+    useOverlayContext();
+  const [id, setId] = useState(1);
+
+  const registerNewOverlay = () => {
+    registerOverlay({
+      element: document.createElement('div'),
+      id: '' + id,
+      ref: createRef<HTMLElement>(),
+    });
+    setId(prev => prev + 1);
+  };
+
+  const removeTopMostOverlay = () => {
+    if (!topMostOverlay) return;
+
+    removeOverlay(topMostOverlay.id);
+    setId(prev => prev - 1);
+  };
+
+  return (
+    <div>
+      <span>Overlays count: {getOverlays().length}</span>
+      <span>Top most overlay ID: {topMostOverlay?.id}</span>
+      <button onClick={registerNewOverlay}>Register overlay</button>
+      <button onClick={removeTopMostOverlay}>Remove overlay</button>
+    </div>
+  );
+}
+
+function renderProviderWithChildren() {
+  return render(
+    <OverlayProvider>
+      <TestComponent />
+    </OverlayProvider>,
+  );
+}
+
+describe('OverlayProvider', () => {
+  test('registers an overlay', () => {
+    const { getByText } = renderProviderWithChildren();
+
+    expect(getByText('Overlays count: 0')).toBeInTheDocument();
+
+    userEvent.click(getByText('Register overlay'));
+
+    expect(getByText('Overlays count: 1')).toBeInTheDocument();
+    expect(getByText('Top most overlay ID: 1')).toBeInTheDocument();
+  });
+
+  test('removes an overlay', () => {
+    const { getByText } = renderProviderWithChildren();
+
+    userEvent.click(getByText('Register overlay'));
+    userEvent.click(getByText('Register overlay'));
+
+    expect(getByText('Overlays count: 2')).toBeInTheDocument();
+    expect(getByText('Top most overlay ID: 2')).toBeInTheDocument();
+
+    userEvent.click(getByText('Remove overlay'));
+
+    expect(getByText('Overlays count: 1')).toBeInTheDocument();
+    expect(getByText('Top most overlay ID: 1')).toBeInTheDocument();
+  });
+});
+
+describe('useOverlayContext', () => {
+  test('provides expected context values when used within OverlayProvider', () => {
+    const wrapper = ({ children }) => (
+      <OverlayProvider>{children}</OverlayProvider>
+    );
+    const { result } = renderHook(useOverlayContext, { wrapper });
+
+    expect(result.current).toHaveProperty('getOverlays');
+    expect(result.current).toHaveProperty('registerOverlay');
+    expect(result.current).toHaveProperty('removeOverlay');
+    expect(result.current).toHaveProperty('topMostOverlay');
+  });
+
+  test('throws an error when used outside of OverlayProvider', () => {
+    try {
+      renderHook(useOverlayContext);
+    } catch (error) {
+      expect(error).toBeInstanceOf(Error);
+      expect(error.message).toBe(
+        'useOverlayContext must be used within an OverlayProvider',
+      );
+    }
+  });
+});

packages/leafygreen-provider/src/OverlayContext/OverlayContext.tsx

@@ -0,0 +1,79 @@
+import React, {
+  createContext,
+  ReactNode,
+  RefObject,
+  useCallback,
+  useContext,
+  useMemo,
+  useState,
+} from 'react';
+import PropTypes from 'prop-types';
+
+interface OverlayItem<TOverlayElement extends HTMLElement = HTMLElement> {
+  element: TOverlayElement;
+  ref: RefObject<TOverlayElement>;
+  id: string;
+}
+
+export interface OverlayContextProps {
+  getOverlays: () => Array<OverlayItem>;
+  registerOverlay: (overlay: OverlayItem) => void;
+  removeOverlay: (idToRemove: string) => void;
+  topMostOverlay: OverlayItem | undefined;
+}
+
+export const OverlayContext = createContext<OverlayContextProps | undefined>(
+  undefined,
+);
+
+export const useOverlayContext = (): OverlayContextProps => {
+  const context = useContext(OverlayContext);
+
+  if (!context) {
+    throw new Error('useOverlayContext must be used within an OverlayProvider');
+  }
+
+  return context;
+};
+
+/**
+ * Global provider tracking overlays that stack on z-axis
+ */
+const OverlayProvider = ({ children }: { children: ReactNode }) => {
+  const [overlays, setOverlays] = useState<Array<OverlayItem>>([]);
+
+  const topMostOverlay = useMemo(
+    () => overlays[overlays.length - 1],
+    [overlays],
+  );
+
+  const registerOverlay = useCallback((overlay: OverlayItem) => {
+    if (overlays.some(item => item.id === overlay.id)) return;
+    setOverlays(prev => [...prev, { ...overlay }]);
+  }, []);
+
+  const removeOverlay = useCallback((idToRemove: string) => {
+    setOverlays(prev => prev.filter(item => item.id !== idToRemove));
+  }, []);
+
+  return (
+    <OverlayContext.Provider
+      value={{
+        getOverlays: () => overlays,
+        registerOverlay,
+        removeOverlay,
+        topMostOverlay,
+      }}
+    >
+      {children}
+    </OverlayContext.Provider>
+  );
+};
+
+OverlayProvider.displayName = 'OverlayProvider';
+
+OverlayProvider.propTypes = {
+  children: PropTypes.node,
+};
+
+export default OverlayProvider;

packages/leafygreen-provider/src/OverlayContext/index.ts

@@ -0,0 +1,7 @@
+export {
+  OverlayContext,
+  type OverlayContextProps,
+  default as OverlayProvider,
+  useOverlayContext,
+} from './OverlayContext';
+export { useOverlay } from './useOverlay';

packages/leafygreen-provider/src/OverlayContext/useOverlay.spec.tsx

@@ -0,0 +1,92 @@
+import { renderHook } from '@leafygreen-ui/testing-lib';
+
+import * as OverlayContextModule from './OverlayContext';
+import OverlayProvider from './OverlayContext';
+import { useOverlay } from './useOverlay';
+
+jest.mock('@leafygreen-ui/hooks', () => ({
+  ...jest.requireActual('@leafygreen-ui/hooks'),
+  useForwardedRef: jest.fn(ref => ref),
+  useIdAllocator: jest.fn(() => 'mock-id'),
+}));
+
+describe('useOverlay', () => {
+  let mockRegisterOverlay;
+  let mockRemoveOverlay;
+  let mockTopMostOverlay;
+
+  beforeEach(() => {
+    mockRegisterOverlay = jest.fn();
+    mockRemoveOverlay = jest.fn();
+    mockTopMostOverlay = undefined;
+
+    jest
+      .spyOn(OverlayContextModule, 'useOverlayContext')
+      .mockImplementation(() => ({
+        getOverlays: jest.fn(),
+        registerOverlay: mockRegisterOverlay,
+        removeOverlay: mockRemoveOverlay,
+        topMostOverlay: mockTopMostOverlay,
+      }));
+
+    document.contains = jest.fn(() => true);
+  });
+
+  afterEach(() => {
+    jest.clearAllMocks();
+  });
+
+  test('registers overlay when component mounts', () => {
+    const ref = { current: document.createElement('div') };
+    renderHook(() => useOverlay(ref), { wrapper: OverlayProvider });
+
+    expect(mockRegisterOverlay).toHaveBeenCalledTimes(1);
+    expect(mockRegisterOverlay).toHaveBeenCalledWith({
+      element: ref.current,
+      id: 'mock-id',
+      ref,
+    });
+  });
+
+  test('does not register overlay when ref is not attached', () => {
+    const ref = { current: null };
+    renderHook(() => useOverlay(ref), { wrapper: OverlayProvider });
+
+    expect(mockRegisterOverlay).not.toHaveBeenCalled();
+  });
+
+  test('removes overlay on unmount', () => {
+    mockRemoveOverlay = jest
+      .spyOn(OverlayContextModule.useOverlayContext(), 'removeOverlay')
+      .mockImplementation(jest.fn);
+    const ref = { current: document.createElement('div') };
+    const { unmount } = renderHook(() => useOverlay(ref), {
+      wrapper: OverlayProvider,
+    });
+
+    unmount();
+
+    expect(mockRemoveOverlay).toHaveBeenCalledTimes(1);
+    expect(mockRemoveOverlay).toHaveBeenCalledWith('mock-id');
+  });
+
+  test('returns id and ref', () => {
+    const ref = { current: document.createElement('div') };
+    const { result } = renderHook(() => useOverlay(ref), {
+      wrapper: OverlayProvider,
+    });
+
+    expect(result.current.id).toBe('mock-id');
+    expect(result.current.ref).toBe(ref);
+  });
+
+  test('returns correct isTopMostOverlay value', () => {
+    const ref = { current: document.createElement('div') };
+    mockTopMostOverlay = { id: 'mock-id' };
+    const { result } = renderHook(() => useOverlay(ref), {
+      wrapper: OverlayProvider,
+    });
+
+    expect(result.current.isTopMostOverlay).toBe(true);
+  });
+});

packages/leafygreen-provider/src/OverlayContext/useOverlay.ts

@@ -0,0 +1,49 @@
+import { ForwardedRef, RefObject, useMemo } from 'react';
+
+import {
+  useForwardedRef,
+  useIdAllocator,
+  useIsomorphicLayoutEffect,
+} from '@leafygreen-ui/hooks';
+
+import { useOverlayContext } from './OverlayContext';
+
+/**
+ * Registers a component as an overlay in the {@link OverlayContext} on mount.
+ * Removes the overlay component on unmount. Checks if `isTopMostOverlay`
+ */
+export const useOverlay = <TOverlayElement extends HTMLElement>(
+  fwdRef: RefObject<TOverlayElement> | ForwardedRef<TOverlayElement>,
+) => {
+  const ref: RefObject<TOverlayElement> = useForwardedRef(fwdRef, null);
+  const { topMostOverlay, registerOverlay, removeOverlay } =
+    useOverlayContext();
+  const id = useIdAllocator({ prefix: 'overlay' });
+
+  useIsomorphicLayoutEffect(() => {
+    const refAttached = ref.current !== null;
+    const refExists = document.contains(ref.current);
+
+    if (refAttached && refExists) {
+      registerOverlay({
+        element: ref.current,
+        id,
+        ref,
+      });
+    }
+
+    return () => {
+      removeOverlay(id);
+    };
+  }, [ref.current]);
+
+  const isTopMostOverlay = useMemo(() => {
+    return topMostOverlay ? topMostOverlay.id === id : false;
+  }, [topMostOverlay]);
+
+  return {
+    id,
+    isTopMostOverlay,
+    ref,
+  };
+};

packages/leafygreen-provider/src/index.ts

@@ -1,5 +1,12 @@
 export { useDarkMode, useDarkModeContext } from './DarkModeContext';
 export { default, type LeafyGreenProviderProps } from './LeafyGreenContext';
+export {
+  OverlayContext,
+  type OverlayContextProps,
+  OverlayProvider,
+  useOverlay,
+  useOverlayContext,
+} from './OverlayContext';
 export {
   PopoverContext,
   PopoverProvider,