feat(macros/volar): allow define styles after the ReturnStatement

Author: zhiyuanzmj

packages/macros/src/volar/define-style.ts

@@ -1,9 +1,10 @@
 import { allCodeFeatures, type Code } from 'ts-macro'
-import type { DefineStyle, TransformOptions } from '.'
+import type { DefineStyle, RootKey, TransformOptions } from '.'
 
 export function transformDefineStyle(
   { expression, isCssModules }: DefineStyle,
   index: number,
+  root: RootKey,
   options: TransformOptions,
 ): void {
   const { ts, codes, ast } = options
@@ -34,6 +35,17 @@ export function transformDefineStyle(
       ),
       '\n}>',
     )
+  } else if (root?.body) {
+    const returnNode = root.body.forEachChild((node) =>
+      ts.isReturnStatement(node) ? node : undefined,
+    )
+    if (returnNode) {
+      codes.replaceRange(expression.getStart(ast), expression.getEnd())
+      codes.replaceRange(returnNode!.pos, returnNode!.pos, ';', [
+        expression.getText(ast),
+        expression.getStart(ast),
+      ])
+    }
   }
 
   addEmbeddedCode(expression, index, options)

packages/macros/src/volar/index.ts

@@ -24,13 +24,13 @@ export type JsxMacros = {
   defineStyle?: DefineStyle[]
   defineComponent?: true
 }
-export type RootMap = Map<
+
+export type RootKey =
   | import('typescript').ArrowFunction
   | import('typescript').FunctionExpression
   | import('typescript').FunctionDeclaration
-  | undefined,
-  JsxMacros
->
+  | undefined
+export type RootMap = Map<RootKey, JsxMacros>
 
 function getMacro(
   node: import('typescript').Node | undefined,

packages/macros/src/volar/transform.ts

@@ -11,7 +11,7 @@ export function transformJsxMacros(
   let defineStyleIndex = 0
   for (const [root, macros] of rootMap) {
     macros.defineStyle?.forEach((defaultStyle) =>
-      transformDefineStyle(defaultStyle, defineStyleIndex++, options),
+      transformDefineStyle(defaultStyle, defineStyleIndex++, root, options),
     )
 
     if (!root?.body) continue