Preserve hook state when Suspense re-suspends

compat/src/suspense.js

@@ -17,6 +17,10 @@ function initSuspenseHooks() {
 
 			while ((vnode = vnode._parent)) {
 				if ((component = vnode._component) && component._childDidSuspend) {
+					// A component that suspends before ever committing has no state
+					// worth keeping; mounted ones keep their hooks while parked.
+					if (oldVNode && !oldVNode._component)
+						newVNode._component.__hooks = UNDEFINED;
 					// Don't call oldCatchError if we found a Suspense
 					return component._childDidSuspend(error, newVNode);
 				}
@@ -40,12 +44,20 @@ function initSuspenseHooks() {
 
 function detachedClone(vnode, detachedParent, parentDom) {
 	if (vnode) {
-		if (vnode._component && vnode._component.__hooks) {
-			vnode._component.__hooks._list.forEach(effect => {
-				if (typeof effect._cleanup == 'function') effect._cleanup();
+		const hooks = vnode._component && vnode._component.__hooks;
+		if (hooks) {
+			hooks._list.forEach(effect => {
+				// Only effects carry `_passive`; clearing `_args` makes them run
+				// again when the tree is revealed, memo/ref state stays intact.
+				if (effect._passive != null) {
+					if (typeof effect._cleanup == 'function') effect._cleanup();
+					effect._cleanup = effect._args = UNDEFINED;
+				}
 			});
-
-			vnode._component.__hooks = null;
+			// Drop effects queued by the aborted render; `options._render` swaps in
+			// a fresh `_pendingEffects` array before anything is pushed again, so
+			// sharing one empty array here is safe.
+			hooks._pendingEffects = vnode._component._renderCallbacks = [];
 		}
 
 		vnode = assign({ constructor: UNDEFINED }, vnode);

compat/test/browser/suspense-hydration.test.jsx

@@ -2,8 +2,10 @@ import { setupRerender } from 'preact/test-utils';
 import React, {
 	createElement,
 	hydrate,
+	render,
 	Fragment,
 	Suspense,
+	use,
 	memo,
 	useState,
 	useSyncExternalStore
@@ -874,6 +876,75 @@ describe('suspense hydration', () => {
 			});
 	});
 
+	it('should preserve component state when re-suspending after streaming-style hydration', async () => {
+		scratch.innerHTML =
+			'<!--$s:2--><div><p>Hello</p><button>Count: 0</button></div><!--/$s:2-->';
+
+		let promise = Promise.resolve('Hello');
+		let increment;
+		function App() {
+			const message = use(promise);
+			const [count, setCount] = useState(0);
+			increment = () => setCount(value => value + 1);
+			return (
+				<div>
+					<p>{message}</p>
+					<button>Count: {count}</button>
+				</div>
+			);
+		}
+
+		hydrate(
+			<Suspense fallback={<div>Fallback</div>}>
+				<App />
+			</Suspense>,
+			scratch
+		);
+		await promise;
+		rerender();
+		rerender();
+
+		increment();
+		rerender();
+		expect(scratch.querySelector('button').textContent).to.equal('Count: 1');
+
+		promise = Promise.resolve('Hello');
+		increment();
+		rerender();
+		await promise;
+		rerender();
+		rerender();
+
+		expect(scratch.querySelector('button').textContent).to.equal('Count: 2');
+	});
+
+	it('should preserve component state when re-suspending after client render', async () => {
+		let promise = Promise.resolve('Hello');
+		let increment;
+		function App() {
+			use(promise);
+			const [count, setCount] = useState(0);
+			increment = () => setCount(value => value + 1);
+			return <button>Count: {count}</button>;
+		}
+
+		render(<Suspense fallback="Fallback"><App /></Suspense>, scratch);
+		await promise;
+		rerender();
+		rerender();
+		increment();
+		rerender();
+		expect(scratch.textContent).to.equal('Count: 1');
+
+		promise = Promise.resolve('Hello');
+		increment();
+		rerender();
+		await promise;
+		rerender();
+		rerender();
+		expect(scratch.textContent).to.equal('Count: 2');
+	});
+
 	it('should correctly hydrate and rerender a memoized lazy data loader', () => {
 		const originalHtml = '<p>Count: 5</p>';
 		scratch.innerHTML = originalHtml;

compat/test/browser/suspense.test.jsx

@@ -8,6 +8,7 @@ import React, {
 	Fragment,
 	createContext,
 	useState,
+	useRef,
 	useEffect,
 	useLayoutEffect,
 	memo
@@ -202,9 +203,43 @@ describe('suspense', () => {
 		resolve().then(assert).catch(assert);
 	});
 
-	it('should reset hooks of components', () => {
+	it('should reset hooks when a component subtree suspends during mount', async () => {
+		let resolve;
+		let resolved = false;
+		let initializations = 0;
+		const promise = new Promise(r => {
+			resolve = () => {
+				resolved = true;
+				r();
+				return promise;
+			};
+		});
+
+		function App() {
+			const [value] = useState(() => ++initializations);
+			if (!resolved) throw promise;
+			return <p>{value}</p>;
+		}
+
+		render(
+			<Suspense fallback="loading">
+				<App />
+			</Suspense>,
+			scratch
+		);
+		rerender();
+		expect(scratch.textContent).to.equal('loading');
+
+		await resolve();
+		rerender();
+		expect(scratch.innerHTML).to.equal('<p>2</p>');
+		expect(initializations).to.equal(2);
+	});
+
+	it('should preserve hooks of mounted components', () => {
 		/** @type {(v) => void} */
 		let set;
+		let initialRef;
 		const LazyComp = ({ name }) => <div>Hello from {name}</div>;
 
 		/** @type {() => Promise<void>} */
@@ -222,7 +257,10 @@ describe('suspense', () => {
 
 		const Parent = ({ children }) => {
 			const [state, setState] = useState(false);
+			const ref = useRef({});
 			set = setState;
+			if (!initialRef) initialRef = ref;
+			else expect(ref).to.equal(initialRef);
 
 			return (
 				<div>
@@ -249,15 +287,21 @@ describe('suspense', () => {
 
 		return resolve().then(() => {
 			rerender();
-			expect(scratch.innerHTML).to.eql(`<div><p>hi</p></div>`);
+			expect(scratch.innerHTML).to.eql(
+				`<div><p>hi</p><div>Hello from LazyComp</div></div>`
+			);
 		});
 	});
 
-	it('should call effect cleanups', () => {
+	it('should call effect cleanups and setups when hiding and revealing', async () => {
 		/** @type {(v) => void} */
 		let set;
+		const effectSetupSpy = vi.fn();
 		const effectSpy = vi.fn();
+		const effectWithoutCleanupSpy = vi.fn();
+		const layoutEffectSetupSpy = vi.fn();
 		const layoutEffectSpy = vi.fn();
+		const layoutEffectWithoutCleanupSpy = vi.fn();
 		const LazyComp = ({ name }) => <div>Hello from {name}</div>;
 
 		/** @type {() => Promise<void>} */
@@ -277,16 +321,20 @@ describe('suspense', () => {
 			const [state, setState] = useState(false);
 			set = setState;
 			useEffect(() => {
+				effectSetupSpy();
 				return () => {
 					effectSpy();
 				};
-			}, []);
+			}, [state]);
+			useEffect(effectWithoutCleanupSpy, [state]);
 
 			useLayoutEffect(() => {
+				layoutEffectSetupSpy();
 				return () => {
 					layoutEffectSpy();
 				};
 			}, []);
+			useLayoutEffect(layoutEffectWithoutCleanupSpy, []);
 
 			return state ? (
 				<div>{children}</div>
@@ -305,19 +353,90 @@ describe('suspense', () => {
 			</Suspense>,
 			scratch
 		);
+		expect(layoutEffectSetupSpy).toHaveBeenCalledOnce();
+		expect(layoutEffectWithoutCleanupSpy).toHaveBeenCalledOnce();
 
 		set(true);
 		rerender();
 		expect(scratch.innerHTML).to.eql('<div>Suspended...</div>');
+
+		expect(effectSetupSpy).toHaveBeenCalledOnce();
+		expect(effectWithoutCleanupSpy).toHaveBeenCalledOnce();
 		expect(effectSpy).toHaveBeenCalledOnce();
 		expect(layoutEffectSpy).toHaveBeenCalledOnce();
 
-		return resolve().then(() => {
-			rerender();
-			expect(effectSpy).toHaveBeenCalledOnce();
-			expect(layoutEffectSpy).toHaveBeenCalledOnce();
-			expect(scratch.innerHTML).to.eql(`<div><p>hi</p></div>`);
-		});
+		await resolve();
+		await act(() => rerender());
+		expect(effectSpy).toHaveBeenCalledOnce();
+		expect(layoutEffectSpy).toHaveBeenCalledOnce();
+		expect(effectSetupSpy).toHaveBeenCalledTimes(2);
+		expect(effectWithoutCleanupSpy).toHaveBeenCalledTimes(2);
+		expect(layoutEffectSetupSpy).toHaveBeenCalledTimes(2);
+		expect(layoutEffectWithoutCleanupSpy).toHaveBeenCalledTimes(2);
+		expect(scratch.innerHTML).to.eql(
+			`<div><div>Hello from LazyComp</div></div>`
+		);
+	});
+
+	it('should run a layout effect with changed deps once when revealed', async () => {
+		let promise = Promise.resolve();
+		let set;
+		let first = true;
+		const setup = vi.fn();
+		const cleanup = vi.fn();
+		function App() {
+			const [n, setN] = useState(0);
+			set = setN;
+			useLayoutEffect(() => {
+				setup(n);
+				return () => cleanup(n);
+			}, [n]);
+			if (n && first) {
+				first = false;
+				throw promise;
+			}
+			return <p>{n}</p>;
+		}
+
+		render(
+			<Suspense fallback="loading">
+				<App />
+			</Suspense>,
+			scratch
+		);
+		expect(setup).toHaveBeenCalledTimes(1);
+
+		set(1);
+		rerender();
+		expect(scratch.textContent).to.equal('loading');
+		expect(cleanup).toHaveBeenCalledTimes(1);
+
+		await promise;
+		await act(() => rerender());
+		expect(scratch.textContent).to.equal('1');
+		expect(setup.mock.calls).to.deep.equal([[0], [1]]);
+		expect(cleanup).toHaveBeenCalledTimes(1);
+	});
+
+	it('should handle a promise thrown from a layout effect', () => {
+		let threw = false;
+		function App() {
+			useLayoutEffect(() => {
+				if (!threw) {
+					threw = true;
+					throw Promise.resolve();
+				}
+			});
+			return <p>x</p>;
+		}
+
+		render(
+			<Suspense fallback="loading">
+				<App />
+			</Suspense>,
+			scratch
+		);
+		rerender();
 	});
 
 	it('should support a call to setState before rendering the fallback', () => {

hooks/src/index.js

@@ -299,6 +299,7 @@ export function useLayoutEffect(callback, args) {
 	/** @type {import('./internal').EffectHookState} */
 	const state = getHookState(currentIndex++, 4);
 	if (!options._skipEffects && argsChanged(state._args, args)) {
+		state._passive = false;
 		state._value = callback;
 		state._pendingArgs = args;
 

mangle.json

@@ -32,6 +32,7 @@
       "$_hydrationMismatch": "__m",
       "$_list": "__",
       "$_pendingEffects": "__h",
+      "$_passive": "__P",
       "$_value": "__",
       "$_nextValue": "__N",
       "$_original": "__v",