From 34275dae7488bc8776a43bd548d4faf4df4c70a9 Mon Sep 17 00:00:00 2001
From: Jordan Pittman <jordan@cryptica.me>
Date: Wed, 2 Aug 2023 10:35:32 -0400
Subject: [PATCH] Revert "Fix hydration of components inside `<Suspense>`
 (#2633)"

This reverts commit 35012893e959c1c7ec86b1e25a1343c890b46144.
---
 packages/@headlessui-react/CHANGELOG.md       |  1 -
 .../src/components/dialog/dialog.tsx          |  6 +-
 .../src/components/focus-trap/focus-trap.tsx  |  5 ++
 .../src/components/portal/portal.tsx          |  6 +-
 .../transitions/transition.test.tsx           | 73 +++++++++++++++++++
 .../src/components/transitions/transition.tsx | 13 ++++
 .../@headlessui-react/src/hooks/use-id.ts     | 24 +-----
 .../src/hooks/use-root-containers.tsx         |  9 +--
 .../src/hooks/use-server-handoff-complete.ts  | 23 ++++++
 9 files changed, 123 insertions(+), 37 deletions(-)
 create mode 100644 packages/@headlessui-react/src/hooks/use-server-handoff-complete.ts

diff --git a/packages/@headlessui-react/CHANGELOG.md b/packages/@headlessui-react/CHANGELOG.md
index 63b1e7262..5537504bf 100644
--- a/packages/@headlessui-react/CHANGELOG.md
+++ b/packages/@headlessui-react/CHANGELOG.md
@@ -11,7 +11,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Use correct value when resetting `<Listbox multiple>` and `<Combobox multiple>` ([#2626](https://github.com/tailwindlabs/headlessui/pull/2626))
 - Render `<MainTreeNode />` in `Popover.Group` component only ([#2634](https://github.com/tailwindlabs/headlessui/pull/2634))
-- Fix hydration of components inside `<Suspense>` ([#2633](https://github.com/tailwindlabs/headlessui/pull/2633))
 
 ## [1.7.16] - 2023-07-27
 
diff --git a/packages/@headlessui-react/src/components/dialog/dialog.tsx b/packages/@headlessui-react/src/components/dialog/dialog.tsx
index 0463b15d7..d1909e0de 100644
--- a/packages/@headlessui-react/src/components/dialog/dialog.tsx
+++ b/packages/@headlessui-react/src/components/dialog/dialog.tsx
@@ -37,6 +37,7 @@ import { Portal, useNestedPortals } from '../../components/portal/portal'
 import { ForcePortalRoot } from '../../internal/portal-force-root'
 import { ComponentDescription, Description, useDescriptions } from '../description/description'
 import { useOpenClosed, State } from '../../internal/open-closed'
+import { useServerHandoffComplete } from '../../hooks/use-server-handoff-complete'
 import { StackProvider, StackMessage } from '../../internal/stack-context'
 import { useOutsideClick } from '../../hooks/use-outside-click'
 import { useOwnerDocument } from '../../hooks/use-owner'
@@ -204,7 +205,8 @@ function DialogFn<TTag extends ElementType = typeof DEFAULT_DIALOG_TAG>(
 
   let setTitleId = useEvent((id: string | null) => dispatch({ type: ActionTypes.SetTitleId, id }))
 
-  let enabled = __demoMode ? false : dialogState === DialogStates.Open
+  let ready = useServerHandoffComplete()
+  let enabled = ready ? (__demoMode ? false : dialogState === DialogStates.Open) : false
   let hasNestedDialogs = nestedDialogCount > 1 // 1 is the current dialog
   let hasParentDialog = useContext(DialogContext) !== null
   let [portals, PortalWrapper] = useNestedPortals()
@@ -214,7 +216,7 @@ function DialogFn<TTag extends ElementType = typeof DEFAULT_DIALOG_TAG>(
     MainTreeNode,
   } = useRootContainers({
     portals,
-    defaultContainers: [() => state.panelRef.current ?? internalDialogRef.current],
+    defaultContainers: [state.panelRef.current ?? internalDialogRef.current],
   })
 
   // If there are multiple dialogs, then you can be the root, the leaf or one
diff --git a/packages/@headlessui-react/src/components/focus-trap/focus-trap.tsx b/packages/@headlessui-react/src/components/focus-trap/focus-trap.tsx
index 3b816bc0f..88a2362fd 100644
--- a/packages/@headlessui-react/src/components/focus-trap/focus-trap.tsx
+++ b/packages/@headlessui-react/src/components/focus-trap/focus-trap.tsx
@@ -10,6 +10,7 @@ import React, {
 
 import { Props } from '../../types'
 import { forwardRefWithAs, HasDisplayName, RefProp, render } from '../../utils/render'
+import { useServerHandoffComplete } from '../../hooks/use-server-handoff-complete'
 import { useSyncRefs } from '../../hooks/use-sync-refs'
 import { Features as HiddenFeatures, Hidden } from '../../internal/hidden'
 import { focusElement, focusIn, Focus, FocusResult } from '../../utils/focus-management'
@@ -81,6 +82,10 @@ function FocusTrapFn<TTag extends ElementType = typeof DEFAULT_FOCUS_TRAP_TAG>(
   let focusTrapRef = useSyncRefs(container, ref)
   let { initialFocus, containers, features = Features.All, ...theirProps } = props
 
+  if (!useServerHandoffComplete()) {
+    features = Features.None
+  }
+
   let ownerDocument = useOwnerDocument(container)
 
   useRestoreFocus({ ownerDocument }, Boolean(features & Features.RestoreFocus))
diff --git a/packages/@headlessui-react/src/components/portal/portal.tsx b/packages/@headlessui-react/src/components/portal/portal.tsx
index 3b6607184..77d854939 100644
--- a/packages/@headlessui-react/src/components/portal/portal.tsx
+++ b/packages/@headlessui-react/src/components/portal/portal.tsx
@@ -19,6 +19,7 @@ import { Props } from '../../types'
 import { forwardRefWithAs, RefProp, HasDisplayName, render } from '../../utils/render'
 import { useIsoMorphicEffect } from '../../hooks/use-iso-morphic-effect'
 import { usePortalRoot } from '../../internal/portal-force-root'
+import { useServerHandoffComplete } from '../../hooks/use-server-handoff-complete'
 import { optionalRef, useSyncRefs } from '../../hooks/use-sync-refs'
 import { useOnUnmount } from '../../hooks/use-on-unmount'
 import { useOwnerDocument } from '../../hooks/use-owner'
@@ -90,6 +91,7 @@ function PortalFn<TTag extends ElementType = typeof DEFAULT_PORTAL_TAG>(
     env.isServer ? null : ownerDocument?.createElement('div') ?? null
   )
   let parent = useContext(PortalParentContext)
+  let ready = useServerHandoffComplete()
 
   useIsoMorphicEffect(() => {
     if (!target || !element) return
@@ -121,9 +123,7 @@ function PortalFn<TTag extends ElementType = typeof DEFAULT_PORTAL_TAG>(
     }
   })
 
-  let [isFirstRender, setIsFirstRender] = useState(true)
-  useEffect(() => setIsFirstRender(false), [])
-  if (isFirstRender) return null
+  if (!ready) return null
 
   let ourProps = { ref: portalRef }
 
diff --git a/packages/@headlessui-react/src/components/transitions/transition.test.tsx b/packages/@headlessui-react/src/components/transitions/transition.test.tsx
index 61a3562c0..9ee464201 100644
--- a/packages/@headlessui-react/src/components/transitions/transition.test.tsx
+++ b/packages/@headlessui-react/src/components/transitions/transition.test.tsx
@@ -155,6 +155,27 @@ describe('Setup API', () => {
         </span>
       `)
     })
+
+    it(
+      'should yell at us when we forget to forward the ref when using a render prop',
+      suppressConsoleLogs(() => {
+        expect.assertions(1)
+
+        function Dummy(props: any) {
+          return <span {...props}>Children</span>
+        }
+
+        expect(() => {
+          render(
+            <Transition show={true} as={Fragment}>
+              {() => <Dummy />}
+            </Transition>
+          )
+        }).toThrowErrorMatchingInlineSnapshot(
+          `"Did you forget to passthrough the \`ref\` to the actual DOM node?"`
+        )
+      })
+    )
   })
 
   describe('nested', () => {
@@ -328,6 +349,58 @@ describe('Setup API', () => {
         </div>
       `)
     })
+
+    it(
+      'should yell at us when we forgot to forward the ref on one of the Transition.Child components',
+      suppressConsoleLogs(() => {
+        expect.assertions(1)
+
+        function Dummy(props: any) {
+          return <div {...props} />
+        }
+
+        expect(() => {
+          render(
+            <div className="My Page">
+              <Transition show={true}>
+                <Transition.Child as={Fragment}>{() => <Dummy>Sidebar</Dummy>}</Transition.Child>
+                <Transition.Child as={Fragment}>{() => <Dummy>Content</Dummy>}</Transition.Child>
+              </Transition>
+            </div>
+          )
+        }).toThrowErrorMatchingInlineSnapshot(
+          `"Did you forget to passthrough the \`ref\` to the actual DOM node?"`
+        )
+      })
+    )
+
+    it(
+      'should yell at us when we forgot to forward a ref on the Transition component',
+      suppressConsoleLogs(() => {
+        expect.assertions(1)
+
+        function Dummy(props: any) {
+          return <div {...props} />
+        }
+
+        expect(() => {
+          render(
+            <div className="My Page">
+              <Transition show={true} as={Fragment}>
+                {() => (
+                  <Dummy>
+                    <Transition.Child>{() => <aside>Sidebar</aside>}</Transition.Child>
+                    <Transition.Child>{() => <section>Content</section>}</Transition.Child>
+                  </Dummy>
+                )}
+              </Transition>
+            </div>
+          )
+        }).toThrowErrorMatchingInlineSnapshot(
+          `"Did you forget to passthrough the \`ref\` to the actual DOM node?"`
+        )
+      })
+    )
   })
 
   describe('transition classes', () => {
diff --git a/packages/@headlessui-react/src/components/transitions/transition.tsx b/packages/@headlessui-react/src/components/transitions/transition.tsx
index e07425ae2..416ff2350 100644
--- a/packages/@headlessui-react/src/components/transitions/transition.tsx
+++ b/packages/@headlessui-react/src/components/transitions/transition.tsx
@@ -27,6 +27,7 @@ import { match } from '../../utils/match'
 import { useIsMounted } from '../../hooks/use-is-mounted'
 import { useIsoMorphicEffect } from '../../hooks/use-iso-morphic-effect'
 import { useLatestValue } from '../../hooks/use-latest-value'
+import { useServerHandoffComplete } from '../../hooks/use-server-handoff-complete'
 import { useSyncRefs } from '../../hooks/use-sync-refs'
 import { useTransition } from '../../hooks/use-transition'
 import { useEvent } from '../../hooks/use-event'
@@ -347,10 +348,19 @@ function TransitionChildFn<TTag extends ElementType = typeof DEFAULT_TRANSITION_
     afterLeave,
   })
 
+  let ready = useServerHandoffComplete()
+
+  useEffect(() => {
+    if (ready && state === TreeStates.Visible && container.current === null) {
+      throw new Error('Did you forget to passthrough the `ref` to the actual DOM node?')
+    }
+  }, [container, state, ready])
+
   // Skipping initial transition
   let skip = initial && !appear
 
   let transitionDirection = (() => {
+    if (!ready) return 'idle'
     if (skip) return 'idle'
     if (prevShow.current === show) return 'idle'
     return show ? 'enter' : 'leave'
@@ -470,6 +480,9 @@ function TransitionRootFn<TTag extends ElementType = typeof DEFAULT_TRANSITION_C
   let internalTransitionRef = useRef<HTMLElement | null>(null)
   let transitionRef = useSyncRefs(internalTransitionRef, ref)
 
+  // The TransitionChild will also call this hook, and we have to make sure that we are ready.
+  useServerHandoffComplete()
+
   let usesOpenClosedState = useOpenClosed()
 
   if (show === undefined && usesOpenClosedState !== null) {
diff --git a/packages/@headlessui-react/src/hooks/use-id.ts b/packages/@headlessui-react/src/hooks/use-id.ts
index c4ddb9fc3..761d6c1c3 100644
--- a/packages/@headlessui-react/src/hooks/use-id.ts
+++ b/packages/@headlessui-react/src/hooks/use-id.ts
@@ -1,5 +1,6 @@
 import React from 'react'
 import { useIsoMorphicEffect } from './use-iso-morphic-effect'
+import { useServerHandoffComplete } from './use-server-handoff-complete'
 import { env } from '../utils/env'
 
 // We used a "simple" approach first which worked for SSR and rehydration on the client. However we
@@ -22,26 +23,3 @@ export let useId =
 
     return id != null ? '' + id : undefined
   }
-
-// NOTE: Do NOT use this outside of the `useId` hook
-// It is not compatible with `<Suspense>` (which is in React 18 which has its own `useId` hook)
-function useServerHandoffComplete() {
-  let [complete, setComplete] = React.useState(env.isHandoffComplete)
-
-  if (complete && env.isHandoffComplete === false) {
-    // This means we are in a test environment and we need to reset the handoff state
-    // This kinda breaks the rules of React but this is only used for testing purposes
-    // And should theoretically be fine
-    setComplete(false)
-  }
-
-  React.useEffect(() => {
-    if (complete === true) return
-    setComplete(true)
-  }, [complete])
-
-  // Transition from pending to complete (forcing a re-render when server rendering)
-  React.useEffect(() => env.handoff(), [])
-
-  return complete
-}
diff --git a/packages/@headlessui-react/src/hooks/use-root-containers.tsx b/packages/@headlessui-react/src/hooks/use-root-containers.tsx
index cb438558b..e3b8d3430 100644
--- a/packages/@headlessui-react/src/hooks/use-root-containers.tsx
+++ b/packages/@headlessui-react/src/hooks/use-root-containers.tsx
@@ -3,15 +3,12 @@ import { Hidden, Features as HiddenFeatures } from '../internal/hidden'
 import { useEvent } from './use-event'
 import { useOwnerDocument } from './use-owner'
 
-type Container = HTMLElement | null | MutableRefObject<HTMLElement | null>
-type MaybeContainerFn = Container | (() => Container)
-
 export function useRootContainers({
   defaultContainers = [],
   portals,
   mainTreeNodeRef: _mainTreeNodeRef,
 }: {
-  defaultContainers?: MaybeContainerFn[]
+  defaultContainers?: (HTMLElement | null | MutableRefObject<HTMLElement | null>)[]
   portals?: MutableRefObject<HTMLElement[]>
   mainTreeNodeRef?: MutableRefObject<HTMLElement | null>
 } = {}) {
@@ -24,10 +21,6 @@ export function useRootContainers({
 
     // Resolve default containers
     for (let container of defaultContainers) {
-      if (typeof container === 'function') {
-        container = container()
-      }
-
       if (container === null) continue
       if (container instanceof HTMLElement) {
         containers.push(container)
diff --git a/packages/@headlessui-react/src/hooks/use-server-handoff-complete.ts b/packages/@headlessui-react/src/hooks/use-server-handoff-complete.ts
new file mode 100644
index 000000000..b931ea2d6
--- /dev/null
+++ b/packages/@headlessui-react/src/hooks/use-server-handoff-complete.ts
@@ -0,0 +1,23 @@
+import { useState, useEffect } from 'react'
+import { env } from '../utils/env'
+
+export function useServerHandoffComplete() {
+  let [complete, setComplete] = useState(env.isHandoffComplete)
+
+  if (complete && env.isHandoffComplete === false) {
+    // This means we are in a test environment and we need to reset the handoff state
+    // This kinda breaks the rules of React but this is only used for testing purposes
+    // And should theoretically be fine
+    setComplete(false)
+  }
+
+  useEffect(() => {
+    if (complete === true) return
+    setComplete(true)
+  }, [complete])
+
+  // Transition from pending to complete (forcing a re-render when server rendering)
+  useEffect(() => env.handoff(), [])
+
+  return complete
+}