From ff41b27d2746f2e62ba1bf2bc44d3a126148ca97 Mon Sep 17 00:00:00 2001
From: Robin Malfait <malfait.robin@gmail.com>
Date: Tue, 25 Jun 2024 15:16:33 +0200
Subject: [PATCH] Fix initial `anchor="selection"` state (#3324)

* compute `selectedOptionIndex` when using `anchor="selection"`

Instead of relying on the DOM directly, we can compute the
`selectedOptionIndex` and rely on the data directly.

We will also freeze the value while closing to prevent UI changes.

* update changelog
---
 packages/@headlessui-react/CHANGELOG.md       |  1 +
 .../src/components/listbox/listbox.tsx        | 57 ++++++++-----------
 2 files changed, 25 insertions(+), 33 deletions(-)

diff --git a/packages/@headlessui-react/CHANGELOG.md b/packages/@headlessui-react/CHANGELOG.md
index 6ffd78ba9e..499cea008f 100644
--- a/packages/@headlessui-react/CHANGELOG.md
+++ b/packages/@headlessui-react/CHANGELOG.md
@@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Fixed
 
 - Fix issues spreading omitted props onto components ([#3313](https://github.com/tailwindlabs/headlessui/pull/3313))
+- Fix initial `anchor="selection"` positioning ([#3324](https://github.com/tailwindlabs/headlessui/pull/3324))
 
 ## [2.1.0] - 2024-06-21
 
diff --git a/packages/@headlessui-react/src/components/listbox/listbox.tsx b/packages/@headlessui-react/src/components/listbox/listbox.tsx
index 2d551b8a20..7353b38942 100644
--- a/packages/@headlessui-react/src/components/listbox/listbox.tsx
+++ b/packages/@headlessui-react/src/components/listbox/listbox.tsx
@@ -944,26 +944,6 @@ function OptionsFn<TTag extends ElementType = typeof DEFAULT_OPTIONS_TAG>(
     allowed: useEvent(() => [data.buttonRef.current, data.optionsRef.current]),
   })
 
-  let initialOption = useRef<number | null>(null)
-
-  useEffect(() => {
-    if (!anchor?.to?.includes('selection')) return
-
-    if (!visible) {
-      initialOption.current = null
-      return
-    }
-
-    let elements = Array.from(data.listRef.current.values())
-    // TODO: Do not rely on DOM elements here
-    initialOption.current = elements.findIndex((el) => el?.dataset.selected === '')
-    // Default to first option if nothing is selected
-    if (initialOption.current === -1) {
-      initialOption.current = elements.findIndex((el) => el?.dataset.disabled === undefined)
-      actions.goToOption(Focus.First)
-    }
-  }, [visible, data.listRef])
-
   // We keep track whether the button moved or not, we only check this when the menu state becomes
   // closed. If the button moved, then we want to cancel pending transitions to prevent that the
   // attached `MenuItems` is still transitioning while the button moved away.
@@ -980,9 +960,30 @@ function OptionsFn<TTag extends ElementType = typeof DEFAULT_OPTIONS_TAG>(
   // its transitions, or rely on the `visible` state to hide the panel whenever necessary.
   let panelEnabled = didButtonMove ? false : visible
 
+  // We should freeze when the listbox is visible but "closed". This means that
+  // a transition is currently happening and the component is still visible (for
+  // the transition) but closed from a functionality perspective.
+  let shouldFreeze = visible && data.listboxState === ListboxStates.Closed
+
+  // Frozen state, the selected value will only update visually when the user re-opens the <Listbox />
+  let frozenValue = useFrozenData(shouldFreeze, data.value)
+
+  let isSelected = useEvent((compareValue: unknown) => data.compare(frozenValue, compareValue))
+
+  let selectedOptionIndex = useMemo(() => {
+    if (anchor == null) return null
+    if (!anchor?.to?.includes('selection')) return null
+
+    // Only compute the selected option index when using `selection` in the
+    // `anchor` prop.
+    let idx = data.options.findIndex((option) => isSelected(option.dataRef.current.value))
+    // Ensure that if no data is selected, we default to the first item.
+    if (idx === -1) idx = 0
+    return idx
+  }, [anchor, data.options])
+
   let anchorOptions = (() => {
-    if (anchor == null) return undefined
-    if (data.listRef.current.size <= 0) return { ...anchor, inner: undefined }
+    if (selectedOptionIndex === null) return { ...anchor, inner: undefined }
 
     let elements = Array.from(data.listRef.current.values())
 
@@ -990,7 +991,7 @@ function OptionsFn<TTag extends ElementType = typeof DEFAULT_OPTIONS_TAG>(
       ...anchor,
       inner: {
         listRef: { current: elements },
-        index: initialOption.current!,
+        index: selectedOptionIndex,
       },
     }
   })()
@@ -1115,16 +1116,6 @@ function OptionsFn<TTag extends ElementType = typeof DEFAULT_OPTIONS_TAG>(
     ...transitionDataAttributes(transitionData),
   })
 
-  // We should freeze when the listbox is visible but "closed". This means that
-  // a transition is currently happening and the component is still visible (for
-  // the transition) but closed from a functionality perspective.
-  let shouldFreeze = visible && data.listboxState === ListboxStates.Closed
-
-  // Frozen state, the selected value will only update visually when the user re-opens the <Listbox />
-  let frozenValue = useFrozenData(shouldFreeze, data.value)
-
-  let isSelected = useEvent((compareValue: unknown) => data.compare(frozenValue, compareValue))
-
   return (
     <Portal enabled={portal ? props.static || visible : false}>
       <ListboxDataContext.Provider