feat(ast)!: add directives field to TSModuleBlock

Boshen · Boshen · commit f7edee5eb67e · 2024-06-23T02:10:27.000+08:00
closes #3564
diff --git a/crates/oxc_ast/src/ast/ts.rs b/crates/oxc_ast/src/ast/ts.rs
@@ -1004,6 +1004,7 @@ pub enum TSModuleDeclarationBody<'a> {
 pub struct TSModuleBlock<'a> {
     #[cfg_attr(feature = "serialize", serde(flatten))]
     pub span: Span,
+    pub directives: Vec<'a, Directive<'a>>,
     pub body: Vec<'a, Statement<'a>>,
 }
 
diff --git a/crates/oxc_ast/src/ast_builder.rs b/crates/oxc_ast/src/ast_builder.rs
@@ -1761,9 +1761,10 @@ impl<'a> AstBuilder<'a> {
     pub fn ts_module_block(
         self,
         span: Span,
+        directives: Vec<'a, Directive<'a>>,
         body: Vec<'a, Statement<'a>>,
     ) -> Box<'a, TSModuleBlock<'a>> {
-        self.alloc(TSModuleBlock { span, body })
+        self.alloc(TSModuleBlock { span, directives, body })
     }
 
     #[inline]
diff --git a/crates/oxc_codegen/src/gen.rs b/crates/oxc_codegen/src/gen.rs
@@ -51,19 +51,10 @@ impl<'a, const MINIFY: bool> Gen<MINIFY> for Program<'a> {
         if let Some(hashbang) = &self.hashbang {
             hashbang.gen(p, ctx);
         }
-        print_directives_and_statements(p, &self.directives, &self.body, ctx);
+        p.print_directives_and_statements(Some(&self.directives), &self.body, ctx);
     }
 }
 
-fn print_directives_and_statements<const MINIFY: bool>(
-    p: &mut Codegen<{ MINIFY }>,
-    directives: &[Directive],
-    statements: &[Statement<'_>],
-    ctx: Context,
-) {
-    p.print_directives_and_statements(Some(directives), statements, ctx);
-}
-
 impl<'a, const MINIFY: bool> Gen<MINIFY> for Hashbang<'a> {
     fn gen(&self, p: &mut Codegen<{ MINIFY }>, _ctx: Context) {
         p.print_str(b"#!");
@@ -3300,11 +3291,9 @@ impl<'a, const MINIFY: bool> Gen<MINIFY> for TSModuleDeclarationName<'a> {
 
 impl<'a, const MINIFY: bool> Gen<MINIFY> for TSModuleBlock<'a> {
     fn gen(&self, p: &mut Codegen<{ MINIFY }>, ctx: Context) {
-        p.print_curly_braces(self.span, self.body.is_empty(), |p| {
-            for item in &self.body {
-                p.print_semicolon_if_needed();
-                item.gen(p, ctx);
-            }
+        let is_empty = self.directives.is_empty() && self.body.is_empty();
+        p.print_curly_braces(self.span, is_empty, |p| {
+            p.print_directives_and_statements(Some(&self.directives), &self.body, ctx);
         });
     }
 }
diff --git a/crates/oxc_isolated_declarations/src/declaration.rs b/crates/oxc_isolated_declarations/src/declaration.rs
@@ -138,7 +138,7 @@ impl<'a> IsolatedDeclarations<'a> {
         self.scope.enter_scope(ScopeFlags::TsModuleBlock);
         let stmts = self.transform_statements_on_demand(&block.body);
         self.scope.leave_scope();
-        self.ast.ts_module_block(SPAN, stmts)
+        self.ast.ts_module_block(SPAN, self.ast.new_vec(), stmts)
     }
 
     pub fn transform_ts_module_declaration(
diff --git a/crates/oxc_parser/src/js/statement.rs b/crates/oxc_parser/src/js/statement.rs
@@ -37,52 +37,30 @@ impl<'a> ParserImpl<'a> {
 
         let mut expecting_directives = true;
         while !self.at(Kind::Eof) {
-            match self.cur_kind() {
-                Kind::RCurly if !is_top_level => break,
-                Kind::Import if !matches!(self.peek_kind(), Kind::Dot | Kind::LParen) => {
-                    let stmt = self.parse_import_declaration()?;
-                    statements.push(stmt);
-                    expecting_directives = false;
-                }
-                Kind::Export => {
-                    let stmt = self.parse_export_declaration()?;
-                    statements.push(stmt);
-                    expecting_directives = false;
-                }
-                Kind::At => {
-                    self.eat_decorators()?;
-                    expecting_directives = false;
-                    continue;
-                }
-                _ => {
-                    let stmt = self.parse_statement_list_item(StatementContext::StatementList)?;
-
-                    // Section 11.2.1 Directive Prologue
-                    // The only way to get a correct directive is to parse the statement first and check if it is a string literal.
-                    // All other method are flawed, see test cases in [babel](https://github.com/babel/babel/blob/main/packages/babel-parser/test/fixtures/core/categorized/not-directive/input.js)
-                    if expecting_directives {
-                        if let Statement::ExpressionStatement(expr) = &stmt {
-                            if let Expression::StringLiteral(string) = &expr.expression {
-                                // span start will mismatch if they are parenthesized when `preserve_parens = false`
-                                if expr.span.start == string.span.start {
-                                    let src = &self.source_text[string.span.start as usize + 1
-                                        ..string.span.end as usize - 1];
-                                    let directive = self.ast.directive(
-                                        expr.span,
-                                        (*string).clone(),
-                                        Atom::from(src),
-                                    );
-                                    directives.push(directive);
-                                    continue;
-                                }
-                            }
+            if !is_top_level && self.at(Kind::RCurly) {
+                break;
+            }
+            let stmt = self.parse_statement_list_item(StatementContext::StatementList)?;
+            // Section 11.2.1 Directive Prologue
+            // The only way to get a correct directive is to parse the statement first and check if it is a string literal.
+            // All other method are flawed, see test cases in [babel](https://github.com/babel/babel/blob/main/packages/babel-parser/test/fixtures/core/categorized/not-directive/input.js)
+            if expecting_directives {
+                if let Statement::ExpressionStatement(expr) = &stmt {
+                    if let Expression::StringLiteral(string) = &expr.expression {
+                        // span start will mismatch if they are parenthesized when `preserve_parens = false`
+                        if expr.span.start == string.span.start {
+                            let src = &self.source_text
+                                [string.span.start as usize + 1..string.span.end as usize - 1];
+                            let directive =
+                                self.ast.directive(expr.span, (*string).clone(), Atom::from(src));
+                            directives.push(directive);
+                            continue;
                         }
-                        expecting_directives = false;
                     }
-
-                    statements.push(stmt);
                 }
-            };
+                expecting_directives = false;
+            }
+            statements.push(stmt);
         }
 
         Ok((directives, statements))
diff --git a/crates/oxc_parser/src/lib.rs b/crates/oxc_parser/src/lib.rs
@@ -492,7 +492,7 @@ mod test {
         let sources = [
             ("import x from 'foo'; 'use strict';", 2),
             ("export {x} from 'foo'; 'use strict';", 2),
-            ("@decorator 'use strict';", 1),
+            (";'use strict';", 2),
         ];
         for (source, body_length) in sources {
             let ret = Parser::new(&allocator, source, source_type).parse();
diff --git a/crates/oxc_parser/src/ts/statement.rs b/crates/oxc_parser/src/ts/statement.rs
@@ -11,7 +11,7 @@ use crate::{
     js::{FunctionKind, VariableDeclarationContext, VariableDeclarationParent},
     lexer::Kind,
     list::{NormalList, SeparatedList},
-    ParserImpl, StatementContext,
+    ParserImpl,
 };
 
 impl<'a> ParserImpl<'a> {
@@ -217,21 +217,11 @@ impl<'a> ParserImpl<'a> {
 
     fn parse_ts_module_block(&mut self) -> Result<Box<'a, TSModuleBlock<'a>>> {
         let span = self.start_span();
-
-        let mut statements = self.ast.new_vec();
-
         self.expect(Kind::LCurly)?;
-
-        while !self.eat(Kind::RCurly) && !self.at(Kind::Eof) {
-            let stmt = self.parse_ts_module_item()?;
-            statements.push(stmt);
-        }
-
-        Ok(self.ast.ts_module_block(self.end_span(span), statements))
-    }
-
-    fn parse_ts_module_item(&mut self) -> Result<Statement<'a>> {
-        self.parse_statement_list_item(StatementContext::StatementList)
+        let (directives, statements) =
+            self.parse_directives_and_statements(/* is_top_level */ false)?;
+        self.expect(Kind::RCurly)?;
+        Ok(self.ast.ts_module_block(self.end_span(span), directives, statements))
     }
 
     pub(crate) fn parse_ts_namespace_or_module_declaration_body(
diff --git a/crates/oxc_transformer/src/typescript/namespace.rs b/crates/oxc_transformer/src/typescript/namespace.rs
@@ -141,8 +141,15 @@ impl<'a> TypeScript<'a> {
         let symbol_id = ctx.generate_uid(&real_name, scope_id, SymbolFlags::FunctionScopedVariable);
         let name = self.ctx.ast.new_atom(ctx.symbols().get_name(symbol_id));
 
-        let namespace_top_level = match body {
-            TSModuleDeclarationBody::TSModuleBlock(block) => block.unbox().body,
+        let directives;
+        let namespace_top_level;
+
+        match body {
+            TSModuleDeclarationBody::TSModuleBlock(block) => {
+                let block = block.unbox();
+                directives = block.directives;
+                namespace_top_level = block.body;
+            }
             // We handle `namespace X.Y {}` as if it was
             //   namespace X {
             //     export namespace Y {}
@@ -152,9 +159,10 @@ impl<'a> TypeScript<'a> {
                 let export_named_decl =
                     self.ctx.ast.plain_export_named_declaration_declaration(SPAN, declaration);
                 let stmt = Statement::ExportNamedDeclaration(export_named_decl);
-                self.ctx.ast.new_vec_single(stmt)
+                directives = self.ctx.ast.new_vec();
+                namespace_top_level = self.ctx.ast.new_vec_single(stmt);
             }
-        };
+        }
 
         let mut new_stmts = self.ctx.ast.new_vec();
 
@@ -256,7 +264,15 @@ impl<'a> TypeScript<'a> {
             return None;
         }
 
-        Some(self.transform_namespace(name, real_name, new_stmts, parent_export, scope_id, ctx))
+        Some(self.transform_namespace(
+            name,
+            real_name,
+            new_stmts,
+            directives,
+            parent_export,
+            scope_id,
+            ctx,
+        ))
     }
 
     // `namespace Foo { }` -> `let Foo; (function (_Foo) { })(Foo || (Foo = {}));`
@@ -280,35 +296,17 @@ impl<'a> TypeScript<'a> {
 
     // `namespace Foo { }` -> `let Foo; (function (_Foo) { })(Foo || (Foo = {}));`
     //                                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-    #[allow(clippy::needless_pass_by_value)]
+    #[allow(clippy::needless_pass_by_value, clippy::too_many_arguments)]
     fn transform_namespace(
         &self,
         arg_name: Atom<'a>,
         real_name: Atom<'a>,
-        mut stmts: Vec<'a, Statement<'a>>,
+        stmts: Vec<'a, Statement<'a>>,
+        directives: Vec<'a, Directive<'a>>,
         parent_export: Option<Expression<'a>>,
         scope_id: ScopeId,
         ctx: &mut TraverseCtx,
     ) -> Statement<'a> {
-        let mut directives = self.ctx.ast.new_vec();
-
-        // Check if the namespace has a `use strict` directive
-        if stmts.first().is_some_and(|stmt| {
-            matches!(stmt, Statement::ExpressionStatement(es) if
-                matches!(&es.expression, Expression::StringLiteral(literal) if
-                literal.value == "use strict")
-            )
-        }) {
-            stmts.remove(0);
-            let directive = self.ctx.ast.new_atom("use strict");
-            let directive = Directive {
-                span: SPAN,
-                expression: StringLiteral::new(SPAN, directive.clone()),
-                directive,
-            };
-            directives.push(directive);
-        }
-
         // `(function (_N) { var x; })(N || (N = {}))`;
         //  ^^^^^^^^^^^^^^^^^^^^^^^^^^
         let callee = {
diff --git a/crates/oxc_traverse/src/ancestor.rs b/crates/oxc_traverse/src/ancestor.rs
diff --git a/crates/oxc_traverse/src/walk.rs b/crates/oxc_traverse/src/walk.rs

Original file line number	Diff line number	Diff line change
`@@ -1004,6 +1004,7 @@ pub enum TSModuleDeclarationBody<'a> {`
`1004`	`1004`	`pub struct TSModuleBlock<'a> {`
`1005`	`1005`	`#[cfg_attr(feature = "serialize", serde(flatten))]`
`1006`	`1006`	`pub span: Span,`
	`1007`	`+ pub directives: Vec<'a, Directive<'a>>,`
`1007`	`1008`	`pub body: Vec<'a, Statement<'a>>,`
`1008`	`1009`	`}`
`1009`	`1010`
Original file line number	Diff line number	Diff line change
`@@ -138,7 +138,7 @@ impl<'a> IsolatedDeclarations<'a> {`
`138`	`138`	`self.scope.enter_scope(ScopeFlags::TsModuleBlock);`
`139`	`139`	`let stmts = self.transform_statements_on_demand(&block.body);`
`140`	`140`	`self.scope.leave_scope();`
`141`		`- self.ast.ts_module_block(SPAN, stmts)`
	`141`	`+ self.ast.ts_module_block(SPAN, self.ast.new_vec(), stmts)`
`142`	`142`	`}`
`143`	`143`
`144`	`144`	`pub fn transform_ts_module_declaration(`