From 4d11b5656e0aeec3c5964c290193908e2afdd8ca Mon Sep 17 00:00:00 2001
From: ComputerGuy <63362464+Ocean-OS@users.noreply.github.com>
Date: Fri, 28 Mar 2025 19:43:09 -0700
Subject: [PATCH] add tests, add to server

---
 .../3-transform/server/visitors/CallExpression.js  | 14 +++++++++++++-
 .../_expected/client/index.svelte.js               | 14 ++++++++++++++
 .../_expected/server/index.svelte.js               | 14 ++++++++++++++
 .../samples/state-in-return/index.svelte.js        |  8 ++++++++
 4 files changed, 49 insertions(+), 1 deletion(-)
 create mode 100644 packages/svelte/tests/snapshot/samples/state-in-return/_expected/client/index.svelte.js
 create mode 100644 packages/svelte/tests/snapshot/samples/state-in-return/_expected/server/index.svelte.js
 create mode 100644 packages/svelte/tests/snapshot/samples/state-in-return/index.svelte.js

diff --git a/packages/svelte/src/compiler/phases/3-transform/server/visitors/CallExpression.js b/packages/svelte/src/compiler/phases/3-transform/server/visitors/CallExpression.js
index a425bc5ec4..ad573b0329 100644
--- a/packages/svelte/src/compiler/phases/3-transform/server/visitors/CallExpression.js
+++ b/packages/svelte/src/compiler/phases/3-transform/server/visitors/CallExpression.js
@@ -1,4 +1,4 @@
-/** @import { CallExpression, Expression } from 'estree' */
+/** @import { ArrowFunctionExpression, CallExpression, Expression } from 'estree' */
 /** @import { Context } from '../types.js' */
 import { is_ignored } from '../../../../state.js';
 import * as b from '../../../../utils/builders.js';
@@ -37,5 +37,17 @@ export function CallExpression(node, context) {
 		return transform_inspect_rune(node, context);
 	}
 
+	if (
+		rune === '$state' &&
+		(context.path.at(-1)?.type === 'ReturnStatement' ||
+			(context.path.at(-1)?.type === 'ArrowFunctionExpression' &&
+				/** @type {ArrowFunctionExpression} */ (context.path.at(-1)).body === node))
+	) {
+		if (node.arguments[0]) {
+			return context.visit(node.arguments[0]);
+		}
+		return b.void0;
+	}
+
 	context.next();
 }
diff --git a/packages/svelte/tests/snapshot/samples/state-in-return/_expected/client/index.svelte.js b/packages/svelte/tests/snapshot/samples/state-in-return/_expected/client/index.svelte.js
new file mode 100644
index 0000000000..55f38f174e
--- /dev/null
+++ b/packages/svelte/tests/snapshot/samples/state-in-return/_expected/client/index.svelte.js
@@ -0,0 +1,14 @@
+/* index.svelte.js generated by Svelte VERSION */
+import * as $ from 'svelte/internal/client';
+
+export default function proxy(object) {
+	return $.proxy(object);
+}
+
+export function createCounter() {
+	let count = $.state(0);
+
+	$.update(count);
+}
+
+export const proxy_in_arrow = (object) => $.proxy(object);
\ No newline at end of file
diff --git a/packages/svelte/tests/snapshot/samples/state-in-return/_expected/server/index.svelte.js b/packages/svelte/tests/snapshot/samples/state-in-return/_expected/server/index.svelte.js
new file mode 100644
index 0000000000..de52cde883
--- /dev/null
+++ b/packages/svelte/tests/snapshot/samples/state-in-return/_expected/server/index.svelte.js
@@ -0,0 +1,14 @@
+/* index.svelte.js generated by Svelte VERSION */
+import * as $ from 'svelte/internal/server';
+
+export default function proxy(object) {
+	return object;
+}
+
+export function createCounter() {
+	let count = 0;
+
+	count++;
+}
+
+export const proxy_in_arrow = (object) => object;
\ No newline at end of file
diff --git a/packages/svelte/tests/snapshot/samples/state-in-return/index.svelte.js b/packages/svelte/tests/snapshot/samples/state-in-return/index.svelte.js
new file mode 100644
index 0000000000..f30cdd0408
--- /dev/null
+++ b/packages/svelte/tests/snapshot/samples/state-in-return/index.svelte.js
@@ -0,0 +1,8 @@
+export default function proxy(object) {
+    return $state(object);
+}
+export function createCounter() {
+    let count = $state(0);
+    count++;
+}
+export const proxy_in_arrow = object => $state(object);
\ No newline at end of file