Merge pull request #4436 from preactjs/ref-cleanup

feat: support ref cleanup functions

src/diff/index.js

@@ -554,14 +554,26 @@ function diffElementNodes(
 
 /**
  * Invoke or update a ref, depending on whether it is a function or object ref.
- * @param {Ref<any>} ref
+ * @param {Ref<any> & { _unmount?: unknown }} ref
  * @param {any} value
  * @param {VNode} vnode
  */
 export function applyRef(ref, value, vnode) {
 	try {
-		if (typeof ref == 'function') ref(value);
-		else ref.current = value;
+		if (typeof ref == 'function') {
+			let hasRefUnmount = typeof ref._unmount == 'function';
+			if (hasRefUnmount) {
+				// @ts-ignore TS doesn't like moving narrowing checks into variables
+				ref._unmount();
+			}
+
+			if (!hasRefUnmount || value != null) {
+				// Store the cleanup function on the function
+				// instance object itself to avoid shape
+				// transitioning vnode
+				ref._unmount = ref(value);
+			}
+		} else ref.current = value;
 	} catch (e) {
 		options._catchError(e, vnode);
 	}

src/internal.d.ts

@@ -126,7 +126,10 @@ declare global {
 	// We use the `current` property to differentiate between the two kinds of Refs so
 	// internally we'll define `current` on both to make TypeScript happy
 	type RefObject<T> = { current: T | null };
-	type RefCallback<T> = { (instance: T | null): void; current: undefined };
+	type RefCallback<T> = {
+		(instance: T | null): void | (() => void);
+		current: undefined;
+	};
 	type Ref<T> = RefObject<T> | RefCallback<T>;
 
 	export interface VNode<P = {}> extends preact.VNode<P> {

test/browser/refs.test.js

@@ -715,4 +715,36 @@ describe('refs', () => {
 		render(<App />, scratch);
 		render(<App show />, scratch);
 	});
+
+	it('should call ref cleanup on unmount', () => {
+		const cleanup = sinon.spy();
+		const ref = sinon.spy(() => cleanup);
+
+		function App({ show = false }) {
+			return <div>{show && <p ref={ref}>hello</p>}</div>;
+		}
+
+		render(<App show />, scratch);
+		render(<App />, scratch);
+
+		expect(cleanup).to.be.calledOnce;
+
+		// Ref should not be called with `null` when cleanup is present
+		expect(ref).to.be.calledOnce;
+		expect(ref).not.to.be.calledWith(null);
+	});
+
+	it('should call ref cleanup when ref changes', () => {
+		const cleanup = sinon.spy();
+
+		function App({ show = false, count = 0 }) {
+			return <div>{show && <p ref={() => cleanup}>hello {count}</p>}</div>;
+		}
+
+		render(<App show />, scratch);
+		render(<App show count={1} />, scratch);
+
+		// Cleanup should be invoked whenever ref function changes
+		expect(cleanup).to.be.calledOnce;
+	});
 });