From c3e8e8c215881b82e491e84575a2b464fa4cfadc Mon Sep 17 00:00:00 2001
From: Artur <artur@vaadin.com>
Date: Mon, 16 Dec 2024 21:40:45 +0200
Subject: [PATCH] fix: Replace body so fast-refresh function caching works
 (#628)

---
 .changeset/thirty-lies-destroy.md             |  8 +++++
 packages/react-transform/src/index.ts         | 31 +++++++------------
 .../react-transform/test/node/index.test.tsx  |  1 +
 3 files changed, 20 insertions(+), 20 deletions(-)
 create mode 100644 .changeset/thirty-lies-destroy.md

diff --git a/.changeset/thirty-lies-destroy.md b/.changeset/thirty-lies-destroy.md
new file mode 100644
index 00000000..863f4eed
--- /dev/null
+++ b/.changeset/thirty-lies-destroy.md
@@ -0,0 +1,8 @@
+---
+"@preact/signals-react-transform": patch
+---
+
+Avoid cloning the top-level component function when we are
+prepending `useSignals`. This fixes compatability with fast-refresh
+as it requires the function identity to correctly leverage its
+cache.
diff --git a/packages/react-transform/src/index.ts b/packages/react-transform/src/index.ts
index 5574d4c1..f4c992d7 100644
--- a/packages/react-transform/src/index.ts
+++ b/packages/react-transform/src/index.ts
@@ -342,11 +342,10 @@ function wrapInTryFinally(
 	path: NodePath<FunctionLike>,
 	state: PluginPass,
 	hookUsage: HookUsage
-): FunctionLike {
+): BabelTypes.BlockStatement {
 	const stopTrackingIdentifier = path.scope.generateUidIdentifier("effect");
 
-	const newFunction = t.cloneNode(path.node);
-	newFunction.body = t.blockStatement(
+	return t.blockStatement(
 		tryCatchTemplate({
 			STORE_IDENTIFIER: stopTrackingIdentifier,
 			HOOK_IDENTIFIER: get(state, getHookIdentifier)(),
@@ -356,29 +355,26 @@ function wrapInTryFinally(
 				: t.returnStatement(path.node.body),
 		})
 	);
-
-	return newFunction;
 }
 
 function prependUseSignals<T extends FunctionLike>(
 	t: typeof BabelTypes,
 	path: NodePath<T>,
 	state: PluginPass
-): T {
-	const newFunction = t.cloneNode(path.node);
-	newFunction.body = t.blockStatement([
+): BabelTypes.BlockStatement {
+	const body = t.blockStatement([
 		t.expressionStatement(
 			t.callExpression(get(state, getHookIdentifier)(), [])
 		),
 	]);
 	if (t.isBlockStatement(path.node.body)) {
 		// TODO: Is it okay to elide the block statement here?
-		newFunction.body.body.push(...path.node.body.body);
+		body.body.push(...path.node.body.body);
 	} else {
-		newFunction.body.body.push(t.returnStatement(path.node.body));
+		body.body.push(t.returnStatement(path.node.body));
 	}
 
-	return newFunction;
+	return body;
 }
 
 function transformFunction(
@@ -398,20 +394,15 @@ function transformFunction(
 		? MANAGED_COMPONENT
 		: UNMANAGED;
 
-	let newFunction: FunctionLike;
+	let newBody: BabelTypes.BlockStatement;
 	if (hookUsage !== UNMANAGED) {
-		newFunction = wrapInTryFinally(t, path, state, hookUsage);
+		newBody = wrapInTryFinally(t, path, state, hookUsage);
 	} else {
-		newFunction = prependUseSignals(t, path, state);
+		newBody = prependUseSignals(t, path, state);
 	}
 
-	// Using replaceWith keeps the existing leading comments already so
-	// we'll clear our cloned node's leading comments to ensure they
-	// aren't duplicated in the output.
-	newFunction.leadingComments = [];
-
 	setData(path, alreadyTransformed, true);
-	path.replaceWith(newFunction);
+	path.get("body").replaceWith(newBody);
 }
 
 function createImportLazily(
diff --git a/packages/react-transform/test/node/index.test.tsx b/packages/react-transform/test/node/index.test.tsx
index 624fc1cc..a3a92b1b 100644
--- a/packages/react-transform/test/node/index.test.tsx
+++ b/packages/react-transform/test/node/index.test.tsx
@@ -669,6 +669,7 @@ describe("React Signals Babel Transform", () => {
 				};
 			`);
 
+			scope.path.scope.crawl();
 			const signalsBinding = scope.bindings["_useSignals"];
 			expect(signalsBinding).to.exist;
 			expect(signalsBinding.kind).to.equal("module");