From c83439dcb7a3ae076dcb4dc4540070b4a8307700 Mon Sep 17 00:00:00 2001
From: Ryan Christian <rchristian@ryanchristian.dev>
Date: Sun, 12 Jan 2025 11:35:30 -0600
Subject: [PATCH] feat: Forward refs with `lazy()`

---
 src/lazy.js       | 10 ++++++++++
 test/lazy.test.js | 19 ++++++++++++++++++-
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/src/lazy.js b/src/lazy.js
index 564cdae..09b4403 100644
--- a/src/lazy.js
+++ b/src/lazy.js
@@ -1,6 +1,15 @@
 import { h, options } from 'preact';
 import { useState, useRef } from 'preact/hooks';
 
+const oldDiff = options.__b;
+options.__b = (vnode) => {
+	if (vnode.type && vnode.type._forwarded && vnode.ref) {
+		vnode.props.ref = vnode.ref;
+		vnode.ref = null;
+	}
+	if (oldDiff) oldDiff(vnode);
+};
+
 export default function lazy(load) {
 	let p, c;
 
@@ -21,6 +30,7 @@ export default function lazy(load) {
 		return p;
 	}
 
+	LazyComponent._forwarded = true;
 	return LazyComponent;
 }
 
diff --git a/test/lazy.test.js b/test/lazy.test.js
index d16015f..c9abf89 100644
--- a/test/lazy.test.js
+++ b/test/lazy.test.js
@@ -4,7 +4,7 @@ import * as sinon from 'sinon';
 import sinonChai from 'sinon-chai';
 
 import { LocationProvider, Router } from '../src/router.js';
-import lazy from '../src/lazy.js';
+import lazy, { ErrorBoundary } from '../src/lazy.js';
 
 const expect = chai.expect;
 chai.use(sinonChai);
@@ -42,4 +42,21 @@ describe('lazy', () => {
 		await B.preload();
 		expect(loadB).to.have.been.calledOnce;
 	});
+
+	it('should forward refs', async () => {
+		const A = () => <h1>A</h1>;
+		const LazyA = lazy(() => Promise.resolve(A));
+
+		const ref = {};
+
+		render(
+			<ErrorBoundary>
+				<LazyA ref={ref} />
+			</ErrorBoundary>,
+			scratch
+		);
+		await new Promise(r => setTimeout(r, 1))
+
+		expect(ref.current.constructor).to.equal(A);
+	});
 });