From 4e53438708ec24a83879c7b98ac66260bdde7dba Mon Sep 17 00:00:00 2001
From: Marvin Hagemeister <hello@marvinh.dev>
Date: Sun, 7 Jul 2024 00:31:55 +0200
Subject: [PATCH 1/2] feat: support ref cleanup functions

---
 src/diff/index.js         | 17 ++++++++++++++---
 src/internal.d.ts         |  5 ++++-
 test/browser/refs.test.js | 32 ++++++++++++++++++++++++++++++++
 3 files changed, 50 insertions(+), 4 deletions(-)

diff --git a/src/diff/index.js b/src/diff/index.js
index 34d380283b..9f11a182d1 100644
--- a/src/diff/index.js
+++ b/src/diff/index.js
@@ -554,14 +554,25 @@ 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) {
+				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);
 	}
diff --git a/src/internal.d.ts b/src/internal.d.ts
index ab03abc67c..cbf23b3888 100644
--- a/src/internal.d.ts
+++ b/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> {
diff --git a/test/browser/refs.test.js b/test/browser/refs.test.js
index 7e4224af57..b096955ce5 100644
--- a/test/browser/refs.test.js
+++ b/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;
+	});
 });

From 7bb245a3e8f63cfe2457777e90be561969772e91 Mon Sep 17 00:00:00 2001
From: Marvin Hagemeister <hello@marvinh.dev>
Date: Sun, 7 Jul 2024 00:35:29 +0200
Subject: [PATCH 2/2] chore: work around TS complaining

---
 src/diff/index.js | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/diff/index.js b/src/diff/index.js
index 9f11a182d1..3e4b17bd62 100644
--- a/src/diff/index.js
+++ b/src/diff/index.js
@@ -563,6 +563,7 @@ export function applyRef(ref, value, vnode) {
 		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();
 			}