Transition improvements

packages/@headlessui-react/src/components/transitions/transition.tsx

@@ -1,7 +1,6 @@
 import React, {
   Fragment,
   createContext,
-  useCallback,
   useContext,
   useEffect,
   useMemo,
@@ -24,6 +23,7 @@ import { Features, PropsForFeatures, render, RenderStrategy } from '../../utils/
 import { Reason, transition } from './utils/transition'
 import { OpenClosedProvider, State, useOpenClosed } from '../../internal/open-closed'
 import { useServerHandoffComplete } from '../../hooks/use-server-handoff-complete'
+import { useLatestValue } from '../../hooks/use-latest-value'
 
 type ID = ReturnType<typeof useId>
 
@@ -95,8 +95,8 @@ function useParentNesting() {
 
 interface NestingContextValues {
   children: MutableRefObject<{ id: ID; state: TreeStates }[]>
-  register: (id: ID) => () => void
-  unregister: (id: ID, strategy?: RenderStrategy) => void
+  register: MutableRefObject<(id: ID) => () => void>
+  unregister: MutableRefObject<(id: ID, strategy?: RenderStrategy) => void>
 }
 
 let NestingContext = createContext<NestingContextValues | null>(null)
@@ -110,48 +110,38 @@ function hasChildren(
 }
 
 function useNesting(done?: () => void) {
-  let doneRef = useRef(done)
+  let doneRef = useLatestValue(done)
   let transitionableChildren = useRef<NestingContextValues['children']['current']>([])
   let mounted = useIsMounted()
 
-  useEffect(() => {
-    doneRef.current = done
-  }, [done])
-
-  let unregister = useCallback(
-    (childId: ID, strategy = RenderStrategy.Hidden) => {
-      let idx = transitionableChildren.current.findIndex(({ id }) => id === childId)
-      if (idx === -1) return
-
-      match(strategy, {
-        [RenderStrategy.Unmount]() {
-          transitionableChildren.current.splice(idx, 1)
-        },
-        [RenderStrategy.Hidden]() {
-          transitionableChildren.current[idx].state = TreeStates.Hidden
-        },
-      })
-
-      if (!hasChildren(transitionableChildren) && mounted.current) {
-        doneRef.current?.()
-      }
-    },
-    [doneRef, mounted, transitionableChildren]
-  )
+  let unregister = useLatestValue((childId: ID, strategy = RenderStrategy.Hidden) => {
+    let idx = transitionableChildren.current.findIndex(({ id }) => id === childId)
+    if (idx === -1) return
+
+    match(strategy, {
+      [RenderStrategy.Unmount]() {
+        transitionableChildren.current.splice(idx, 1)
+      },
+      [RenderStrategy.Hidden]() {
+        transitionableChildren.current[idx].state = TreeStates.Hidden
+      },
+    })
 
-  let register = useCallback(
-    (childId: ID) => {
-      let child = transitionableChildren.current.find(({ id }) => id === childId)
-      if (!child) {
-        transitionableChildren.current.push({ id: childId, state: TreeStates.Visible })
-      } else if (child.state !== TreeStates.Visible) {
-        child.state = TreeStates.Visible
-      }
-
-      return () => unregister(childId, RenderStrategy.Unmount)
-    },
-    [transitionableChildren, unregister]
-  )
+    if (!hasChildren(transitionableChildren) && mounted.current) {
+      doneRef.current?.()
+    }
+  })
+
+  let register = useLatestValue((childId: ID) => {
+    let child = transitionableChildren.current.find(({ id }) => id === childId)
+    if (!child) {
+      transitionableChildren.current.push({ id: childId, state: TreeStates.Visible })
+    } else if (child.state !== TreeStates.Visible) {
+      child.state = TreeStates.Visible
+    }
+
+    return () => unregister.current(childId, RenderStrategy.Unmount)
+  })
 
   return useMemo(
     () => ({
@@ -227,14 +217,14 @@ function TransitionChild<TTag extends ElementType = typeof DEFAULT_TRANSITION_CH
     // transitioning ourselves. Otherwise we would unmount before the transitions are finished.
     if (!isTransitioning.current) {
       setState(TreeStates.Hidden)
-      unregister(id)
+      unregister.current(id)
       events.current.afterLeave()
     }
   })
 
   useIsoMorphicEffect(() => {
     if (!id) return
-    return register(id)
+    return register.current(id)
   }, [register, id])
 
   useIsoMorphicEffect(() => {
@@ -249,8 +239,8 @@ function TransitionChild<TTag extends ElementType = typeof DEFAULT_TRANSITION_CH
     }
 
     match(state, {
-      [TreeStates.Hidden]: () => unregister(id),
-      [TreeStates.Visible]: () => register(id),
+      [TreeStates.Hidden]: () => unregister.current(id),
+      [TreeStates.Visible]: () => register.current(id),
     })
   }, [state, id, register, unregister, show, strategy])
 
@@ -314,7 +304,7 @@ function TransitionChild<TTag extends ElementType = typeof DEFAULT_TRANSITION_CH
             // ourselves.
             if (!hasChildren(nesting)) {
               setState(TreeStates.Hidden)
-              unregister(id)
+              unregister.current(id)
               events.current.afterLeave()
             }
           }