fix(react): handle ref changes across rerenders before hydration

.changeset/fix-react-ref-before-hydration.md

@@ -0,0 +1,5 @@
+---
+"@lynx-js/react": patch
+---
+
+Fix ref callbacks not being cleaned up or re-applied correctly when the ref at the same element slot changes across rerenders that happen before hydration (e.g. a `useEffect` triggering `setState` during the initial background render).

packages/react/runtime/__test__/snapshot/ref.test.jsx

@@ -1979,6 +1979,179 @@ describe('ui operations', () => {
   });
 });
 
+describe('applyRef before hydration', () => {
+  it('ref is changed across rerenders before hydration', async function() {
+    const oldCb = vi.fn();
+    const newCb = vi.fn();
+
+    function App({ cb }) {
+      return <view ref={cb} />;
+    }
+
+    globalEnvManager.switchToBackground();
+    render(<App cb={oldCb} />, __root);
+    render(<App cb={newCb} />, __root);
+
+    expect(oldCb).toHaveBeenCalledTimes(2);
+    expect(oldCb.mock.calls[0][0]).toBeInstanceOf(RefProxy);
+    expect(oldCb.mock.calls[1][0]).toBeNull();
+
+    expect(newCb).toHaveBeenCalledTimes(1);
+    expect(newCb.mock.calls[0][0]).toBeInstanceOf(RefProxy);
+  });
+
+  it('ref becomes null on rerender before hydration', async function() {
+    const cb = vi.fn();
+
+    function App({ useRef }) {
+      return <view ref={useRef ? cb : null} />;
+    }
+
+    globalEnvManager.switchToBackground();
+    render(<App useRef={true} />, __root);
+    render(<App useRef={false} />, __root);
+
+    expect(cb).toHaveBeenCalledTimes(2);
+    expect(cb.mock.calls[0][0]).toBeInstanceOf(RefProxy);
+    expect(cb.mock.calls[1][0]).toBeNull();
+  });
+
+  it('ref is added on rerender before hydration', async function() {
+    const cb = vi.fn();
+
+    function App({ useRef }) {
+      return <view ref={useRef ? cb : null} />;
+    }
+
+    globalEnvManager.switchToBackground();
+    render(<App useRef={false} />, __root);
+    render(<App useRef={true} />, __root);
+
+    expect(cb).toHaveBeenCalledTimes(1);
+    expect(cb.mock.calls[0][0]).toBeInstanceOf(RefProxy);
+  });
+
+  it('spread ref is removed on rerender before hydration', async function() {
+    const cb = vi.fn();
+
+    function App({ withRef }) {
+      return <view {...(withRef ? { ref: cb } : {})} />;
+    }
+
+    globalEnvManager.switchToBackground();
+    render(<App withRef={true} />, __root);
+    render(<App withRef={false} />, __root);
+
+    expect(cb).toHaveBeenCalledTimes(2);
+    expect(cb.mock.calls[0][0]).toBeInstanceOf(RefProxy);
+    expect(cb.mock.calls[1][0]).toBeNull();
+  });
+
+  it('spread ref is added on rerender before hydration', async function() {
+    const cb = vi.fn();
+
+    function App({ withRef }) {
+      return <view {...(withRef ? { ref: cb } : {})} />;
+    }
+
+    globalEnvManager.switchToBackground();
+    render(<App withRef={false} />, __root);
+    render(<App withRef={true} />, __root);
+
+    expect(cb).toHaveBeenCalledTimes(1);
+    expect(cb.mock.calls[0][0]).toBeInstanceOf(RefProxy);
+  });
+
+  it('object ref (createRef) is changed across rerenders before hydration', async function() {
+    const oldRef = createRef();
+    const newRef = createRef();
+
+    function App({ r }) {
+      return <view ref={r} />;
+    }
+
+    globalEnvManager.switchToBackground();
+    render(<App r={oldRef} />, __root);
+    expect(oldRef.current).toBeInstanceOf(RefProxy);
+    expect(newRef.current).toBeNull();
+
+    render(<App r={newRef} />, __root);
+    expect(oldRef.current).toBeNull();
+    expect(newRef.current).toBeInstanceOf(RefProxy);
+  });
+
+  it('object ref (createRef) becomes null on rerender before hydration', async function() {
+    const ref = createRef();
+
+    function App({ useRef }) {
+      return <view ref={useRef ? ref : null} />;
+    }
+
+    globalEnvManager.switchToBackground();
+    render(<App useRef={true} />, __root);
+    expect(ref.current).toBeInstanceOf(RefProxy);
+
+    render(<App useRef={false} />, __root);
+    expect(ref.current).toBeNull();
+  });
+
+  it('object ref (createRef) is added on rerender before hydration', async function() {
+    const ref = createRef();
+
+    function App({ useRef }) {
+      return <view ref={useRef ? ref : null} />;
+    }
+
+    globalEnvManager.switchToBackground();
+    render(<App useRef={false} />, __root);
+    expect(ref.current).toBeNull();
+
+    render(<App useRef={true} />, __root);
+    expect(ref.current).toBeInstanceOf(RefProxy);
+  });
+
+  it('same ref callback in spread form should not be re-invoked', async function() {
+    const cb = vi.fn();
+
+    function App() {
+      return <view {...{ ref: cb }} />;
+    }
+
+    globalEnvManager.switchToBackground();
+    render(<App />, __root);
+    render(<App />, __root);
+
+    expect(cb).toHaveBeenCalledTimes(1);
+    expect(cb.mock.calls[0][0]).toBeInstanceOf(RefProxy);
+  });
+
+  it('three consecutive rerenders before hydration clean up intermediate refs', async function() {
+    const cb1 = vi.fn();
+    const cb2 = vi.fn();
+    const cb3 = vi.fn();
+
+    function App({ cb }) {
+      return <view ref={cb} />;
+    }
+
+    globalEnvManager.switchToBackground();
+    render(<App cb={cb1} />, __root);
+    render(<App cb={cb2} />, __root);
+    render(<App cb={cb3} />, __root);
+
+    expect(cb1).toHaveBeenCalledTimes(2);
+    expect(cb1.mock.calls[0][0]).toBeInstanceOf(RefProxy);
+    expect(cb1.mock.calls[1][0]).toBeNull();
+
+    expect(cb2).toHaveBeenCalledTimes(2);
+    expect(cb2.mock.calls[0][0]).toBeInstanceOf(RefProxy);
+    expect(cb2.mock.calls[1][0]).toBeNull();
+
+    expect(cb3).toHaveBeenCalledTimes(1);
+    expect(cb3.mock.calls[0][0]).toBeInstanceOf(RefProxy);
+  });
+});
+
 describe('runDelayedUiOps helper', () => {
   it('should reset shouldDelayUiOps when no tasks queued', () => {
     // flush any queued tasks from previous tests

packages/react/runtime/src/snapshot/snapshot/backgroundSnapshot.ts

@@ -13,7 +13,7 @@ import type { Worklet } from '@lynx-js/react/worklet-runtime/bindings';
 import { snapshotManager } from './definition.js';
 import type { Snapshot } from './definition.js';
 import { DynamicPartType } from './dynamicPartType.js';
-import { applyRef, clearQueuedRefs, queueRefAttrUpdate } from './ref.js';
+import { applyRef, clearQueuedRefs, getRefFromValue, queueRefAttrUpdate } from './ref.js';
 import type { Ref } from './ref.js';
 import { snapshotCreatorMap } from './snapshot.js';
 import { hydrationMap } from './snapshotInstanceHydrationMap.js';
@@ -362,14 +362,11 @@ export class BackgroundSnapshotInstance {
         }
       } else {
         this.__snapshot_def.refAndSpreadIndexes?.forEach((index) => {
+          // In first render, this.__values is undefined.
+          // In next rerenders before hydration, this.__values is not undefined.
+          const oldValue: unknown = this.__values?.[index];
           const v = (value as unknown[])[index];
-          if (v && (typeof v === 'object' || typeof v === 'function')) {
-            if ('__spread' in v && 'ref' in v) {
-              queueRefAttrUpdate(null, v.ref as Ref, this.__id, index);
-            } else if ('__ref' in v) {
-              queueRefAttrUpdate(null, v as Ref, this.__id, index);
-            }
-          }
+          queueRefAttrUpdate(getRefFromValue(oldValue), getRefFromValue(v), this.__id, index);
         });
       }
       this.__values = value as unknown[];

packages/react/runtime/src/snapshot/snapshot/ref.ts

@@ -80,6 +80,19 @@ function updateRef(
   }
 }
 
+function getRefFromValue(val: unknown): Ref | null {
+  if (!val || (typeof val !== 'object' && typeof val !== 'function')) {
+    return null;
+  }
+  if ('__spread' in val && 'ref' in val) {
+    return ((val as { ref?: Ref | null }).ref) ?? null;
+  }
+  if ('__ref' in val) {
+    return val as Ref;
+  }
+  return null;
+}
+
 function transformRef(ref: unknown): Ref | null | undefined {
   if (ref === undefined || ref === null) {
     return ref;
@@ -136,4 +149,14 @@ function clearQueuedRefs(): void {
 /**
  * @internal
  */
-export { queueRefAttrUpdate, updateRef, unref, transformRef, applyRef, applyQueuedRefs, clearQueuedRefs, type Ref };
+export {
+  queueRefAttrUpdate,
+  updateRef,
+  unref,
+  transformRef,
+  applyRef,
+  applyQueuedRefs,
+  clearQueuedRefs,
+  getRefFromValue,
+  type Ref,
+};

packages/react/testing-library/src/__tests__/ref.test.jsx

@@ -1,6 +1,6 @@
-import { createRef, Component, useState } from '@lynx-js/react';
+import { createRef, Component, useState, useEffect, useRef } from '@lynx-js/react';
 import { render } from '..';
-import { expect, vi } from 'vitest';
+import { expect, vi, describe, it } from 'vitest';
 import { act } from 'preact/test-utils';
 
 describe('component ref', () => {
@@ -483,3 +483,91 @@ describe('element ref', () => {
     });
   });
 });
+
+describe('applyRef before hydration', () => {
+  it('rerender with same ref callback should not invoke ref callback', () => {
+    const refCallback = vi.fn();
+    let bump;
+
+    function App() {
+      const [, setTick] = useState(0);
+      bump = () => setTick(t => t + 1);
+
+      useEffect(() => {
+        // This will trigger a rerender before hydration
+        bump();
+      }, []);
+
+      return <view ref={refCallback} />;
+    }
+
+    render(<App />);
+    expect(refCallback).toHaveBeenCalledTimes(1);
+    expect(refCallback.mock.calls[0][0]).toMatchObject({
+      refAttr: expect.any(Array),
+    });
+  });
+
+  const forms = [
+    'normal',
+    'spread',
+  ];
+  forms.forEach((key) => {
+    forms.forEach((key2) => {
+      it(`rerender when ref is changed from ${key} to ${key2}`, () => {
+        const oldCb = vi.fn();
+        const newCb = vi.fn();
+        let bump;
+
+        function App() {
+          const [tick, setTick] = useState(0);
+          bump = () => setTick(t => t + 1);
+
+          useEffect(() => {
+            // This will trigger a rerender before hydration
+            bump();
+          }, []);
+
+          const isFirst = tick === 0;
+          const ref = isFirst ? oldCb : newCb;
+          const form = isFirst ? key : key2;
+
+          if (form === 'spread') {
+            return <view {...{ ref }} />;
+          }
+          return <view ref={ref} />;
+        }
+
+        render(<App />);
+
+        expect(oldCb).toHaveBeenCalledTimes(2);
+        expect(oldCb.mock.calls[0][0]).toMatchObject({
+          refAttr: expect.any(Array),
+        });
+        expect(oldCb.mock.calls[1][0]).toBeNull();
+
+        expect(newCb).toHaveBeenCalledTimes(1);
+        expect(newCb.mock.calls[0][0]).toMatchObject({
+          refAttr: expect.any(Array),
+        });
+      });
+    });
+  });
+
+  it('useRef + useEffect + setState host capture is stable (portal-host pattern)', () => {
+    const seenHosts = vi.fn();
+
+    function App() {
+      const hostRef = useRef(null);
+      const [host, setHost] = useState(null);
+      useEffect(() => {
+        setHost(hostRef.current);
+      }, []);
+      if (host) seenHosts(host);
+      return <view ref={hostRef} />;
+    }
+
+    render(<App />);
+    expect(seenHosts).toHaveBeenCalledTimes(1);
+  });
+});

packages/react/testing-library/src/pure.jsx

@@ -71,8 +71,9 @@ export function render(
     globalThis.lynxTestingEnv.switchToBackgroundThread();
     act(() => {
       preactRender(compBackgroundThread, __root);
-      flushDelayedLifecycleEvents();
     });
+    // `OnLifecycleEvent::rLynxFirstScreen` should happen after `useEffect`
+    flushDelayedLifecycleEvents();
   }
 
   return {