Improve performance of Listbox and Menu when closing (#3690)

In a previous PR, we already batched registering options. This PR also
batches unregistering options to make the closing behavior smoother when
there are a lot of items rendered.

Author: RobinMalfait

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

@@ -72,7 +72,7 @@ export enum ActionTypes {
   ClearSearch,
 
   RegisterOptions,
-  UnregisterOption,
+  UnregisterOptions,
 
   SetButtonElement,
   SetOptionsElement,
@@ -124,7 +124,7 @@ type Actions<T> =
       type: ActionTypes.RegisterOptions
       options: { id: string; dataRef: ListboxOptionDataRef<T> }[]
     }
-  | { type: ActionTypes.UnregisterOption; id: string }
+  | { type: ActionTypes.UnregisterOptions; options: string[] }
   | { type: ActionTypes.SetButtonElement; element: HTMLButtonElement | null }
   | { type: ActionTypes.SetOptionsElement; element: HTMLElement | null }
   | { type: ActionTypes.SortOptions }
@@ -328,12 +328,24 @@ let reducers: {
       pendingShouldSort: true,
     }
   },
-  [ActionTypes.UnregisterOption]: (state, action) => {
+  [ActionTypes.UnregisterOptions]: (state, action) => {
     let options = state.options
-    let idx = options.findIndex((a) => a.id === action.id)
-    if (idx !== -1) {
+
+    let idxs = []
+    let ids = new Set(action.options)
+    for (let [idx, option] of options.entries()) {
+      if (ids.has(option.id)) {
+        idxs.push(idx)
+        ids.delete(option.id)
+        if (ids.size === 0) break
+      }
+    }
+
+    if (idxs.length > 0) {
       options = options.slice()
-      options.splice(idx, 1)
+      for (let idx of idxs.reverse()) {
+        options.splice(idx, 1)
+      }
     }
 
     return {
@@ -421,6 +433,15 @@ export class ListboxMachine<T> extends Machine<State<T>, Actions<T>> {
         },
       ]
     }),
+    unregisterOption: batch(() => {
+      let options: string[] = []
+      return [
+        (id: string) => options.push(id),
+        () => {
+          this.send({ type: ActionTypes.UnregisterOptions, options: options.splice(0) })
+        },
+      ]
+    }),
     goToOption: batch(() => {
       let last: Extract<Actions<unknown>, { type: ActionTypes.GoToOption }> | null = null
       return [

packages/@headlessui-react/src/components/listbox/listbox.tsx

@@ -850,9 +850,7 @@ function OptionFn<
   useIsoMorphicEffect(() => {
     if (usedInSelectedOption) return
     machine.actions.registerOption(id, bag)
-    return () => {
-      machine.send({ type: ActionTypes.UnregisterOption, id })
-    }
+    return () => machine.actions.unregisterOption(id)
   }, [bag, id, usedInSelectedOption])
 
   let handleClick = useEvent((event: { preventDefault: Function }) => {

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

@@ -45,7 +45,7 @@ export enum ActionTypes {
   Search,
   ClearSearch,
   RegisterItems,
-  UnregisterItem,
+  UnregisterItems,
 
   SetButtonElement,
   SetItemsElement,
@@ -95,7 +95,7 @@ export type Actions =
   | { type: ActionTypes.Search; value: string }
   | { type: ActionTypes.ClearSearch }
   | { type: ActionTypes.RegisterItems; items: { id: string; dataRef: MenuItemDataRef }[] }
-  | { type: ActionTypes.UnregisterItem; id: string }
+  | { type: ActionTypes.UnregisterItems; items: string[] }
   | { type: ActionTypes.SetButtonElement; element: HTMLButtonElement | null }
   | { type: ActionTypes.SetItemsElement; element: HTMLElement | null }
   | { type: ActionTypes.SortItems }
@@ -283,12 +283,24 @@ let reducers: {
       pendingShouldSort: true,
     }
   },
-  [ActionTypes.UnregisterItem]: (state, action) => {
+  [ActionTypes.UnregisterItems]: (state, action) => {
     let items = state.items
-    let idx = items.findIndex((a) => a.id === action.id)
-    if (idx !== -1) {
+
+    let idxs = []
+    let ids = new Set(action.items)
+    for (let [idx, item] of items.entries()) {
+      if (ids.has(item.id)) {
+        idxs.push(idx)
+        ids.delete(item.id)
+        if (ids.size === 0) break
+      }
+    }
+
+    if (idxs.length > 0) {
       items = items.slice()
-      items.splice(idx, 1)
+      for (let idx of idxs.reverse()) {
+        items.splice(idx, 1)
+      }
     }
 
     return {
@@ -297,6 +309,7 @@ let reducers: {
       activationTrigger: ActivationTrigger.Other,
     }
   },
+
   [ActionTypes.SetButtonElement]: (state, action) => {
     if (state.buttonElement === action.element) return state
     return { ...state, buttonElement: action.element }
@@ -359,6 +372,14 @@ export class MenuMachine extends Machine<State, Actions> {
         () => this.send({ type: ActionTypes.RegisterItems, items: items.splice(0) }),
       ]
     }),
+    unregisterItem: batch(() => {
+      let items: string[] = []
+
+      return [
+        (id: string) => items.push(id),
+        () => this.send({ type: ActionTypes.UnregisterItems, items: items.splice(0) }),
+      ]
+    }),
   }
 
   selectors = {

packages/@headlessui-react/src/components/menu/menu.tsx

@@ -620,7 +620,7 @@ function ItemFn<TTag extends ElementType = typeof DEFAULT_ITEM_TAG>(
 
   useIsoMorphicEffect(() => {
     machine.actions.registerItem(id, bag)
-    return () => machine.send({ type: ActionTypes.UnregisterItem, id })
+    return () => machine.actions.unregisterItem(id)
   }, [bag, id])
 
   let close = useEvent(() => {