TransitionGroup.tsx

// Released under the MIT License

import * as React from 'react';

export enum TransitionStates {
  // `MOUNTED` means the child was mounted along with the parent transition group
  MOUNTED,
  // `ENTERED` means the child item was added AFTER the parent was mounted
  ENTERED,
  // `YEETED` implies the element has been removed, but we are still waiting
  // for its cleanup callback to be fired
  YEETED,
}

function wrapChildrenDefault(children: React.ReactNode): React.ReactNode {
  return children;
}

export type TransitionGroupRenderItem<T> = (
  key: string,
  item: T,
  state: TransitionStates,
  cleanUp: () => void,
) => React.ReactNode;

interface ItemCache<T> {
  item: T;
  children: React.ReactNode | null;
  state: TransitionStates;
  cleanUp: () => void;
  renderItem: TransitionGroupRenderItem<T>;
}

type ItemCacheMap<T> = Map<string, ItemCache<T>>;

export interface TransitionGroupProps<T> {
  items: T[];
  renderItem: TransitionGroupRenderItem<T>;
  getItemKey(item: T): string;
  wrapChildren?: (children: React.ReactNode, items: T[]) => React.ReactNode;
  cleanUpDebounceTimeout?: number;
}

export function TransitionGroup<T>({
  items,
  renderItem,
  getItemKey,
  wrapChildren = wrapChildrenDefault,
  cleanUpDebounceTimeout,
}: TransitionGroupProps<T>) {
  const debounceRef = React.useRef(-1);
  React.useLayoutEffect(() => () => clearTimeout(debounceRef.current));
  const [, forceUpdate] = React.useState({});
  // `null` value here allows us to determine the MOUNTED render
  const previousItemCache = React.useRef<ItemCacheMap<T> | null>(null);
  const itemCache = React.useMemo<ItemCacheMap<T>>(() => {
    // Create id Set to track deleted items
    const previousItems = new Set(previousItemCache.current?.keys());
    // Duplicate itemCache to maintain sort order -- items order does not matter
    const itemCache: ItemCacheMap<T> = new Map(previousItemCache.current);
    for (const item of items) {
      const key = getItemKey(item);
      let itemData = itemCache.get(key);
      // If the item doesn't exist, we should create it
      if (itemData == null) {
        const state = previousItemCache.current != null ? TransitionStates.ENTERED : TransitionStates.MOUNTED;
        const cleanUp = () => {
          const item = previousItemCache.current?.get(key);
          if (item == null) {
            // Silently do nothing if the element has already been YEETED - in
            // practice when using things like animation completion handlers,
            // it's often possible to have multiple cleanUp triggers, and it
            // would be much nicer to not have to fix this at the call sites
          } else if (item.state === TransitionStates.YEETED) {
            // Actually clean up the item and re-render
            previousItemCache.current?.delete(key);
            if (cleanUpDebounceTimeout != null) {
              clearTimeout(debounceRef.current);
              debounceRef.current = setTimeout(() => forceUpdate({}), cleanUpDebounceTimeout);
            } else {
              forceUpdate({});
            }
          } else {
            // If we got in here it's because cleanUp was used improperly
            process.env.NODE_ENV === 'development' &&
              console.warn(`TransitionGroup.cleanUp: Attempted to remove an item that isn't yeetable: ${key}`);
          }
        };
        const children = renderItem(key, item, state, cleanUp);
        itemData = {item, children, state, cleanUp, renderItem};
      }
      // Only re-render children if item, renderItem or state has changed
      else if (
        itemData.item !== item ||
        itemData.renderItem !== renderItem ||
        // If this equates to true, it means a YEETED item was come back
        itemData.state === TransitionStates.YEETED
      ) {
        const {cleanUp} = itemData;
        const state = itemData.state === TransitionStates.YEETED ? TransitionStates.ENTERED : itemData.state;
        const children = renderItem(key, item, state, itemData.cleanUp);
        itemData = {item, children, state, cleanUp, renderItem};
      }
      itemCache.set(key, itemData);
      previousItems.delete(key);
    }
    // Iterate through all the items to yeet
    for (const key of previousItems) {
      let itemData = itemCache.get(key);
      if (itemData == null) continue;
      // The item needs to be converted into a yeeted state and/or re-rendered
      if (itemData.state !== TransitionStates.YEETED || itemData.renderItem !== renderItem) {
        const {item, cleanUp} = itemData;
        const children = renderItem(key, itemData.item, TransitionStates.YEETED, itemData.cleanUp);
        const state = TransitionStates.YEETED;
        itemData = {item, children, state, cleanUp, renderItem};
        if (itemData.children != null) {
          itemCache.set(key, itemData);
        } else {
          itemCache.delete(key);
        }
      }
      // An item that has been YEETED but not yet run cleanUp -- we need to
      // hold onto it but there's nothing to do otherwise
      else {
        itemCache.set(key, itemData);
      }
    }
    return itemCache;
  }, [items, getItemKey, renderItem, cleanUpDebounceTimeout]);
  React.useInsertionEffect(() => {
    previousItemCache.current = itemCache;
    // Lets try to ensure we keep leeks to a minimum, ok?
    return () => previousItemCache.current?.clear();
  }, [itemCache]);
  const children: React.ReactNode[] = [];
  for (const [, item] of itemCache) {
    children.push(item.children);
  }
  return <>{children.length > 0 ? wrapChildren(children, items) : null}</>;
}