From 39e184e0d76bc3863d1b1d2a605045dfda601df3 Mon Sep 17 00:00:00 2001
From: Kevin Deng <sxzz@sxzz.moe>
Date: Mon, 15 Dec 2025 22:03:00 +0800
Subject: [PATCH] feat: returning new string / magic-string instance

---
 src/index.ts        | 27 +++++++++++++++++++--------
 tests/index.test.ts | 23 ++++++++++++++++++++++-
 2 files changed, 41 insertions(+), 9 deletions(-)

diff --git a/src/index.ts b/src/index.ts
index adca245..fd501ef 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -4,7 +4,6 @@ import type { BindingMagicString } from 'rolldown'
 type ObjectIntersection<A, B> = {
   [K in keyof A & keyof B]: A[K]
 }
-
 export type RolldownString = ObjectIntersection<MagicString, BindingMagicString>
 
 export function rolldownString(
@@ -16,14 +15,18 @@ export function rolldownString(
 }
 
 type Awaitable<T> = T | Promise<T>
-
+export type HandlerReturn =
+  | string
+  | MagicString
+  | BindingMagicString
+  | RolldownString
+  | void
+  | undefined
 export type Handler<Meta> = (
   s: RolldownString,
   id: string,
   meta: Meta,
-) => Awaitable<
-  MagicString | BindingMagicString | RolldownString | void | undefined
->
+) => Awaitable<HandlerReturn>
 
 export function withMagicString<Meta>(
   handler: Handler<Meta>,
@@ -35,10 +38,17 @@ export function withMagicString<Meta>(
   return (code: string, id: string, meta: Meta) => {
     const s = rolldownString(code, id, meta)
     const res = handler(s, id, meta)
+    const callback = (res: HandlerReturn) => {
+      if (typeof res === 'string') {
+        return { code: res }
+      }
+      return generateTransform(res || s, id, !!res)
+    }
+
     if (res instanceof Promise) {
-      return res.then((res) => generateTransform(res || s, id))
+      return res.then(callback)
     }
-    return generateTransform(res || s, id)
+    return callback(res)
   }
 }
 
@@ -56,12 +66,13 @@ export interface CodeTransform {
 export function generateTransform(
   s: MagicString | BindingMagicString | RolldownString | undefined,
   id: string,
+  force?: boolean,
 ): CodeTransform | undefined {
   if (s?.constructor.name === 'BindingMagicString') {
     return { code: s }
   }
 
-  if (s?.hasChanged()) {
+  if (s && (force || s.hasChanged())) {
     return {
       code: s.toString(),
       get map() {
diff --git a/tests/index.test.ts b/tests/index.test.ts
index ab4e492..f619d21 100644
--- a/tests/index.test.ts
+++ b/tests/index.test.ts
@@ -49,7 +49,28 @@ test('withMagicString', async () => {
     }),
   }
   const { snapshot } = await rolldownBuild('/entry', [entry, plugin])
-  expect(snapshot).include('43')
+  expect(snapshot).includes('43')
+})
+
+test('withMagicString with returning Promise', async () => {
+  const plugin: Plugin = {
+    name: 'test',
+    transform: withMagicString(async () => {
+      await new Promise((resolve) => setTimeout(resolve, 10))
+      return new MagicString('console.log(43)')
+    }),
+  }
+  const { snapshot } = await rolldownBuild('/entry', [entry, plugin])
+  expect(snapshot).includes('console.log(43)')
+})
+
+test('withMagicString with returning string', async () => {
+  const plugin: Plugin = {
+    name: 'test',
+    transform: withMagicString(() => Promise.resolve('console.log(43)')),
+  }
+  const { snapshot } = await rolldownBuild('/entry', [entry, plugin])
+  expect(snapshot).includes('console.log(43)')
 })
 
 test('withMagicString with new magic-string instance', async () => {
-- 
2.51.2