Improve re-render checking using memo in TransitionPresence

Custom checks that happened internally are now correctly handled by
memo. Which also performs better, since those checks don't have to be
done during internal updates.

I left in a bunch of logging (commented out) since we'll be changing
this component more often in the near future.

Fixes #92

packages/react-transition-presence/src/TransitionPresence/TransitionPresence.stories.tsx

@@ -30,23 +30,30 @@ function useRerender(): () => void {
 function Child({ background, onClick, duration = 1, delayIn = 0 }: ChildProps): ReactElement {
   const ref = useRef<HTMLButtonElement>(null);
 
+  const debugLog = useCallback(
+    (message: string): void => {
+      console.log(`%c${message}: ${background}`, `color: ${background}`);
+    },
+    [background],
+  );
+
   useAnimation(() => {
-    console.log('animate-in', background);
+    debugLog(` >> animate-in`);
     return gsap.fromTo(ref.current, { opacity: 0 }, { opacity: 1, duration, delay: delayIn });
   }, []);
 
   // show visible animation during "before unmount" lifecycle
   useBeforeUnmount(async () => {
-    console.log('animate-out', background);
+    debugLog(` << animate-out`);
     return gsap.fromTo(ref.current, { opacity: 1 }, { opacity: 0, duration });
   }, []);
 
   // show when mounted/unmounted
   useEffect(() => {
-    console.log('mounted', background);
+    debugLog('-> mounted');
 
     return () => {
-      console.log('unmounted', background);
+      debugLog('<- unmounted');
     };
   }, []);
 
@@ -94,6 +101,15 @@ export function DeferFlow(): ReactElement {
       </TransitionPresence>
 
       <div style={{ marginTop: 24 }}>Click the square (isRedVisible: {String(isRedVisible)})</div>
+      <button
+        type="button"
+        /* eslint-disable-next-line react/jsx-no-bind */
+        onClick={(): void => {
+          setIsRedVisible((previous) => !previous);
+        }}
+      >
+        Or click here to switch
+      </button>
     </>
   );
 }

packages/react-transition-presence/src/TransitionPresence/TransitionPresence.tsx

@@ -7,7 +7,7 @@ import {
   useState,
   type ReactElement,
   type Key,
-  useRef,
+  memo,
 } from 'react';
 import {
   useBeforeUnmount,
@@ -27,16 +27,12 @@ export type TransitionPresenceProps = {
  * crossFlow is enabled new children are added immediately when deferred
  * children are not the same.
  */
-export function TransitionPresence({
+function TransitionPresenceInternal({
   children,
   crossFlow,
   onStart,
   onComplete,
 }: TransitionPresenceProps): ReactElement {
-  const lastRenderChildren = useRef<
-    [children: ReactElement | null, deferredChildren: ReactElement | null]
-  >([null, null]);
-  const isReRender = useRef(false);
   const beforeUnmountCallbacks = useMemo(() => new Set<BeforeUnmountCallback>(), []);
   const [deferredChildren, setDeferredChildren] = useState<ReactElement | null>(children);
 
@@ -53,24 +49,9 @@ export function TransitionPresence({
     [beforeUnmountCallbacks],
   );
 
-  // Check if the children or deferredChildren have changed since the previous render
-  isReRender.current = ((): boolean => {
-    const [previousChildren, previousDeferredChildren] = lastRenderChildren.current;
-    lastRenderChildren.current = [children, deferredChildren];
-    return (
-      areChildrenEqual(previousChildren, children) &&
-      areChildrenEqual(previousDeferredChildren, deferredChildren)
-    );
-  })();
-
   useEffect(() => {
-    if (isReRender.current) {
-      // The children and the deferred children are the same,
-      // so we should not trigger a new transition.
-      // If we do, we accidentally animate out new children.
-      return;
-    }
     if (areChildrenEqual(children, deferredChildren)) {
+      // console.log('    setDeferredChildren early return');
       setDeferredChildren(children);
       return;
     }
@@ -81,6 +62,7 @@ export function TransitionPresence({
       // Defer children update for before unmount lifecycle
       await beforeUnmount(abortController.signal);
 
+      // console.log('    setDeferredChildren after unMount children');
       setDeferredChildren(children);
       onComplete?.();
     })();
@@ -98,6 +80,17 @@ export function TransitionPresence({
     Children.only(children);
   }
 
+  // console.group('TransitionPresence::render');
+  // console.log(
+  //   `children %c${children?.key}`,
+  //   `color: ${String(children?.key).replaceAll(/\d/gu, '')}`,
+  // );
+  // console.log(
+  //   `deferredChildren %c${deferredChildren?.key}`,
+  //   `color: ${String(deferredChildren?.key).replaceAll(/\d/gu, '')}`,
+  // );
+  // console.groupEnd();
+
   const shouldRenderOldChildren = crossFlow && !areChildrenEqual(children, deferredChildren, false);
 
   return (
@@ -108,6 +101,26 @@ export function TransitionPresence({
   );
 }
 
+// The beforeUnmount behaviour in the useEffect hook is relying on the fact that
+// the TransitionPresence only renders when the children are different.
+// However, the parent component can still re-render and pass the same children (with the same key),
+// but these will have different references.
+// If we don't prevent the render in that case, the beforeUnmount will be called on the new children.
+// This is why we need a custom comparison function that compares the keys of the children,
+// instead of the references.
+export const TransitionPresence = memo(TransitionPresenceInternal, (prevProps, nextProps) => {
+  const { children, ...propsToCompare } = prevProps;
+  const areEqual =
+    // custom children comparison, on keys
+    areChildrenEqual(children, nextProps.children) &&
+    // compare all other props like react does
+    Object.keys(propsToCompare).every((key) =>
+      Object.is(prevProps[key as keyof typeof prevProps], nextProps[key as keyof typeof nextProps]),
+    );
+  // console.log('TransitionPresence::memo', areEqual);
+  return areEqual;
+});
+
 function areChildrenEqual(
   childrenA: ReactElement | null,
   childrenB: ReactElement | null,