Recreate disposed machines after Activity resumes

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

@@ -1,5 +1,5 @@
-import { createContext, useContext, useMemo } from 'react'
-import { useOnUnmount } from '../../hooks/use-on-unmount'
+import { createContext, useContext } from 'react'
+import { useMachine } from '../../hooks/use-machine'
 import { ComboboxMachine } from './combobox-machine'
 
 export const ComboboxContext = createContext<ComboboxMachine<unknown> | null>(null)
@@ -18,7 +18,5 @@ export function useComboboxMachine({
   virtual = null,
   __demoMode = false,
 }: Parameters<typeof ComboboxMachine.new>[0]) {
-  let machine = useMemo(() => ComboboxMachine.new({ id, virtual, __demoMode }), [])
-  useOnUnmount(() => machine.dispose())
-  return machine
+  return useMachine(() => ComboboxMachine.new({ id, virtual, __demoMode }))
 }

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

@@ -1,5 +1,5 @@
-import { createContext, useContext, useMemo } from 'react'
-import { useOnUnmount } from '../../hooks/use-on-unmount'
+import { createContext, useContext } from 'react'
+import { useMachine } from '../../hooks/use-machine'
 import { ListboxMachine } from './listbox-machine'
 
 export const ListboxContext = createContext<ListboxMachine<unknown> | null>(null)
@@ -20,7 +20,5 @@ export function useListboxMachine({
   id: string
   __demoMode?: boolean
 }) {
-  let machine = useMemo(() => ListboxMachine.new({ id, __demoMode }), [])
-  useOnUnmount(() => machine.dispose())
-  return machine
+  return useMachine(() => ListboxMachine.new({ id, __demoMode }))
 }

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

@@ -1,4 +1,4 @@
-import { act, render, waitFor } from '@testing-library/react'
+import { act, render, renderHook, waitFor } from '@testing-library/react'
 import React, { Fragment, createElement, createRef, useEffect, useState } from 'react'
 import {
   ListboxMode,
@@ -36,6 +36,7 @@ import {
 import { suppressConsoleLogs } from '../../test-utils/suppress-console-logs'
 import { Transition } from '../transition/transition'
 import { Listbox, ListboxButton, ListboxOption, ListboxOptions } from './listbox'
+import { useListboxMachine } from './listbox-machine-glue'
 
 beforeAll(() => {
   jest.spyOn(window, 'requestAnimationFrame').mockImplementation(setImmediate as any)
@@ -85,6 +86,19 @@ describe('safeguards', () => {
       assertListbox({ state: ListboxState.InvisibleUnmounted })
     })
   )
+
+  it('should recreate a disposed Listbox machine when rendered again', () => {
+    let { result, rerender } = renderHook(() =>
+      useListboxMachine({ id: 'headlessui-listbox-test' })
+    )
+    let machine = result.current
+
+    act(() => machine.dispose())
+    rerender()
+
+    expect(result.current).not.toBe(machine)
+    expect(result.current.disposed).toBe(false)
+  })
 })
 
 describe('Rendering', () => {

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

@@ -1,5 +1,5 @@
-import { createContext, useContext, useMemo } from 'react'
-import { useOnUnmount } from '../../hooks/use-on-unmount'
+import { createContext, useContext } from 'react'
+import { useMachine } from '../../hooks/use-machine'
 import { MenuMachine } from './menu-machine'
 
 export const MenuContext = createContext<MenuMachine | null>(null)
@@ -14,7 +14,5 @@ export function useMenuMachineContext(component: string) {
 }
 
 export function useMenuMachine({ id, __demoMode = false }: { id: string; __demoMode?: boolean }) {
-  let machine = useMemo(() => MenuMachine.new({ id, __demoMode }), [])
-  useOnUnmount(() => machine.dispose())
-  return machine
+  return useMachine(() => MenuMachine.new({ id, __demoMode }))
 }

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

@@ -1,5 +1,5 @@
-import { createContext, useContext, useMemo } from 'react'
-import { useOnUnmount } from '../../hooks/use-on-unmount'
+import { createContext, useContext } from 'react'
+import { useMachine } from '../../hooks/use-machine'
 import { PopoverMachine } from './popover-machine'
 
 export const PopoverContext = createContext<PopoverMachine | null>(null)
@@ -20,7 +20,5 @@ export function usePopoverMachine({
   id: string
   __demoMode?: boolean
 }) {
-  let machine = useMemo(() => PopoverMachine.new({ id, __demoMode }), [])
-  useOnUnmount(() => machine.dispose())
-  return machine
+  return useMachine(() => PopoverMachine.new({ id, __demoMode }))
 }

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

@@ -0,0 +1,20 @@
+import { useRef } from 'react'
+import { useOnUnmount } from './use-on-unmount'
+
+interface DisposableMachine {
+  readonly disposed: boolean
+  dispose(): void
+}
+
+export function useMachine<T extends DisposableMachine>(createMachine: () => T) {
+  let machine = useRef<T | null>(null)
+
+  // Activity preserves hook state while tearing down effects and their machine subscriptions.
+  if (machine.current === null || machine.current.disposed) {
+    machine.current = createMachine()
+  }
+
+  let current = machine.current
+  useOnUnmount(() => current.dispose())
+  return current
+}

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

@@ -4,6 +4,7 @@ import { env } from './utils/env'
 
 export abstract class Machine<State, Event extends { type: number | string }> {
   #state: State = {} as State
+  #disposed = false
   #eventSubscribers = new DefaultMap<
     Event['type'],
     Set<(state: State, event: Extract<Event, { type: any }>) => void>
@@ -24,9 +25,14 @@ export abstract class Machine<State, Event extends { type: number | string }> {
   }
 
   dispose() {
+    this.#disposed = true
     this.disposables.dispose()
   }
 
+  get disposed() {
+    return this.#disposed
+  }
+
   get state(): Readonly<State> {
     return this.#state
   }