fix(transformer): create new scopes for new blocks in TS transform (#3908)

Create scopes for new `BlockStatement`s inserted in TS transform, and update scope tree.

crates/oxc_semantic/src/builder.rs

@@ -417,27 +417,8 @@ impl<'a> Visit<'a> for SemanticBuilder<'a> {
             if flags.contains(ScopeFlags::Top) { None } else { Some(self.current_scope_id) };
 
         let mut flags = flags;
-        // Inherit strict mode for functions
-        // https://tc39.es/ecma262/#sec-strict-mode-code
         if let Some(parent_scope_id) = parent_scope_id {
-            let mut strict_mode = self.scope.root_flags().is_strict_mode();
+            flags = self.scope.get_new_scope_flags(flags, parent_scope_id);
-            let parent_scope_flags = self.scope.get_flags(parent_scope_id);
-
-            if !strict_mode
-                && (parent_scope_flags.is_function() || parent_scope_flags.is_ts_module_block())
-                && parent_scope_flags.is_strict_mode()
-            {
-                strict_mode = true;
-            }
-
-            // inherit flags for non-function scopes
-            if !flags.contains(ScopeFlags::Function) {
-                flags |= parent_scope_flags & ScopeFlags::Modifiers;
-            };
-
-            if strict_mode {
-                flags |= ScopeFlags::StrictMode;
-            }
         }
 
         self.current_scope_id = self.scope.add_scope(parent_scope_id, flags);

crates/oxc_semantic/src/scope.rs

@@ -70,6 +70,10 @@ impl ScopeTree {
         self.child_ids.get(&scope_id)
     }
 
+    pub fn get_child_ids_mut(&mut self, scope_id: ScopeId) -> Option<&mut Vec<ScopeId>> {
+        self.child_ids.get_mut(&scope_id)
+    }
+
     pub fn descendants_from_root(&self) -> impl Iterator<Item = ScopeId> + '_ {
         self.parent_ids.iter_enumerated().map(|(scope_id, _)| scope_id)
     }
@@ -95,10 +99,43 @@ impl ScopeTree {
         &mut self.flags[scope_id]
     }
 
+    pub fn get_new_scope_flags(&self, flags: ScopeFlags, parent_scope_id: ScopeId) -> ScopeFlags {
+        let mut strict_mode = self.root_flags().is_strict_mode();
+        let parent_scope_flags = self.get_flags(parent_scope_id);
+
+        // Inherit strict mode for functions
+        // https://tc39.es/ecma262/#sec-strict-mode-code
+        if !strict_mode
+            && (parent_scope_flags.is_function() || parent_scope_flags.is_ts_module_block())
+            && parent_scope_flags.is_strict_mode()
+        {
+            strict_mode = true;
+        }
+
+        // inherit flags for non-function scopes
+        let mut flags = flags;
+        if !flags.contains(ScopeFlags::Function) {
+            flags |= parent_scope_flags & ScopeFlags::Modifiers;
+        };
+
+        if strict_mode {
+            flags |= ScopeFlags::StrictMode;
+        }
+
+        flags
+    }
+
     pub fn get_parent_id(&self, scope_id: ScopeId) -> Option<ScopeId> {
         self.parent_ids[scope_id]
     }
 
+    pub fn set_parent_id(&mut self, scope_id: ScopeId, parent_id: Option<ScopeId>) {
+        self.parent_ids[scope_id] = parent_id;
+        if let Some(parent_id) = parent_id {
+            self.child_ids.entry(parent_id).or_default().push(scope_id);
+        }
+    }
+
     /// Get a variable binding by name that was declared in the top-level scope
     pub fn get_root_binding(&self, name: &str) -> Option<SymbolId> {
         self.get_binding(self.root_scope_id(), name)
@@ -143,7 +180,7 @@ impl ScopeTree {
         &mut self.bindings[scope_id]
     }
 
-    pub(crate) fn add_scope(&mut self, parent_id: Option<ScopeId>, flags: ScopeFlags) -> ScopeId {
+    pub fn add_scope(&mut self, parent_id: Option<ScopeId>, flags: ScopeFlags) -> ScopeId {
         let scope_id = self.parent_ids.push(parent_id);
         _ = self.flags.push(flags);
         _ = self.bindings.push(Bindings::default());

crates/oxc_transformer/src/lib.rs

@@ -280,24 +280,24 @@ impl<'a> Traverse<'a> for Transformer<'a> {
         self.x3_es2015.transform_declaration_on_exit(decl);
     }
 
-    fn enter_if_statement(&mut self, stmt: &mut IfStatement<'a>, _ctx: &mut TraverseCtx<'a>) {
+    fn enter_if_statement(&mut self, stmt: &mut IfStatement<'a>, ctx: &mut TraverseCtx<'a>) {
-        self.x0_typescript.transform_if_statement(stmt);
+        self.x0_typescript.transform_if_statement(stmt, ctx);
     }
 
-    fn enter_while_statement(&mut self, stmt: &mut WhileStatement<'a>, _ctx: &mut TraverseCtx<'a>) {
+    fn enter_while_statement(&mut self, stmt: &mut WhileStatement<'a>, ctx: &mut TraverseCtx<'a>) {
-        self.x0_typescript.transform_while_statement(stmt);
+        self.x0_typescript.transform_while_statement(stmt, ctx);
     }
 
     fn enter_do_while_statement(
         &mut self,
         stmt: &mut DoWhileStatement<'a>,
-        _ctx: &mut TraverseCtx<'a>,
+        ctx: &mut TraverseCtx<'a>,
     ) {
-        self.x0_typescript.transform_do_while_statement(stmt);
+        self.x0_typescript.transform_do_while_statement(stmt, ctx);
     }
 
-    fn enter_for_statement(&mut self, stmt: &mut ForStatement<'a>, _ctx: &mut TraverseCtx<'a>) {
+    fn enter_for_statement(&mut self, stmt: &mut ForStatement<'a>, ctx: &mut TraverseCtx<'a>) {
-        self.x0_typescript.transform_for_statement(stmt);
+        self.x0_typescript.transform_for_statement(stmt, ctx);
     }
 
     fn enter_ts_export_assignment(

crates/oxc_transformer/src/typescript/annotations.rs

@@ -1,11 +1,13 @@
 #![allow(clippy::unused_self)]
 
-use std::rc::Rc;
+use std::{cell::Cell, rc::Rc};
 
 use oxc_allocator::Vec as ArenaVec;
 use oxc_ast::ast::*;
 use oxc_span::{Atom, GetSpan, Span, SPAN};
-use oxc_syntax::{operator::AssignmentOperator, reference::ReferenceFlag, symbol::SymbolId};
+use oxc_syntax::{
+    operator::AssignmentOperator, reference::ReferenceFlag, scope::ScopeFlags, symbol::SymbolId,
+};
 use oxc_traverse::TraverseCtx;
 use rustc_hash::FxHashSet;
 
@@ -370,16 +372,23 @@ impl<'a> TypeScriptAnnotations<'a> {
     /// // to
     /// if (true) { super() } else { super() }
     /// ```
-    pub fn transform_if_statement(&mut self, stmt: &mut IfStatement<'a>) {
+    pub fn transform_if_statement(
+        &mut self,
+        stmt: &mut IfStatement<'a>,
+        ctx: &mut TraverseCtx<'a>,
+    ) {
         if !self.assignments.is_empty() {
-            if let Statement::ExpressionStatement(expr) = &stmt.consequent {
+            let consequent_span = match &stmt.consequent {
-                if expr.expression.is_super_call_expression() {
+                Statement::ExpressionStatement(expr)
-                    // TODO: Need to create a scope for this block
+                    if expr.expression.is_super_call_expression() =>
-                    stmt.consequent = self.ctx.ast.block_statement(self.ctx.ast.block(
+                {
-                        expr.span,
+                    Some(expr.span)
-                        self.ctx.ast.new_vec_single(self.ctx.ast.copy(&stmt.consequent)),
-                    ));
                 }
+                _ => None,
+            };
+            if let Some(span) = consequent_span {
+                let consequent = ctx.ast.move_statement(&mut stmt.consequent);
+                stmt.consequent = Self::create_block_with_statement(consequent, span, ctx);
             }
 
             let alternate_span = match &stmt.alternate {
@@ -392,52 +401,64 @@ impl<'a> TypeScriptAnnotations<'a> {
             };
             if let Some(span) = alternate_span {
                 let alternate = stmt.alternate.take().unwrap();
-                // TODO: Need to create a scope for this block
+                stmt.alternate = Some(Self::create_block_with_statement(alternate, span, ctx));
-                stmt.alternate = Some(self.ctx.ast.block_statement(
-                    self.ctx.ast.block(span, self.ctx.ast.new_vec_single(alternate)),
-                ));
             }
         }
 
-        if stmt.consequent.is_typescript_syntax() {
+        Self::replace_with_empty_block_if_ts(&mut stmt.consequent, ctx);
-            // TODO: Need to create a scope for this block
-            stmt.consequent = self.ctx.ast.block_statement(
-                self.ctx.ast.block(stmt.consequent.span(), self.ctx.ast.new_vec()),
-            );
-        }
 
         if stmt.alternate.as_ref().is_some_and(Statement::is_typescript_syntax) {
             stmt.alternate = None;
         }
     }
 
-    pub fn transform_for_statement(&mut self, stmt: &mut ForStatement<'a>) {
+    fn create_block_with_statement(
-        if stmt.body.is_typescript_syntax() {
+        stmt: Statement<'a>,
-            // TODO: Need to create a scope for this block
+        span: Span,
-            stmt.body = self
+        ctx: &mut TraverseCtx<'a>,
-                .ctx
+    ) -> Statement<'a> {
-                .ast
+        let scope_id = ctx.insert_scope_below_statement(&stmt, ScopeFlags::empty());
-                .block_statement(self.ctx.ast.block(stmt.body.span(), self.ctx.ast.new_vec()));
+        let block = BlockStatement {
-        }
+            span,
+            body: ctx.ast.new_vec_single(stmt),
+            scope_id: Cell::new(Some(scope_id)),
+        };
+        Statement::BlockStatement(ctx.ast.alloc(block))
     }
 
-    pub fn transform_while_statement(&mut self, stmt: &mut WhileStatement<'a>) {
+    pub fn transform_for_statement(
-        if stmt.body.is_typescript_syntax() {
+        &mut self,
-            // TODO: Need to create a scope for this block
+        stmt: &mut ForStatement<'a>,
-            stmt.body = self
+        ctx: &mut TraverseCtx<'a>,
-                .ctx
+    ) {
-                .ast
+        Self::replace_with_empty_block_if_ts(&mut stmt.body, ctx);
-                .block_statement(self.ctx.ast.block(stmt.body.span(), self.ctx.ast.new_vec()));
-        }
     }
 
-    pub fn transform_do_while_statement(&mut self, stmt: &mut DoWhileStatement<'a>) {
+    pub fn transform_while_statement(
-        if stmt.body.is_typescript_syntax() {
+        &mut self,
-            // TODO: Need to create a scope for this block
+        stmt: &mut WhileStatement<'a>,
-            stmt.body = self
+        ctx: &mut TraverseCtx<'a>,
-                .ctx
+    ) {
-                .ast
+        Self::replace_with_empty_block_if_ts(&mut stmt.body, ctx);
-                .block_statement(self.ctx.ast.block(stmt.body.span(), self.ctx.ast.new_vec()));
+    }
+
+    pub fn transform_do_while_statement(
+        &mut self,
+        stmt: &mut DoWhileStatement<'a>,
+        ctx: &mut TraverseCtx<'a>,
+    ) {
+        Self::replace_with_empty_block_if_ts(&mut stmt.body, ctx);
+    }
+
+    fn replace_with_empty_block_if_ts(stmt: &mut Statement<'a>, ctx: &mut TraverseCtx<'a>) {
+        if stmt.is_typescript_syntax() {
+            let scope_id = ctx.create_scope_child_of_current(ScopeFlags::empty());
+            let block = BlockStatement {
+                span: stmt.span(),
+                body: ctx.ast.new_vec(),
+                scope_id: Cell::new(Some(scope_id)),
+            };
+            *stmt = Statement::BlockStatement(ctx.ast.alloc(block));
         }
     }
 

crates/oxc_transformer/src/typescript/mod.rs

@@ -171,20 +171,36 @@ impl<'a> TypeScript<'a> {
         self.r#enum.transform_statement(stmt, ctx);
     }
 
-    pub fn transform_if_statement(&mut self, stmt: &mut IfStatement<'a>) {
+    pub fn transform_if_statement(
-        self.annotations.transform_if_statement(stmt);
+        &mut self,
+        stmt: &mut IfStatement<'a>,
+        ctx: &mut TraverseCtx<'a>,
+    ) {
+        self.annotations.transform_if_statement(stmt, ctx);
     }
 
-    pub fn transform_while_statement(&mut self, stmt: &mut WhileStatement<'a>) {
+    pub fn transform_while_statement(
-        self.annotations.transform_while_statement(stmt);
+        &mut self,
+        stmt: &mut WhileStatement<'a>,
+        ctx: &mut TraverseCtx<'a>,
+    ) {
+        self.annotations.transform_while_statement(stmt, ctx);
     }
 
-    pub fn transform_do_while_statement(&mut self, stmt: &mut DoWhileStatement<'a>) {
+    pub fn transform_do_while_statement(
-        self.annotations.transform_do_while_statement(stmt);
+        &mut self,
+        stmt: &mut DoWhileStatement<'a>,
+        ctx: &mut TraverseCtx<'a>,
+    ) {
+        self.annotations.transform_do_while_statement(stmt, ctx);
     }
 
-    pub fn transform_for_statement(&mut self, stmt: &mut ForStatement<'a>) {
+    pub fn transform_for_statement(
-        self.annotations.transform_for_statement(stmt);
+        &mut self,
+        stmt: &mut ForStatement<'a>,
+        ctx: &mut TraverseCtx<'a>,
+    ) {
+        self.annotations.transform_for_statement(stmt, ctx);
     }
 
     pub fn transform_tagged_template_expression(

crates/oxc_traverse/src/context/mod.rs

@@ -1,5 +1,8 @@
 use oxc_allocator::{Allocator, Box};
-use oxc_ast::AstBuilder;
+use oxc_ast::{
+    ast::{Expression, Statement},
+    AstBuilder,
+};
 use oxc_semantic::{ScopeTree, SymbolTable};
 use oxc_span::CompactStr;
 use oxc_syntax::{
@@ -209,6 +212,14 @@ impl<'a> TraverseCtx<'a> {
         self.scoping.current_scope_id()
     }
 
+    /// Get current scope flags.
+    ///
+    /// Shortcut for `ctx.scoping.current_scope_flags`.
+    #[inline]
+    pub fn current_scope_flags(&self) -> ScopeFlags {
+        self.scoping.current_scope_flags()
+    }
+
     /// Get scopes tree.
     ///
     /// Shortcut for `ctx.scoping.scopes`.
@@ -275,6 +286,45 @@ impl<'a> TraverseCtx<'a> {
         self.scoping.find_scope_by_flags(finder)
     }
 
+    /// Create new scope as child of current scope.
+    ///
+    /// `flags` provided are amended to inherit from parent scope's flags.
+    ///
+    /// This is a shortcut for `ctx.scoping.create_scope_child_of_current`.
+    pub fn create_scope_child_of_current(&mut self, flags: ScopeFlags) -> ScopeId {
+        self.scoping.create_scope_child_of_current(flags)
+    }
+
+    /// Insert a scope into scope tree below a statement.
+    ///
+    /// Statement must be in current scope.
+    /// New scope is created as child of current scope.
+    /// All child scopes of the statement are reassigned to be children of the new scope.
+    ///
+    /// `flags` provided are amended to inherit from parent scope's flags.
+    ///
+    /// This is a shortcut for `ctx.scoping.insert_scope_below_statement`.
+    pub fn insert_scope_below_statement(&mut self, stmt: &Statement, flags: ScopeFlags) -> ScopeId {
+        self.scoping.insert_scope_below_statement(stmt, flags)
+    }
+
+    /// Insert a scope into scope tree below an expression.
+    ///
+    /// Expression must be in current scope.
+    /// New scope is created as child of current scope.
+    /// All child scopes of the expression are reassigned to be children of the new scope.
+    ///
+    /// `flags` provided are amended to inherit from parent scope's flags.
+    ///
+    /// This is a shortcut for `ctx.scoping.insert_scope_below_expression`.
+    pub fn insert_scope_below_expression(
+        &mut self,
+        expr: &Expression,
+        flags: ScopeFlags,
+    ) -> ScopeId {
+        self.scoping.insert_scope_below_expression(expr, flags)
+    }
+
     /// Generate UID.
     ///
     /// This is a shortcut for `ctx.scoping.generate_uid`.

crates/oxc_traverse/src/context/scoping.rs

@@ -1,6 +1,11 @@
 use std::str;
 
 use compact_str::{format_compact, CompactString};
+#[allow(clippy::wildcard_imports)]
+use oxc_ast::{
+    ast::*,
+    visit::{walk, Visit},
+};
 use oxc_semantic::{AstNodeId, Reference, ScopeTree, SymbolTable};
 use oxc_span::{CompactStr, SPAN};
 use oxc_syntax::{
@@ -31,6 +36,12 @@ impl TraverseScoping {
         self.current_scope_id
     }
 
+    /// Get current scope flags
+    #[inline]
+    pub fn current_scope_flags(&self) -> ScopeFlags {
+        self.scopes.get_flags(self.current_scope_id)
+    }
+
     /// Get scopes tree
     #[inline]
     pub fn scopes(&self) -> &ScopeTree {
@@ -101,6 +112,62 @@ impl TraverseScoping {
         })
     }
 
+    /// Create new scope as child of current scope.
+    ///
+    /// `flags` provided are amended to inherit from parent scope's flags.
+    pub fn create_scope_child_of_current(&mut self, flags: ScopeFlags) -> ScopeId {
+        let flags = self.scopes.get_new_scope_flags(flags, self.current_scope_id);
+        self.scopes.add_scope(Some(self.current_scope_id), flags)
+    }
+
+    /// Insert a scope into scope tree below a statement.
+    ///
+    /// Statement must be in current scope.
+    /// New scope is created as child of current scope.
+    /// All child scopes of the statement are reassigned to be children of the new scope.
+    ///
+    /// `flags` provided are amended to inherit from parent scope's flags.
+    pub fn insert_scope_below_statement(&mut self, stmt: &Statement, flags: ScopeFlags) -> ScopeId {
+        let mut collector = ChildScopeCollector::new();
+        collector.visit_statement(stmt);
+        self.insert_scope_below(&collector.scope_ids, flags)
+    }
+
+    /// Insert a scope into scope tree below an expression.
+    ///
+    /// Expression must be in current scope.
+    /// New scope is created as child of current scope.
+    /// All child scopes of the expression are reassigned to be children of the new scope.
+    ///
+    /// `flags` provided are amended to inherit from parent scope's flags.
+    pub fn insert_scope_below_expression(
+        &mut self,
+        expr: &Expression,
+        flags: ScopeFlags,
+    ) -> ScopeId {
+        let mut collector = ChildScopeCollector::new();
+        collector.visit_expression(expr);
+        self.insert_scope_below(&collector.scope_ids, flags)
+    }
+
+    fn insert_scope_below(&mut self, child_scope_ids: &[ScopeId], flags: ScopeFlags) -> ScopeId {
+        // Remove these scopes from parent's children
+        if let Some(current_child_scope_ids) = self.scopes.get_child_ids_mut(self.current_scope_id)
+        {
+            current_child_scope_ids.retain(|scope_id| !child_scope_ids.contains(scope_id));
+        }
+
+        // Create new scope as child of parent
+        let new_scope_id = self.create_scope_child_of_current(flags);
+
+        // Set scopes as children of new scope instead
+        for &child_id in child_scope_ids {
+            self.scopes.set_parent_id(child_id, Some(new_scope_id));
+        }
+
+        new_scope_id
+    }
+
     /// Generate UID.
     ///
     /// Finds a unique variable name which does clash with any other variables used in the program.
@@ -372,3 +439,90 @@ fn create_uid_name_base(name: &str) -> CompactString {
     str.push_str(name);
     str
 }
+
+/// Visitor that locates all child scopes.
+/// NB: Child scopes only, not grandchild scopes.
+/// Does not do full traversal - stops each time it hits a node with a scope.
+struct ChildScopeCollector {
+    scope_ids: Vec<ScopeId>,
+}
+
+impl ChildScopeCollector {
+    fn new() -> Self {
+        Self { scope_ids: vec![] }
+    }
+}
+
+impl<'a> Visit<'a> for ChildScopeCollector {
+    fn visit_block_statement(&mut self, stmt: &BlockStatement<'a>) {
+        self.scope_ids.push(stmt.scope_id.get().unwrap());
+    }
+
+    fn visit_for_statement(&mut self, stmt: &ForStatement<'a>) {
+        if let Some(scope_id) = stmt.scope_id.get() {
+            self.scope_ids.push(scope_id);
+        } else {
+            walk::walk_for_statement(self, stmt);
+        }
+    }
+
+    fn visit_for_in_statement(&mut self, stmt: &ForInStatement<'a>) {
+        if let Some(scope_id) = stmt.scope_id.get() {
+            self.scope_ids.push(scope_id);
+        } else {
+            walk::walk_for_in_statement(self, stmt);
+        }
+    }
+
+    fn visit_for_of_statement(&mut self, stmt: &ForOfStatement<'a>) {
+        if let Some(scope_id) = stmt.scope_id.get() {
+            self.scope_ids.push(scope_id);
+        } else {
+            walk::walk_for_of_statement(self, stmt);
+        }
+    }
+
+    fn visit_switch_statement(&mut self, stmt: &SwitchStatement<'a>) {
+        self.scope_ids.push(stmt.scope_id.get().unwrap());
+    }
+
+    fn visit_catch_clause(&mut self, clause: &CatchClause<'a>) {
+        self.scope_ids.push(clause.scope_id.get().unwrap());
+    }
+
+    fn visit_finally_clause(&mut self, clause: &BlockStatement<'a>) {
+        self.scope_ids.push(clause.scope_id.get().unwrap());
+    }
+
+    fn visit_function(&mut self, func: &Function<'a>, _flags: Option<ScopeFlags>) {
+        self.scope_ids.push(func.scope_id.get().unwrap());
+    }
+
+    fn visit_class(&mut self, class: &Class<'a>) {
+        if let Some(scope_id) = class.scope_id.get() {
+            self.scope_ids.push(scope_id);
+        } else {
+            walk::walk_class(self, class);
+        }
+    }
+
+    fn visit_static_block(&mut self, block: &StaticBlock<'a>) {
+        self.scope_ids.push(block.scope_id.get().unwrap());
+    }
+
+    fn visit_arrow_expression(&mut self, expr: &ArrowFunctionExpression<'a>) {
+        self.scope_ids.push(expr.scope_id.get().unwrap());
+    }
+
+    fn visit_enum(&mut self, decl: &TSEnumDeclaration<'a>) {
+        self.scope_ids.push(decl.scope_id.get().unwrap());
+    }
+
+    fn visit_ts_module_declaration(&mut self, decl: &TSModuleDeclaration<'a>) {
+        self.scope_ids.push(decl.scope_id.get().unwrap());
+    }
+
+    fn visit_ts_type_parameter(&mut self, ty: &TSTypeParameter<'a>) {
+        self.scope_ids.push(ty.scope_id.get().unwrap());
+    }
+}