Performance improvement: only re-render top-level component (#3722)

This PR fixes a performance issue where all components using the
`useIsTopLayer` hook will re-render when the hook changes.

For context, the internal hook is used to know which component is the
top most component. This is important in a situation like this:

```
<Dialog>
  <Menu />
</Dialog>
```

If the Menu inside the Dialog is open, it is considered the top most
component. Clicking outside of the Menu or pressing escape should only
close the Menu and not the Dialog.

This behavior is similar to the native `#top-layer` you see when using
native dialogs for example.

The issue however is that the `useIsTopLayer` subscribes to an external
store which is shared across all components. This means that when the
store changes, all components using the hook will re-render.

To make things worse, since we can't use these hooks unconditionally,
they will all be subscribed to the store even if the Menu component(s)
are not open.

To solve this, we will use a new state machine and use the `useMachine`
hook. This internally uses a `useSyncExternalStoreWithSelector` to
subscribe to the store.

This means that the component will only re-render if the state computed
by the selector changes.

This now means that at most 2 components will re-render when the store
changes:

1. The component that _was_ in the top most position
2. The component that is going to be in the top most position

Fixes: #3630
Closes: #3662


# Test plan

Behavior before: notice how all Menu components re-render:


https://github.com/user-attachments/assets/3172b632-0fa4-42db-970c-39efc827dd84

After this change, only the Menu that was opened / closed will
re-render:


https://github.com/user-attachments/assets/5d254bfc-5233-47a7-94d3-eb7a8593e14f

Author: RobinMalfait

packages/@headlessui-react/CHANGELOG.md

@@ -17,6 +17,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Ensure clicking on interactive elements inside `Label` component works ([#3709](https://github.com/tailwindlabs/headlessui/pull/3709))
 - Fix focus not returned to SVG Element ([#3704](https://github.com/tailwindlabs/headlessui/pull/3704))
 - Fix `Listbox` not focusing first or last option on ArrowUp / ArrowDown ([#3721](https://github.com/tailwindlabs/headlessui/pull/3721))
+- Performance improvement: only re-render top-level component when nesting components e.g.: `Menu` inside a `Dialog` ([#3722](https://github.com/tailwindlabs/headlessui/pull/3722))
 
 ## [2.2.2] - 2025-04-17
 

packages/@headlessui-react/src/hooks/use-is-top-layer.ts

@@ -1,27 +1,7 @@
-import { useId } from 'react'
-import { DefaultMap } from '../utils/default-map'
-import { createStore } from '../utils/store'
+import { useCallback, useId } from 'react'
+import { stackMachines } from '../machines/stack-machine'
+import { useSlice } from '../react-glue'
 import { useIsoMorphicEffect } from './use-iso-morphic-effect'
-import { useStore } from './use-store'
-
-/**
- * Map of stable hierarchy stores based on a given scope.
- */
-let hierarchyStores = new DefaultMap(() =>
-  createStore(() => [] as string[], {
-    ADD(id: string) {
-      if (this.includes(id)) return this
-      return [...this, id]
-    },
-    REMOVE(id: string) {
-      let idx = this.indexOf(id)
-      if (idx === -1) return this
-      let copy = this.slice()
-      copy.splice(idx, 1)
-      return copy
-    },
-  })
-)
 
 /**
  * A hook that returns whether the current node is on the top of the hierarchy,
@@ -46,32 +26,41 @@ let hierarchyStores = new DefaultMap(() =>
  * </Dialog>
  * ```
  */
-export function useIsTopLayer(enabled: boolean, scope: string) {
-  let hierarchyStore = hierarchyStores.get(scope)
+export function useIsTopLayer(enabled: boolean, scope: string | null) {
   let id = useId()
-  let hierarchy = useStore(hierarchyStore)
+  let stackMachine = stackMachines.get(scope)
+
+  let [isTop, onStack] = useSlice(
+    stackMachine,
+    useCallback(
+      (state) => [
+        stackMachine.selectors.isTop(state, id),
+        stackMachine.selectors.inStack(state, id),
+      ],
+      [stackMachine, id, enabled]
+    )
+  )
 
+  // Depending on the enable state, push/pop the current `id` to/from the
+  // hierarchy.
   useIsoMorphicEffect(() => {
     if (!enabled) return
+    stackMachine.actions.push(id)
+    return () => stackMachine.actions.pop(id)
+  }, [stackMachine, enabled, id])
 
-    hierarchyStore.dispatch('ADD', id)
-    return () => hierarchyStore.dispatch('REMOVE', id)
-  }, [hierarchyStore, enabled])
-
+  // If the hook is not enabled, we know for sure it is not going to tbe the
+  // top-most item.
   if (!enabled) return false
 
-  let idx = hierarchy.indexOf(id)
-  let hierarchyLength = hierarchy.length
-
-  // Not in the hierarchy yet
-  if (idx === -1) {
-    // Assume that it will be inserted at the end, then it means that the `idx`
-    // will be the length of the current hierarchy.
-    idx = hierarchyLength
-
-    // Increase the hierarchy length as-if the node is already in the hierarchy.
-    hierarchyLength += 1
-  }
+  // If the hook is enabled, and it's on the stack, we can rely on the `isTop`
+  // derived state to determine if it's the top-most item.
+  if (onStack) return isTop
 
-  return idx === hierarchyLength - 1
+  // In this scenario, the hook is enabled, but we are not on the stack yet. In
+  // this case we assume that we will be the top-most item, so we return
+  // `true`. However, if that's not the case, and once we are on the stack (or
+  // other items are pushed) this hook will be re-evaluated and the `isTop`
+  // derived state will be used instead.
+  return true
 }

packages/@headlessui-react/src/machine.ts

@@ -8,10 +8,16 @@ export abstract class Machine<State, Event extends { type: number | string }> {
   )
   #subscribers: Set<Subscriber<State, any>> = new Set()
 
+  disposables = disposables()
+
   constructor(initialState: State) {
     this.#state = initialState
   }
 
+  dispose() {
+    this.disposables.dispose()
+  }
+
   get state(): Readonly<State> {
     return this.#state
   }
@@ -29,20 +35,23 @@ export abstract class Machine<State, Event extends { type: number | string }> {
     }
     this.#subscribers.add(subscriber)
 
-    return () => {
+    return this.disposables.add(() => {
       this.#subscribers.delete(subscriber)
-    }
+    })
   }
 
   on(type: Event['type'], callback: (state: State, event: Event) => void) {
     this.#eventSubscribers.get(type).add(callback)
-    return () => {
+    return this.disposables.add(() => {
       this.#eventSubscribers.get(type).delete(callback)
-    }
+    })
   }
 
   send(event: Event) {
-    this.#state = this.reduce(this.#state, event)
+    let newState = this.reduce(this.#state, event)
+    if (newState === this.#state) return // No change
+
+    this.#state = newState
 
     for (let subscriber of this.#subscribers) {
       let slice = subscriber.selector(this.#state)

packages/@headlessui-react/src/machines/stack-machine.ts

@@ -0,0 +1,72 @@
+import { Machine } from '../machine'
+import { DefaultMap } from '../utils/default-map'
+import { match } from '../utils/match'
+
+type Scope = string | null
+type Id = string
+
+interface State {
+  stack: Id[]
+}
+
+export enum ActionTypes {
+  Push,
+  Pop,
+}
+
+export type Actions = { type: ActionTypes.Push; id: Id } | { type: ActionTypes.Pop; id: Id }
+
+let reducers: {
+  [P in ActionTypes]: (state: State, action: Extract<Actions, { type: P }>) => State
+} = {
+  [ActionTypes.Push](state, action) {
+    let id = action.id
+    let stack = state.stack
+    let idx = state.stack.indexOf(id)
+
+    // Already in the stack, move it to the top
+    if (idx !== -1) {
+      let copy = state.stack.slice()
+      copy.splice(idx, 1)
+      copy.push(id)
+
+      stack = copy
+      return { ...state, stack }
+    }
+
+    // Not in the stack, add it to the top
+    return { ...state, stack: [...state.stack, id] }
+  },
+  [ActionTypes.Pop](state, action) {
+    let id = action.id
+    let idx = state.stack.indexOf(id)
+    if (idx === -1) return state // Not in the stack
+
+    let copy = state.stack.slice()
+    copy.splice(idx, 1)
+
+    return { ...state, stack: copy }
+  },
+}
+
+class StackMachine extends Machine<State, Actions> {
+  static new() {
+    return new StackMachine({ stack: [] })
+  }
+
+  reduce(state: Readonly<State>, action: Actions): State {
+    return match(action.type, reducers, state, action)
+  }
+
+  actions = {
+    push: (id: Id) => this.send({ type: ActionTypes.Push, id }),
+    pop: (id: Id) => this.send({ type: ActionTypes.Pop, id }),
+  }
+
+  selectors = {
+    isTop: (state: State, id: Id) => state.stack[state.stack.length - 1] === id,
+    inStack: (state: State, id: Id) => state.stack.includes(id),
+  }
+}
+
+export const stackMachines = new DefaultMap<Scope, StackMachine>(() => StackMachine.new())

playgrounds/react/pages/dialog/dialog.tsx

@@ -95,6 +95,7 @@ export default function Home() {
 
               <Transition.Child
                 as="div"
+                className="relative"
                 enter="ease-out transform duration-300"
                 enterFrom="opacity-0 translate-y-4 sm:translate-y-0 sm:scale-95"
                 enterTo="opacity-100 translate-y-0 sm:scale-100"