Rewrite useImperativeHandle to use useLayoutEffect (-35 B) (#2003)

* Rewrite useImperativeHandle to use useLayoutEffect (-41 B)

* Add new useContext test

* Rerun imperative handler creator if ref changes (+3 B)

* Improve hooks JSDoc

* Assign the return value of push to _afterPaintQueued (-2 B)

* Consistenly assign a number to _afterPaintQueued (-1 B)

* Support no deps array in useImperativeHandle (+12 B)

* Rename `safeRaf` to `afterNextFrame`

...since its implementation does more than a real `safeRaf` implementation. In other words, `safeRaf` invokes the callback after a rAF (or timeout) and another timeout (line 241) so its not a pure `safeRaf`. So to avoid confustion, I renamed the function to be clearer and added a comment describing its purpose

* Reuse `prevRaf` variable (-4 B)

hooks/src/index.js

@@ -30,7 +30,6 @@ options.diffed = vnode => {
 
 	const hooks = c.__hooks;
 	if (hooks) {
-		hooks._handles = bindHandles(hooks._handles);
 		hooks._pendingLayoutEffects = handleEffects(hooks._pendingLayoutEffects);
 	}
 };
@@ -61,18 +60,27 @@ function getHookState(index) {
 	// * https://github.com/michael-klein/funcy.js/blob/650beaa58c43c33a74820a3c98b3c7079cf2e333/src/renderer.mjs
 	// Other implementations to look at:
 	// * https://codesandbox.io/s/mnox05qp8
-	const hooks = currentComponent.__hooks || (currentComponent.__hooks = { _list: [], _pendingEffects: [], _pendingLayoutEffects: [], _handles: [] });
+	const hooks = currentComponent.__hooks || (currentComponent.__hooks = { _list: [], _pendingEffects: [], _pendingLayoutEffects: [] });
 
 	if (index >= hooks._list.length) {
 		hooks._list.push({});
 	}
 	return hooks._list[index];
 }
 
+/**
+ * @param {import('./index').StateUpdater<any>} initialState
+ */
 export function useState(initialState) {
 	return useReducer(invokeOrReturn, initialState);
 }
 
+/**
+ * @param {import('./index').Reducer<any, any>} reducer
+ * @param {import('./index').StateUpdater<any>} initialState
+ * @param {(initialState: any) => void} [init]
+ * @returns {[ any, (state: any) => void ]}
+ */
 export function useReducer(reducer, initialState, init) {
 
 	/** @type {import('./internal').ReducerHookState} */
@@ -132,20 +140,16 @@ export function useRef(initialValue) {
 	return useMemo(() => ({ current: initialValue }), []);
 }
 
+/**
+ * @param {object} ref
+ * @param {() => object} createHandle
+ * @param {any[]} args
+ */
 export function useImperativeHandle(ref, createHandle, args) {
-	const state = getHookState(currentIndex++);
-	if (argsChanged(state._args, args)) {
-		state._args = args;
-		currentComponent.__hooks._handles.push({ ref, createHandle });
-	}
-}
-
-function bindHandles(handles) {
-	handles.some(handle => {
-		if (typeof handle.ref === 'function') handle.ref(handle.createHandle());
-		else if (handle.ref) handle.ref.current = handle.createHandle();
-	});
-	return [];
+	useLayoutEffect(() => {
+		if (typeof ref === 'function') ref(createHandle());
+		else if (ref) ref.current = createHandle();
+	}, args == null ? args : args.concat(ref));
 }
 
 /**
@@ -213,7 +217,7 @@ let afterPaint = () => {};
  */
 function flushAfterPaintEffects() {
 	afterPaintEffects.some(component => {
-		component._afterPaintQueued = false;
+		component._afterPaintQueued = 0;
 		if (component._parentDom) {
 			component.__hooks._pendingEffects = handleEffects(component.__hooks._pendingEffects);
 		}
@@ -224,9 +228,14 @@ function flushAfterPaintEffects() {
 const RAF_TIMEOUT = 100;
 
 /**
- * requestAnimationFrame with a timeout in case it doesn't fire (for example if the browser tab is not visible)
+ * Schedule a callback to be invoked after the browser has a chance to paint a new frame.
+ * Do this by combining requestAnimationFrame (rAF) + setTimeout to invoke a callback after
+ * the next browser frame.
+ *
+ * Also, schedule a timeout in parallel to the the rAF to ensure the callback is invoked
+ * even if RAF doesn't fire (for example if the browser tab is not visible)
  */
-function safeRaf(callback) {
+function afterNextFrame(callback) {
 	const done = () => {
 		clearTimeout(timeout);
 		cancelAnimationFrame(raf);
@@ -240,22 +249,30 @@ function safeRaf(callback) {
 if (typeof window !== 'undefined') {
 	let prevRaf = options.requestAnimationFrame;
 	afterPaint = (component) => {
-		if ((!component._afterPaintQueued && (component._afterPaintQueued = true) && afterPaintEffects.push(component) === 1) ||
-		    prevRaf !== options.requestAnimationFrame) {
+		if (
+			(!component._afterPaintQueued && (component._afterPaintQueued = afterPaintEffects.push(component)) === 1)
+			|| prevRaf !== options.requestAnimationFrame
+		) {
 			prevRaf = options.requestAnimationFrame;
 
 			/* istanbul ignore next */
-			(options.requestAnimationFrame || safeRaf)(flushAfterPaintEffects);
+			(prevRaf || afterNextFrame)(flushAfterPaintEffects);
 		}
 	};
 }
 
+/**
+ * @param {import('./internal').EffectHookState[]} effects
+ */
 function handleEffects(effects) {
 	effects.forEach(invokeCleanup);
 	effects.forEach(invokeEffect);
 	return [];
 }
 
+/**
+ * @param {import('./internal').EffectHookState} hook
+ */
 function invokeCleanup(hook) {
 	if (hook._cleanup) hook._cleanup();
 }
@@ -269,6 +286,10 @@ function invokeEffect(hook) {
 	if (typeof result === 'function') hook._cleanup = result;
 }
 
+/**
+ * @param {any[]} oldArgs
+ * @param {any[]} newArgs
+ */
 function argsChanged(oldArgs, newArgs) {
 	return !oldArgs || newArgs.some((arg, index) => arg !== oldArgs[index]);
 }

hooks/test/browser/useContext.test.js

@@ -165,6 +165,28 @@ describe('useContext', () => {
 		 expect(unmountspy).not.to.be.called;
 	});
 
+	it('should only subscribe a component once', () => {
+		const values = [];
+		const Context = createContext(13);
+		let provider, subSpy;
+
+		function Comp() {
+			const value = useContext(Context);
+			values.push(value);
+			return null;
+		}
+
+		render(<Comp />, scratch);
+
+		render(<Context.Provider ref={p => provider = p} value={42}><Comp /></Context.Provider>, scratch);
+		subSpy = sinon.spy(provider, 'sub');
+
+		render(<Context.Provider value={69}><Comp /></Context.Provider>, scratch);
+		expect(subSpy).to.not.have.been.called;
+
+		expect(values).to.deep.equal([13, 42, 69]);
+	});
+
 	it('should maintain context', done => {
 		const context = createContext(null);
 		const { Provider } = context;

hooks/test/browser/useImperativeHandle.test.js

@@ -1,6 +1,7 @@
 import { createElement as h, render } from 'preact';
 import { setupScratch, teardown } from '../../../test/_util/helpers';
-import { useImperativeHandle, useRef } from '../../src';
+import { useImperativeHandle, useRef, useState } from '../../src';
+import { setupRerender } from '../../../test-utils';
 
 /** @jsx h */
 
@@ -10,8 +11,12 @@ describe('useImperativeHandle', () => {
 	/** @type {HTMLDivElement} */
 	let scratch;
 
+	/** @type {() => void} */
+	let rerender;
+
 	beforeEach(() => {
 		scratch = setupScratch();
+		rerender = setupRerender();
 	});
 
 	afterEach(() => {
@@ -32,37 +37,131 @@ describe('useImperativeHandle', () => {
 		expect(ref.current.test()).to.equal('test');
 	});
 
-	it('Updates given ref with args', () => {
-		let ref;
+	it('calls createHandle after every render by default', () => {
+		let ref, createHandleSpy = sinon.spy();
+
+		function Comp() {
+			ref = useRef({});
+			useImperativeHandle(ref, createHandleSpy);
+			return <p>Test</p>;
+		}
+
+		render(<Comp />, scratch);
+		expect(createHandleSpy).to.have.been.calledOnce;
+
+		render(<Comp />, scratch);
+		expect(createHandleSpy).to.have.been.calledTwice;
+
+		render(<Comp />, scratch);
+		expect(createHandleSpy).to.have.been.calledThrice;
+	});
+
+	it('calls createHandle only on mount if an empty array is passed', () => {
+		let ref, createHandleSpy = sinon.spy();
+
+		function Comp() {
+			ref = useRef({});
+			useImperativeHandle(ref, createHandleSpy, []);
+			return <p>Test</p>;
+		}
+
+		render(<Comp />, scratch);
+		expect(createHandleSpy).to.have.been.calledOnce;
+
+		render(<Comp />, scratch);
+		expect(createHandleSpy).to.have.been.calledOnce;
+	});
+
+	it('Updates given ref when args change', () => {
+		let ref, createHandleSpy = sinon.spy();
 
 		function Comp({ a }) {
 			ref = useRef({});
-			useImperativeHandle(ref, () => ({ test: () => 'test' + a }), [a]);
+			useImperativeHandle(ref, () => {
+				createHandleSpy();
+				return { test: () => 'test' + a };
+			}, [a]);
 			return <p>Test</p>;
 		}
 
 		render(<Comp a={0} />, scratch);
+		expect(createHandleSpy).to.have.been.calledOnce;
 		expect(ref.current).to.have.property('test');
 		expect(ref.current.test()).to.equal('test0');
 
 		render(<Comp a={1} />, scratch);
+		expect(createHandleSpy).to.have.been.calledTwice;
 		expect(ref.current).to.have.property('test');
 		expect(ref.current.test()).to.equal('test1');
+
+		render(<Comp a={0} />, scratch);
+		expect(createHandleSpy).to.have.been.calledThrice;
+		expect(ref.current).to.have.property('test');
+		expect(ref.current.test()).to.equal('test0');
+	});
+
+	it('Updates given ref when passed-in ref changes', () => {
+		let ref1, ref2;
+
+		/** @type {(arg: any) => void} */
+		let setRef;
+
+		/** @type {() => void} */
+		let updateState;
+
+		const createHandleSpy = sinon.spy(() => ({
+			test: () => 'test'
+		}));
+
+		function Comp() {
+			ref1 = useRef({});
+			ref2 = useRef({});
+
+			const [ref, setRefInternal] = useState(ref1);
+			setRef = setRefInternal;
+
+			let [value, setState] = useState(0);
+			updateState = () => setState((value + 1) % 2);
+
+			useImperativeHandle(ref, createHandleSpy, []);
+			return <p>Test</p>;
+		}
+
+		render(<Comp a={0} />, scratch);
+		expect(createHandleSpy).to.have.been.calledOnce;
+
+		updateState();
+		rerender();
+		expect(createHandleSpy).to.have.been.calledOnce;
+
+		setRef(ref2);
+		rerender();
+		expect(createHandleSpy).to.have.been.calledTwice;
+
+		updateState();
+		rerender();
+		expect(createHandleSpy).to.have.been.calledTwice;
+
+		setRef(ref1);
+		rerender();
+		expect(createHandleSpy).to.have.been.calledThrice;
 	});
 
 	it('should not update ref when args have not changed', () => {
-		let ref;
+		let ref, createHandleSpy = sinon.spy(() => ({ test: () => 'test' }));
 
 		function Comp() {
 			ref = useRef({});
-			useImperativeHandle(ref, () => ({ test: () => 'test' }), [1]);
+			useImperativeHandle(ref, createHandleSpy, [1]);
 			return <p>Test</p>;
 		}
 
 		render(<Comp />, scratch);
+		expect(createHandleSpy).to.have.been.calledOnce;
 		expect(ref.current.test()).to.equal('test');
 
 		render(<Comp />, scratch);
+		expect(createHandleSpy).to.have.been.calledOnce;
 		expect(ref.current.test()).to.equal('test');
 	});
 

src/internal.d.ts

@@ -64,6 +64,7 @@ export interface Component<P = {}, S = {}> extends preact.Component<P, S> {
 
 	_dirty: boolean;
 	_force?: boolean | null;
+	_afterPaintQueued: number; // For useEffect hook
 	_renderCallbacks: Array<() => void>; // Only class components
 	_context?: any;
 	_vnode?: VNode<P> | null;