From 81a7819535c4382ba7817c817722bac6d41921d8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=9E=97=E7=83=81=E5=A3=95?= <1138674510@qq.com>
Date: Thu, 18 Aug 2022 16:07:55 +0800
Subject: [PATCH] fix(compiler-sfc): rewriteDefault for class with decorators
 (#6320)

fix #6318
---
 .../__tests__/rewriteDefault.spec.ts          | 51 ++++++++++++++++++-
 packages/compiler-sfc/src/rewriteDefault.ts   |  7 ++-
 2 files changed, 56 insertions(+), 2 deletions(-)

diff --git a/packages/compiler-sfc/__tests__/rewriteDefault.spec.ts b/packages/compiler-sfc/__tests__/rewriteDefault.spec.ts
index ce198fbc236..9fb4c64bbb3 100644
--- a/packages/compiler-sfc/__tests__/rewriteDefault.spec.ts
+++ b/packages/compiler-sfc/__tests__/rewriteDefault.spec.ts
@@ -190,7 +190,56 @@ describe('compiler sfc: rewriteDefault', () => {
     ).toMatchInlineSnapshot(`
       "/*
       export default class Foo {}*/
-      const script = class Bar {}"
+      class Bar {}
+      const script = Bar"
+    `)
+  })
+
+  test('@Component\nexport default class', async () => {
+    expect(rewriteDefault(`@Component\nexport default class Foo {}`, 'script'))
+      .toMatchInlineSnapshot(`
+      "@Component
+      class Foo {}
+      const script = Foo"
+    `)
+  })
+
+  test('@Component\nexport default class w/ comments', async () => {
+    expect(
+      rewriteDefault(`// export default\n@Component\nexport default class Foo {}`, 'script')
+    ).toMatchInlineSnapshot(`
+      "// export default
+      @Component
+      class Foo {}
+      const script = Foo"
+    `)
+  })
+
+  test('@Component\nexport default class w/ comments 2', async () => {
+    expect(
+      rewriteDefault(
+        `export default {}\n` + `// @Component\n// export default class Foo {}`,
+        'script'
+      )
+    ).toMatchInlineSnapshot(`
+      "const script = {}
+      // @Component
+      // export default class Foo {}"
+    `)
+  })
+
+  test('@Component\nexport default class w/ comments 3', async () => {
+    expect(
+      rewriteDefault(
+        `/*\n@Component\nexport default class Foo {}*/\n` + `export default class Bar {}`,
+        'script'
+      )
+    ).toMatchInlineSnapshot(`
+      "/*
+      @Component
+      export default class Foo {}*/
+      class Bar {}
+      const script = Bar"
     `)
   })
 })
diff --git a/packages/compiler-sfc/src/rewriteDefault.ts b/packages/compiler-sfc/src/rewriteDefault.ts
index da7c3a042e9..b3f57a0bc4b 100644
--- a/packages/compiler-sfc/src/rewriteDefault.ts
+++ b/packages/compiler-sfc/src/rewriteDefault.ts
@@ -42,7 +42,12 @@ export function rewriteDefault(
   }).program.body
   ast.forEach(node => {
     if (node.type === 'ExportDefaultDeclaration') {
-      s.overwrite(node.start!, node.declaration.start!, `const ${as} = `)
+      if (node.declaration.type === 'ClassDeclaration') {
+        s.overwrite(node.start!, node.declaration.id.start!, `class `)
+        s.append(`\nconst ${as} = ${node.declaration.id.name}`)
+      } else {
+        s.overwrite(node.start!, node.declaration.start!, `const ${as} = `)
+      }
     }
     if (node.type === 'ExportNamedDeclaration') {
       for (const specifier of node.specifiers) {