Fix handling of comments with decorators before export (#15032)

Co-authored-by: liuxingbaoyu <30521560+liuxingbaoyu@users.noreply.github.com>

packages/babel-generator/src/generators/modules.ts

@@ -186,6 +186,7 @@ export function ExportDefaultDeclaration(
   }
 
   this.word("export");
+  this.printInnerComments(node);
   this.space();
   this.word("default");
   this.space();

packages/babel-generator/test/fixtures/comments/decorators-after-export/input.js

@@ -0,0 +1,5 @@
+/* 1 */ export /* 2 */ @dec1 /* 3 */ @dec2
+/* 4 */ class /* 5 */ C /* 6 */ { /* 7 */ } /* 8 */
+
+/* A */ export /* B */ default /* C */ @dec1 /* D */ @dec2
+/* E */ class /* F */ { /* G */ } /* H */

packages/babel-generator/test/fixtures/comments/decorators-after-export/options.json

@@ -0,0 +1,4 @@
+{
+  "plugins": [["decorators", { "decoratorsBeforeExport": false }]],
+  "decoratorsBeforeExport": false
+}

packages/babel-generator/test/fixtures/comments/decorators-after-export/output.js

@@ -0,0 +1,9 @@
+/* 1 */export /* 2 */@dec1
+/* 3 */@dec2
+/* 4 */class /* 5 */C /* 6 */ {/* 7 */} /* 8 */
+
+/* A */
+export /* B */
+default /* C */@dec1
+/* D */@dec2
+/* E */class /* F */{/* G */} /* H */

packages/babel-generator/test/fixtures/comments/decorators-before-export/input.js

@@ -0,0 +1,5 @@
+/* 1 */ @dec1 /* 2 */ @dec2 /* 3 */
+export /* 4 */ class /* 5 */ C /* 6 */ { /* 7 */ } /* 8 */
+
+/* A */ @dec1 /* B */ @dec2 /* C */
+export /* D */ default /* E */ class /* F */ { /* G */ } /* H */

packages/babel-generator/test/fixtures/comments/decorators-before-export/options.json

@@ -0,0 +1,5 @@
+{
+  "BABEL_8_BREAKING": false,
+  "plugins": [["decorators", { "decoratorsBeforeExport": true }]],
+  "decoratorsBeforeExport": true
+}

packages/babel-generator/test/fixtures/comments/decorators-before-export/output.js

@@ -0,0 +1,9 @@
+/* 1 */@dec1
+/* 2 */@dec2
+/* 3 */export /* 4 */class /* 5 */C /* 6 */ {/* 7 */} /* 8 */
+
+/* A */
+@dec1
+/* B */@dec2
+/* C */export
+  /* D */ default /* E */class /* F */{/* G */} /* H */

packages/babel-generator/test/fixtures/comments/decorators-legacy-before-export/input.js

@@ -0,0 +1,5 @@
+/* 1 */ @dec1 /* 2 */ @dec2 /* 3 */
+export /* 4 */ class /* 5 */ C /* 6 */ { /* 7 */ } /* 8 */
+
+/* A */ @dec1 /* B */ @dec2 /* C */
+export /* D */ default /* E */ class /* F */ { /* G */ } /* H */

packages/babel-generator/test/fixtures/comments/decorators-legacy-before-export/options.json

@@ -0,0 +1,5 @@
+{
+  "BABEL_8_BREAKING": false,
+  "plugins": ["decorators-legacy"],
+  "decoratorsBeforeExport": true
+}

packages/babel-generator/test/fixtures/comments/decorators-legacy-before-export/output.js

@@ -0,0 +1,9 @@
+/* 1 */@dec1
+/* 2 */@dec2
+/* 3 */export /* 4 */class /* 5 */C /* 6 */ {/* 7 */} /* 8 */
+
+/* A */
+@dec1
+/* B */@dec2
+/* C */export
+  /* D */ default /* E */class /* F */{/* G */} /* H */

packages/babel-parser/src/parser/expression.ts

@@ -100,7 +100,6 @@ export default abstract class ExpressionParser extends LValParser {
     node: N.Function,
     allowModifiers?: boolean,
   ): void;
-  abstract takeDecorators(node: N.HasDecorators): void;
   abstract parseBlockOrModuleBlockBody(
     body: N.Statement[],
     directives: N.Directive[] | null | undefined,
@@ -1102,6 +1101,7 @@ export default abstract class ExpressionParser extends LValParser {
     refExpressionErrors?: ExpressionErrors | null,
   ): N.Expression {
     let node;
+    let decorators: N.Decorator[] | null = null;
 
     const { type } = this.state;
     switch (type) {
@@ -1198,12 +1198,13 @@ export default abstract class ExpressionParser extends LValParser {
         return this.parseFunctionOrFunctionSent();
 
       case tt.at:
-        this.parseDecorators();
+        decorators = this.parseDecorators();
       // fall through
       case tt._class:
-        node = this.startNode<N.Class>();
-        this.takeDecorators(node);
-        return this.parseClass(node, false);
+        return this.parseClass(
+          this.maybeTakeDecorators(decorators, this.startNode()),
+          false,
+        );
 
       case tt._new:
         return this.parseNewOrNewTarget();

packages/babel-parser/src/parser/statement.ts

@@ -340,16 +340,19 @@ export default abstract class StatementParser extends ExpressionParser {
     context?: string | null,
     topLevel?: boolean,
   ): N.Statement {
+    let decorators: N.Decorator[] | null = null;
+
     if (this.match(tt.at)) {
-      this.parseDecorators(true);
+      decorators = this.parseDecorators(true);
     }
-    return this.parseStatementContent(context, topLevel);
+    return this.parseStatementContent(context, topLevel, decorators);
   }
 
   parseStatementContent(
     this: Parser,
     context?: string | null,
     topLevel?: boolean | null,
+    decorators?: N.Decorator[] | null,
   ): N.Statement {
     let starttype = this.state.type;
     const node = this.startNode();
@@ -392,7 +395,13 @@ export default abstract class StatementParser extends ExpressionParser {
 
       case tt._class:
         if (context) this.unexpected();
-        return this.parseClass(node as Undone<N.ClassDeclaration>, true);
+        return this.parseClass(
+          this.maybeTakeDecorators(
+            decorators,
+            node as Undone<N.ClassDeclaration>,
+          ),
+          true,
+        );
 
       case tt._if:
         return this.parseIfStatement(node as Undone<N.IfStatement>);
@@ -462,6 +471,7 @@ export default abstract class StatementParser extends ExpressionParser {
               | N.ExportDefaultDeclaration
               | N.ExportDefaultDeclaration
             >,
+            decorators,
           );
 
           if (
@@ -521,6 +531,7 @@ export default abstract class StatementParser extends ExpressionParser {
       return this.parseExpressionStatement(
         node as Undone<N.ExpressionStatement>,
         expr,
+        decorators,
       );
     }
   }
@@ -531,44 +542,58 @@ export default abstract class StatementParser extends ExpressionParser {
     }
   }
 
-  takeDecorators(node: N.HasDecorators): void {
-    const decorators =
-      this.state.decoratorStack[this.state.decoratorStack.length - 1];
-    if (decorators.length) {
-      node.decorators = decorators;
-      this.resetStartLocationFromNode(node, decorators[0]);
-      this.state.decoratorStack[this.state.decoratorStack.length - 1] = [];
+  decoratorsEnabledBeforeExport(): boolean {
+    if (this.hasPlugin("decorators-legacy")) return true;
+    return (
+      this.hasPlugin("decorators") &&
+      !!this.getPluginOption("decorators", "decoratorsBeforeExport")
+    );
+  }
+
+  // Attach the decorators to the given class.
+  // NOTE: This method changes the .start location of the class, and thus
+  // can affect comment attachment. Calling it before or after finalizing
+  // the class node (and thus finalizing its comments) changes how comments
+  // before the `class` keyword or before the final .start location of the
+  // class are attached.
+  maybeTakeDecorators<T extends Undone<N.Class>>(
+    maybeDecorators: N.Decorator[] | null,
+    classNode: T,
+    exportNode?: Undone<N.ExportDefaultDeclaration | N.ExportNamedDeclaration>,
+  ): T {
+    if (maybeDecorators) {
+      classNode.decorators = maybeDecorators;
+      this.resetStartLocationFromNode(classNode, maybeDecorators[0]);
+      if (exportNode) this.resetStartLocationFromNode(exportNode, classNode);
     }
+    return classNode;
   }
 
   canHaveLeadingDecorator(): boolean {
     return this.match(tt._class);
   }
 
-  parseDecorators(this: Parser, allowExport?: boolean): void {
-    const currentContextDecorators =
-      this.state.decoratorStack[this.state.decoratorStack.length - 1];
-    while (this.match(tt.at)) {
-      const decorator = this.parseDecorator();
-      currentContextDecorators.push(decorator);
-    }
+  parseDecorators(this: Parser, allowExport?: boolean): N.Decorator[] {
+    const decorators = [];
+    do {
+      decorators.push(this.parseDecorator());
+    } while (this.match(tt.at));
 
     if (this.match(tt._export)) {
       if (!allowExport) {
         this.unexpected();
       }
 
-      if (
-        this.hasPlugin("decorators") &&
-        !this.getPluginOption("decorators", "decoratorsBeforeExport")
-      ) {
+      if (!this.decoratorsEnabledBeforeExport()) {
         this.raise(Errors.DecoratorExportClass, { at: this.state.startLoc });
       }
     } else if (!this.canHaveLeadingDecorator()) {
       throw this.raise(Errors.UnexpectedLeadingDecorator, {
         at: this.state.startLoc,
       });
     }
+
+    return decorators;
   }
 
   parseDecorator(this: Parser): N.Decorator {
@@ -578,10 +603,6 @@ export default abstract class StatementParser extends ExpressionParser {
     this.next();
 
     if (this.hasPlugin("decorators")) {
-      // Every time a decorator class expression is evaluated, a new empty array is pushed onto the stack
-      // So that the decorators of any nested class expressions will be dealt with separately
-      this.state.decoratorStack.push([]);
-
       const startLoc = this.state.startLoc;
       let expr: N.Expression;
 
@@ -624,8 +645,6 @@ export default abstract class StatementParser extends ExpressionParser {
 
         node.expression = this.parseMaybeDecoratorArguments(expr);
       }
-
-      this.state.decoratorStack.pop();
     } else {
       node.expression = this.parseExprSubscripts();
     }
@@ -1102,6 +1121,8 @@ export default abstract class StatementParser extends ExpressionParser {
   parseExpressionStatement(
     node: Undone<N.ExpressionStatement>,
     expr: N.Expression,
+    /* eslint-disable-next-line @typescript-eslint/no-unused-vars -- used in TypeScript parser */
+    decorators: N.Decorator[] | null,
   ) {
     node.expression = expr;
     this.semicolon();
@@ -1480,8 +1501,7 @@ export default abstract class StatementParser extends ExpressionParser {
     isStatement: /* T === ClassDeclaration */ boolean,
     optionalId?: boolean,
   ): T {
-    this.next();
-    this.takeDecorators(node);
+    this.next(); // 'class'
 
     // A class definition is always strict mode code.
     const oldStrict = this.state.strict;
@@ -2110,6 +2130,7 @@ export default abstract class StatementParser extends ExpressionParser {
       | N.ExportAllDeclaration
       | N.ExportNamedDeclaration
     >,
+    decorators: N.Decorator[] | null,
   ): N.AnyExport {
     const hasDefault = this.maybeParseExportDefaultSpecifier(
       // @ts-expect-error todo(flow->ts)
@@ -2134,6 +2155,9 @@ export default abstract class StatementParser extends ExpressionParser {
 
     if (hasStar && !hasNamespace) {
       if (hasDefault) this.unexpected();
+      if (decorators) {
+        throw this.raise(Errors.UnsupportedDecoratorExport, { at: node });
+      }
       this.parseExportFrom(node as Undone<N.ExportNamedDeclaration>, true);
 
       return this.finishNode(node, "ExportAllDeclaration");
@@ -2154,6 +2178,9 @@ export default abstract class StatementParser extends ExpressionParser {
     let hasDeclaration;
     if (isFromRequired || hasSpecifiers) {
       hasDeclaration = false;
+      if (decorators) {
+        throw this.raise(Errors.UnsupportedDecoratorExport, { at: node });
+      }
       this.parseExportFrom(
         node as Undone<N.ExportNamedDeclaration>,
         isFromRequired,
@@ -2165,22 +2192,31 @@ export default abstract class StatementParser extends ExpressionParser {
     }
 
     if (isFromRequired || hasSpecifiers || hasDeclaration) {
-      this.checkExport(
-        node as Undone<N.ExportNamedDeclaration>,
-        true,
-        false,
-        !!(node as Undone<N.ExportNamedDeclaration>).source,
-      );
-      return this.finishNode(node, "ExportNamedDeclaration");
+      const node2 = node as Undone<N.ExportNamedDeclaration>;
+      this.checkExport(node2, true, false, !!node2.source);
+      if (node2.declaration?.type === "ClassDeclaration") {
+        this.maybeTakeDecorators(decorators, node2.declaration, node2);
+      } else if (decorators) {
+        throw this.raise(Errors.UnsupportedDecoratorExport, { at: node });
+      }
+      return this.finishNode(node2, "ExportNamedDeclaration");
     }
 
     if (this.eat(tt._default)) {
+      const node2 = node as Undone<N.ExportDefaultDeclaration>;
       // export default ...
-      (node as Undone<N.ExportDefaultDeclaration>).declaration =
-        this.parseExportDefaultExpression();
-      this.checkExport(node as Undone<N.ExportDefaultDeclaration>, true, true);
+      const decl = this.parseExportDefaultExpression();
+      node2.declaration = decl;
 
-      return this.finishNode(node, "ExportDefaultDeclaration");
+      if (decl.type === "ClassDeclaration") {
+        this.maybeTakeDecorators(decorators, decl as N.ClassDeclaration, node2);
+      } else if (decorators) {
+        throw this.raise(Errors.UnsupportedDecoratorExport, { at: node });
+      }
+
+      this.checkExport(node2, true, true);
+
+      return this.finishNode(node2, "ExportDefaultDeclaration");
     }
 
     throw this.unexpected(null, tt.braceL);
@@ -2291,8 +2327,14 @@ export default abstract class StatementParser extends ExpressionParser {
       ) {
         this.raise(Errors.DecoratorBeforeExport, { at: this.state.startLoc });
       }
-      this.parseDecorators(false);
-      return this.parseClass(expr as Undone<N.ClassExpression>, true, true);
+      return this.parseClass(
+        this.maybeTakeDecorators(
+          this.parseDecorators(false),
+          this.startNode<N.ClassDeclaration>(),
+        ),
+        true,
+        true,
+      );
     }
 
     if (this.match(tt._const) || this.match(tt._var) || this.isLet()) {
@@ -2311,6 +2353,14 @@ export default abstract class StatementParser extends ExpressionParser {
     // eslint-disable-next-line @typescript-eslint/no-unused-vars
     node: Undone<N.ExportNamedDeclaration>,
   ): N.Declaration | undefined | null {
+    if (this.match(tt._class)) {
+      const node = this.parseClass(
+        this.startNode<N.ClassDeclaration>(),
+        true,
+        false,
+      );
+      return node;
+    }
     return this.parseStatement(null) as N.Declaration;
   }
 
@@ -2474,14 +2524,6 @@ export default abstract class StatementParser extends ExpressionParser {
         }
       }
     }
-
-    const currentContextDecorators =
-      this.state.decoratorStack[this.state.decoratorStack.length - 1];
-    // If node.declaration is a class, it will take all decorators in the current context.
-    // Thus we should throw if we see non-empty decorators here.
-    if (currentContextDecorators.length) {
-      throw this.raise(Errors.UnsupportedDecoratorExport, { at: node });
-    }
   }
 
   checkDeclaration(node: N.Pattern | N.ObjectProperty): void {