Start cleanup phase of the Dialog component when going into the Closing state (#2264)

* introduce `opening` and `closing` states

Also represent them as bits so that we can easily combine them while we
are transitioning from one state to the other.

* update `open/closed` state checks

Instead of checking whether it is in one state or an other, we can check
if the current state contains some potential sub-state.

This allows us to still check if we are in the `Open` state, while also
`Closing` because the state will be `S.Open | S.Closing`.

* expose `flags` from the `useFlags` hook

* add the `Closing` and `Opening` states to the Open/Closed state

* create dedicated `abcEnabled` variables

* keep the `State.Closing` into account for `scroll locking` and `inert others`

* add a test for the `Closing` state impacting the `Dialog` component

* cleanup unused imports

* add `unmount` util to the Vue Test renderer

* update changelog

Author: RobinMalfait

packages/@headlessui-react/CHANGELOG.md

@@ -14,6 +14,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Re-focus `Combobox.Input` when a `Combobox.Option` is selected ([#2272](https://github.com/tailwindlabs/headlessui/pull/2272))
 - Ensure we reset the `activeOptionIndex` if the active option is unmounted ([#2274](https://github.com/tailwindlabs/headlessui/pull/2274))
 - Improve Ref type for forwarded `Switch`'s ref ([#2277](https://github.com/tailwindlabs/headlessui/pull/2277))
+- Start cleanup phase of the Dialog component when going into the Closing state ([#2264](https://github.com/tailwindlabs/headlessui/pull/2264))
 
 ## [1.7.10] - 2023-02-06
 

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

@@ -1174,7 +1174,7 @@ let Options = forwardRefWithAs(function Options<
   let usesOpenClosedState = useOpenClosed()
   let visible = (() => {
     if (usesOpenClosedState !== null) {
-      return usesOpenClosedState === State.Open
+      return (usesOpenClosedState & State.Open) === State.Open
     }
 
     return data.comboboxState === ComboboxState.Open

packages/@headlessui-react/src/components/dialog/dialog.test.tsx

@@ -1,3 +1,4 @@
+import { createPortal } from 'react-dom'
 import React, { createElement, useRef, useState, Fragment, useEffect, useCallback } from 'react'
 import { render } from '@testing-library/react'
 
@@ -24,7 +25,7 @@ import {
 import { click, mouseDrag, press, Keys, shift } from '../../test-utils/interactions'
 import { PropsOf } from '../../types'
 import { Transition } from '../transitions/transition'
-import { createPortal } from 'react-dom'
+import { OpenClosedProvider, State } from '../../internal/open-closed'
 
 jest.mock('../../hooks/use-id')
 
@@ -411,6 +412,34 @@ describe('Rendering', () => {
         expect(document.documentElement.style.overflow).toBe('')
       })
     )
+
+    it(
+      'should remove the scroll lock when the open closed state is `Closing`',
+      suppressConsoleLogs(async () => {
+        function Example({ value = State.Open }) {
+          return (
+            <OpenClosedProvider value={value}>
+              <Dialog open={true} onClose={() => {}}>
+                <input id="a" type="text" />
+                <input id="b" type="text" />
+                <input id="c" type="text" />
+              </Dialog>
+            </OpenClosedProvider>
+          )
+        }
+
+        let { rerender } = render(<Example value={State.Open} />)
+
+        // The overflow should be there
+        expect(document.documentElement.style.overflow).toBe('hidden')
+
+        // Re-render but with the `Closing` state
+        rerender(<Example value={State.Open | State.Closing} />)
+
+        // The moment the dialog is closing, the overflow should be gone
+        expect(document.documentElement.style.overflow).toBe('')
+      })
+    )
   })
 
   describe('Dialog.Overlay', () => {

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

@@ -140,10 +140,7 @@ let DialogRoot = forwardRefWithAs(function Dialog<
   let usesOpenClosedState = useOpenClosed()
   if (open === undefined && usesOpenClosedState !== null) {
     // Update the `open` prop based on the open closed state
-    open = match(usesOpenClosedState, {
-      [State.Open]: true,
-      [State.Closed]: false,
-    })
+    open = (usesOpenClosedState & State.Open) === State.Open
   }
 
   let containers = useRef<Set<MutableRefObject<HTMLElement | null>>>(new Set())
@@ -209,8 +206,21 @@ let DialogRoot = forwardRefWithAs(function Dialog<
   // in between. We only care abou whether you are the top most one or not.
   let position = !hasNestedDialogs ? 'leaf' : 'parent'
 
+  // When the `Dialog` is wrapped in a `Transition` (or another Headless UI component that exposes
+  // the OpenClosed state) then we get some information via context about its state. When the
+  // `Transition` is about to close, then the `State.Closing` state will be exposed. This allows us
+  // to enable/disable certain functionality in the `Dialog` upfront instead of waiting until the
+  // `Transition` is done transitioning.
+  let isClosing =
+    usesOpenClosedState !== null ? (usesOpenClosedState & State.Closing) === State.Closing : false
+
   // Ensure other elements can't be interacted with
-  useInertOthers(internalDialogRef, hasNestedDialogs ? enabled : false)
+  let inertOthersEnabled = (() => {
+    if (!hasNestedDialogs) return false
+    if (isClosing) return false
+    return enabled
+  })()
+  useInertOthers(internalDialogRef, inertOthersEnabled)
 
   let resolveContainers = useEvent(() => {
     // Third party roots
@@ -229,25 +239,36 @@ let DialogRoot = forwardRefWithAs(function Dialog<
   })
 
   // Close Dialog on outside click
-  useOutsideClick(() => resolveContainers(), close, enabled && !hasNestedDialogs)
+  let outsideClickEnabled = (() => {
+    if (!enabled) return false
+    if (hasNestedDialogs) return false
+    return true
+  })()
+  useOutsideClick(() => resolveContainers(), close, outsideClickEnabled)
 
   // Handle `Escape` to close
+  let escapeToCloseEnabled = (() => {
+    if (hasNestedDialogs) return false
+    if (dialogState !== DialogStates.Open) return false
+    return true
+  })()
   useEventListener(ownerDocument?.defaultView, 'keydown', (event) => {
+    if (!escapeToCloseEnabled) return
     if (event.defaultPrevented) return
     if (event.key !== Keys.Escape) return
-    if (dialogState !== DialogStates.Open) return
-    if (hasNestedDialogs) return
     event.preventDefault()
     event.stopPropagation()
     close()
   })
 
   // Scroll lock
-  useScrollLock(
-    ownerDocument,
-    dialogState === DialogStates.Open && !hasParentDialog,
-    resolveContainers
-  )
+  let scrollLockEnabled = (() => {
+    if (isClosing) return false
+    if (dialogState !== DialogStates.Open) return false
+    if (hasParentDialog) return false
+    return true
+  })()
+  useScrollLock(ownerDocument, scrollLockEnabled, resolveContainers)
 
   // Trigger close when the FocusTrap gets hidden
   useEffect(() => {

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

@@ -375,7 +375,7 @@ let Panel = forwardRefWithAs(function Panel<TTag extends ElementType = typeof DE
   let usesOpenClosedState = useOpenClosed()
   let visible = (() => {
     if (usesOpenClosedState !== null) {
-      return usesOpenClosedState === State.Open
+      return (usesOpenClosedState & State.Open) === State.Open
     }
 
     return state.disclosureState === DisclosureStates.Open

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

@@ -756,7 +756,7 @@ let Options = forwardRefWithAs(function Options<
   let usesOpenClosedState = useOpenClosed()
   let visible = (() => {
     if (usesOpenClosedState !== null) {
-      return usesOpenClosedState === State.Open
+      return (usesOpenClosedState & State.Open) === State.Open
     }
 
     return data.listboxState === ListboxStates.Open

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

@@ -419,7 +419,7 @@ let Items = forwardRefWithAs(function Items<TTag extends ElementType = typeof DE
   let usesOpenClosedState = useOpenClosed()
   let visible = (() => {
     if (usesOpenClosedState !== null) {
-      return usesOpenClosedState === State.Open
+      return (usesOpenClosedState & State.Open) === State.Open
     }
 
     return state.menuState === MenuStates.Open

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

@@ -616,7 +616,7 @@ let Overlay = forwardRefWithAs(function Overlay<
   let usesOpenClosedState = useOpenClosed()
   let visible = (() => {
     if (usesOpenClosedState !== null) {
-      return usesOpenClosedState === State.Open
+      return (usesOpenClosedState & State.Open) === State.Open
     }
 
     return popoverState === PopoverStates.Open
@@ -693,7 +693,7 @@ let Panel = forwardRefWithAs(function Panel<TTag extends ElementType = typeof DE
   let usesOpenClosedState = useOpenClosed()
   let visible = (() => {
     if (usesOpenClosedState !== null) {
-      return usesOpenClosedState === State.Open
+      return (usesOpenClosedState & State.Open) === State.Open
     }
 
     return state.popoverState === PopoverStates.Open

packages/@headlessui-react/src/components/transitions/transition.tsx

@@ -32,6 +32,7 @@ import { useEvent } from '../../hooks/use-event'
 import { useDisposables } from '../../hooks/use-disposables'
 import { classNames } from '../../utils/class-names'
 import { env } from '../../utils/env'
+import { useFlags } from '../../hooks/use-flags'
 
 type ContainerElement = MutableRefObject<HTMLElement | null>
 
@@ -359,18 +360,32 @@ let TransitionChild = forwardRefWithAs(function TransitionChild<
     return show ? 'enter' : 'leave'
   })() as TransitionDirection
 
+  let transitionStateFlags = useFlags(0)
+
   let beforeEvent = useEvent((direction: TransitionDirection) => {
     return match(direction, {
-      enter: () => events.current.beforeEnter(),
-      leave: () => events.current.beforeLeave(),
+      enter: () => {
+        transitionStateFlags.addFlag(State.Opening)
+        events.current.beforeEnter()
+      },
+      leave: () => {
+        transitionStateFlags.addFlag(State.Closing)
+        events.current.beforeLeave()
+      },
       idle: () => {},
     })
   })
 
   let afterEvent = useEvent((direction: TransitionDirection) => {
     return match(direction, {
-      enter: () => events.current.afterEnter(),
-      leave: () => events.current.afterLeave(),
+      enter: () => {
+        transitionStateFlags.removeFlag(State.Opening)
+        events.current.afterEnter()
+      },
+      leave: () => {
+        transitionStateFlags.removeFlag(State.Closing)
+        events.current.afterLeave()
+      },
       idle: () => {},
     })
   })
@@ -425,10 +440,12 @@ let TransitionChild = forwardRefWithAs(function TransitionChild<
   return (
     <NestingContext.Provider value={nesting}>
       <OpenClosedProvider
-        value={match(state, {
-          [TreeStates.Visible]: State.Open,
-          [TreeStates.Hidden]: State.Closed,
-        })}
+        value={
+          match(state, {
+            [TreeStates.Visible]: State.Open,
+            [TreeStates.Hidden]: State.Closed,
+          }) | transitionStateFlags.flags
+        }
       >
         {render({
           ourProps,
@@ -457,10 +474,7 @@ let TransitionRoot = forwardRefWithAs(function Transition<
   let usesOpenClosedState = useOpenClosed()
 
   if (show === undefined && usesOpenClosedState !== null) {
-    show = match(usesOpenClosedState, {
-      [State.Open]: true,
-      [State.Closed]: false,
-    })
+    show = (usesOpenClosedState & State.Open) === State.Open
   }
 
   if (![true, false].includes(show as unknown as boolean)) {

packages/@headlessui-react/src/hooks/use-flags.ts

@@ -8,5 +8,5 @@ export function useFlags(initialFlags = 0) {
   let removeFlag = useCallback((flag: number) => setFlags((flags) => flags & ~flag), [setFlags])
   let toggleFlag = useCallback((flag: number) => setFlags((flags) => flags ^ flag), [setFlags])
 
-  return { addFlag, hasFlag, removeFlag, toggleFlag }
+  return { flags, addFlag, hasFlag, removeFlag, toggleFlag }
 }