From 3ba0032dedeec33fb3bc4289bc916ba98de4d0bb Mon Sep 17 00:00:00 2001
From: jdecroock <decroockjovi@gmail.com>
Date: Wed, 8 Jan 2025 11:19:18 +0100
Subject: [PATCH] Support ref unmounting on imperative handles

---
 hooks/src/index.js                            |  7 ++++--
 .../test/browser/useImperativeHandle.test.js  | 24 +++++++++++++++++++
 2 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/hooks/src/index.js b/hooks/src/index.js
index d6c434e400..1f40937fbd 100644
--- a/hooks/src/index.js
+++ b/hooks/src/index.js
@@ -320,8 +320,11 @@ export function useImperativeHandle(ref, createHandle, args) {
 	useLayoutEffect(
 		() => {
 			if (typeof ref == 'function') {
-				ref(createHandle());
-				return () => ref(null);
+				const result = ref(createHandle());
+				return () => {
+					ref(null);
+					if (result && typeof result == 'function') result();
+				};
 			} else if (ref) {
 				ref.current = createHandle();
 				return () => (ref.current = null);
diff --git a/hooks/test/browser/useImperativeHandle.test.js b/hooks/test/browser/useImperativeHandle.test.js
index ac71065696..f05df6212c 100644
--- a/hooks/test/browser/useImperativeHandle.test.js
+++ b/hooks/test/browser/useImperativeHandle.test.js
@@ -35,6 +35,30 @@ describe('useImperativeHandle', () => {
 		expect(ref.current.test()).to.equal('test');
 	});
 
+	it('Calls ref unmounting function', () => {
+		let ref;
+		const unmount = sinon.spy();
+
+		function Comp() {
+			useImperativeHandle(
+				r => {
+					ref = r;
+					return unmount;
+				},
+				() => ({ test: () => 'test' }),
+				[]
+			);
+			return <p>Test</p>;
+		}
+
+		render(<Comp />, scratch);
+		expect(ref).to.have.property('test');
+		expect(ref.test()).to.equal('test');
+		render(null, scratch);
+		expect(unmount).to.be.calledOnce;
+		expect(ref).to.equal(null);
+	});
+
 	it('calls createHandle after every render by default', () => {
 		let ref,
 			createHandleSpy = sinon.spy();