fix: handle ref cleanup in React 19 (#5)

Author: CodyJasonBennett

src/index.ts

@@ -31,6 +31,7 @@ export interface Fiber<P = any, I = any, R = any> extends VNode<P> {
   container: FiberNode<R>
   props: P & { children: ComponentChildren }
   memoizedProps?: P & { children: ComponentChildren }
+  refCleanup: null | (() => void)
   sibling: Fiber | null
   flags: number
 }
@@ -168,9 +169,14 @@ class FiberNode extends HTMLElement {
           },
           set(value) {
             ref = (self) => {
-              const publicInstance = self === null ? null : HostConfig.getPublicInstance(fiber.stateNode)
-              if (value && 'current' in value) value.current = publicInstance
-              else value?.(publicInstance)
+              const isMounted = self != null
+              const publicInstance = isMounted ? HostConfig.getPublicInstance(fiber.stateNode) : null
+              if (value && 'current' in value) {
+                value.current = publicInstance
+              } else {
+                if (isMounted) fiber.refCleanup = value?.(publicInstance)
+                else fiber.refCleanup?.()
+              }
             }
           },
         })

tests/index.test.tsx

@@ -175,7 +175,7 @@ for (const label of ['react', 'preact']) {
     textInstance: ReconcilerNode
     suspenseInstance: ReconcilerNode
     hydratableInstance: never
-    publicInstance: null
+    publicInstance: {}
     formInstance: never
     hostContext: {}
     childSet: never
@@ -198,6 +198,7 @@ for (const label of ['react', 'preact']) {
   }
 
   const NO_CONTEXT: HostConfig['hostContext'] = {}
+  const PUBLIC_INSTANCE: HostConfig['publicInstance'] = {}
 
   let currentUpdatePriority: number = NoEventPriority
 
@@ -238,7 +239,7 @@ for (const label of ['react', 'preact']) {
       parent.children.splice(parent.children.indexOf(beforeChild), 0, child),
     removeChild: (parent, child) => parent.children.splice(parent.children.indexOf(child), 1),
     removeChildFromContainer: (container) => (container.head = null),
-    getPublicInstance: () => null,
+    getPublicInstance: () => PUBLIC_INSTANCE,
     getRootHostContext: () => NO_CONTEXT,
     getChildHostContext: () => NO_CONTEXT,
     shouldSetTextContent: () => false,
@@ -335,17 +336,19 @@ for (const label of ['react', 'preact']) {
     it('should go through lifecycle', async () => {
       const lifecycle: string[] = []
 
+      const refCallback = vi.fn()
+      const refCleanup = vi.fn()
+
       function Test() {
         React.useState(() => lifecycle.push('useState'))
         const ref = React.useRef<any>(null)
         ref.current ??= lifecycle.push('render')
         React.useImperativeHandle(ref, () => void lifecycle.push('ref'))
         React.useLayoutEffect(() => void lifecycle.push('useLayoutEffect'), [])
         React.useEffect(() => void lifecycle.push('useEffect'), [])
-        return null
+        return <element ref={(self) => (refCallback(self), refCleanup)} />
       }
-      let container!: HostContainer
-      await act(async () => void (container = render(<Test />)))
+      await act(async () => void render(<Test />))
 
       expect(lifecycle).toStrictEqual([
         'useState',
@@ -355,7 +358,14 @@ for (const label of ['react', 'preact']) {
         'useLayoutEffect',
         'useEffect',
       ])
-      expect(container.head).toBe(null)
+      expect(refCallback).toHaveBeenCalledWith(PUBLIC_INSTANCE)
+      expect(refCleanup).not.toHaveBeenCalled()
+
+      refCallback.mockClear()
+
+      await act(async () => void render(null))
+      expect(refCallback).not.toHaveBeenCalled()
+      expect(refCleanup).toHaveBeenCalled()
     })
 
     it('should render JSX', async () => {