From f35214db4c70c38d19a18abefbed33b6b51359f1 Mon Sep 17 00:00:00 2001
From: Robin Malfait <malfait.robin@gmail.com>
Date: Tue, 30 Apr 2024 19:31:40 +0200
Subject: [PATCH] calculate the size of an element as soon as possible (#3153)

Instead of waiting for a `useEffect` to trigger, we can use `useMemo` to
always compute the size the moment a DOM element is available (or
changes).

We also make sure to force a re-render when resizes are detected on the
stable DOM node, which in turn causes the `useMemo` to recompute.
---
 .../src/components/combobox/combobox.tsx         |  1 +
 .../src/hooks/use-element-size.ts                | 16 ++++++++++------
 2 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/packages/@headlessui-react/src/components/combobox/combobox.tsx b/packages/@headlessui-react/src/components/combobox/combobox.tsx
index 1fdb28f..832dff3 100644
--- a/packages/@headlessui-react/src/components/combobox/combobox.tsx
+++ b/packages/@headlessui-react/src/components/combobox/combobox.tsx
@@ -1644,6 +1644,7 @@ function OptionsFn<TTag extends ElementType = typeof DEFAULT_OPTIONS_TAG>(
       option: undefined,
     } satisfies OptionsRenderPropArg
   }, [data])
+
   let ourProps = mergeProps(anchor ? getFloatingPanelProps() : {}, {
     'aria-labelledby': labelledBy,
     role: 'listbox',
diff --git a/packages/@headlessui-react/src/hooks/use-element-size.ts b/packages/@headlessui-react/src/hooks/use-element-size.ts
index c05fea9..1928fae 100644
--- a/packages/@headlessui-react/src/hooks/use-element-size.ts
+++ b/packages/@headlessui-react/src/hooks/use-element-size.ts
@@ -1,4 +1,4 @@
-import { useState } from 'react'
+import { useMemo, useReducer } from 'react'
 import { useIsoMorphicEffect } from './use-iso-morphic-effect'
 
 function computeSize(element: HTMLElement | null) {
@@ -12,15 +12,19 @@ export function useElementSize(
   unit = false
 ) {
   let element = ref === null ? null : 'current' in ref ? ref.current : ref
-  let [size, setSize] = useState(() => computeSize(element))
+  let [identity, forceRerender] = useReducer(() => ({}), {})
+
+  // When the element changes during a re-render, we want to make sure we
+  // compute the correct size as soon as possible. However, once the element is
+  // stable, we also want to watch for changes to the element. The `identity`
+  // state can be used to recompute the size.
+  let size = useMemo(() => computeSize(element), [element, identity])
 
   useIsoMorphicEffect(() => {
     if (!element) return
 
-    let observer = new ResizeObserver(() => {
-      setSize(computeSize(element))
-    })
-
+    // Trigger a re-render whenever the element resizes
+    let observer = new ResizeObserver(forceRerender)
     observer.observe(element)
 
     return () => {