perf: accelerate variable scope lookup

crates/tree_shaker/src/module.rs

@@ -5,7 +5,7 @@ use oxc::{
   allocator::FromIn,
   ast::ast::{ImportDeclaration, Program, Statement},
   parser::Parser,
-  semantic::{Semantic, SemanticBuilder, SymbolId},
+  semantic::{ScopeId, Semantic, SemanticBuilder, SymbolId},
   span::{Atom, SourceType},
 };
 use oxc_index::{IndexVec, define_index_type};
@@ -30,6 +30,8 @@ pub struct ModuleInfo<'a> {
   pub semantic: Rc<Semantic<'a>>,
   pub call_id: DepAtom,
 
+  pub scopes_depth: FxHashMap<ScopeId, usize>,
+
   pub named_exports: FxHashMap<Atom<'a>, (VariableScopeId, SymbolId)>,
   pub default_export: Option<Entity<'a>>,
 
@@ -91,13 +93,16 @@ impl<'a> Analyzer<'a> {
     }
     let semantic = SemanticBuilder::new().build(unsafe { &*program.get() }).semantic;
     let semantic = Rc::new(semantic);
+    let scopes_depth = FxHashMap::from_iter([(semantic.scoping().root_scope_id(), 0)]);
     let module_id = self.modules.modules.push(ModuleInfo {
       path: Atom::from_in(path.clone(), self.allocator),
       line_index,
       program,
       semantic,
       call_id: DepAtom::from_counter(),
 
+      scopes_depth,
+
       named_exports: Default::default(),
       default_export: Default::default(),
 

crates/tree_shaker/src/nodes/misc/catch_clause.rs

@@ -8,6 +8,7 @@ use crate::{analyzer::Analyzer, ast::DeclarationKind, entity::Entity, transforme
 impl<'a> Analyzer<'a> {
   pub fn exec_catch_clause(&mut self, node: &'a CatchClause<'a>, value: Entity<'a>) {
     self.push_indeterminate_cf_scope();
+    self.set_variable_scope_depth(node.scope_id());
 
     if let Some(param) = &node.param {
       self.declare_binding_pattern(&param.pattern, false, DeclarationKind::Caught);

crates/tree_shaker/src/nodes/misc/class.rs

@@ -174,6 +174,7 @@ impl<'a> Analyzer<'a> {
       let function = constructor.value.as_ref();
       let dep = self.factory.dep(AstKind2::Function(function));
       self.cf_scope_mut().push_dep(dep);
+      self.set_variable_scope_depth(constructor.value.scope_id());
       self.exec_formal_parameters(&function.params, args, DeclarationKind::FunctionParameter);
       self.exec_function_body(function.body.as_ref().unwrap());
       if consume {

crates/tree_shaker/src/nodes/stmt/block_statement.rs

@@ -6,6 +6,8 @@ impl<'a> Analyzer<'a> {
   pub fn exec_block_statement(&mut self, node: &'a BlockStatement) {
     let data = self.load_data::<StatementVecData>(AstKind2::BlockStatement(node));
 
+    self.set_variable_scope_depth(node.scope_id());
+
     self.exec_statement_vec(data, &node.body);
   }
 }

crates/tree_shaker/src/nodes/stmt/for_in_statement.rs

@@ -9,6 +9,8 @@ impl<'a> Analyzer<'a> {
   pub fn exec_for_in_statement(&mut self, node: &'a ForInStatement<'a>) {
     let right = self.exec_expression(&node.right);
 
+    self.set_variable_scope_depth(node.scope_id());
+
     if let Some(keys) = right.get_own_keys(self) {
       let dep = self.factory.dep((right.shallow_dep(), AstKind2::ForInStatement(node)));
       self.push_cf_scope_with_deps(CfScopeKind::LoopBreak, self.factory.vec1(dep), Some(false));

crates/tree_shaker/src/nodes/stmt/for_of_statement.rs

@@ -16,6 +16,7 @@ impl<'a> Analyzer<'a> {
       right
     };
 
+    self.set_variable_scope_depth(node.scope_id());
     self.declare_for_statement_left(&node.left);
 
     let Some(iterated) = right.iterate_result_union(self, AstKind2::ForOfStatement(node)) else {

crates/tree_shaker/src/nodes/stmt/for_statement.rs

@@ -7,6 +7,8 @@ use crate::{analyzer::Analyzer, ast::AstKind2, scope::CfScopeKind, transformer::
 
 impl<'a> Analyzer<'a> {
   pub fn exec_for_statement(&mut self, node: &'a ForStatement<'a>) {
+    self.set_variable_scope_depth(node.scope_id());
+
     if let Some(init) = &node.init {
       match init {
         ForStatementInit::VariableDeclaration(node) => {

crates/tree_shaker/src/nodes/stmt/switch_statement.rs

@@ -10,6 +10,8 @@ use crate::{
 
 impl<'a> Analyzer<'a> {
   pub fn exec_switch_statement(&mut self, node: &'a SwitchStatement<'a>) {
+    self.set_variable_scope_depth(node.scope_id());
+
     // 1. discriminant
     let discriminant = self.exec_expression(&node.discriminant);
     self.push_dependent_cf_scope(discriminant);

crates/tree_shaker/src/scope/mod.rs

@@ -9,6 +9,7 @@ pub mod variable_scope;
 use call_scope::CallScope;
 use cf_scope::CfScope;
 pub use cf_scope::{CfScopeId, CfScopeKind};
+use oxc::semantic::ScopeId;
 use scope_tree::ScopeTree;
 use try_scope::TryScope;
 use variable_scope::VariableScope;
@@ -34,9 +35,7 @@ pub struct Scoping<'a> {
 
 impl<'a> Scoping<'a> {
   pub fn new(factory: &Factory<'a>) -> Self {
-    let mut variable = ScopeTree::new();
-    variable.push(VariableScope::new_with_this(factory.unknown));
-    let mut cf = ScopeTree::new();
+    let mut cf = ScopeTree::default();
     cf.push(CfScope::new(CfScopeKind::Root, factory.vec(), Some(false)));
     Scoping {
       call: vec![CallScope::new_in(
@@ -55,7 +54,7 @@ impl<'a> Scoping<'a> {
         false,
         factory.allocator,
       )],
-      variable,
+      variable: ScopeTree::default(),
       cf,
       pure: 0,
 
@@ -135,7 +134,7 @@ impl<'a> Analyzer<'a> {
 
     self.module_stack.push(callee.module_id);
     let old_variable_scope_stack = self.replace_variable_scope_stack(variable_scope_stack);
-    let body_variable_scope = self.push_variable_scope();
+    let body_variable_scope = self.push_variable_scope(callee.scope_id());
     let cf_scope_depth = self.push_cf_scope_with_deps(
       CfScopeKind::Function,
       self.factory.vec1(self.dep((call_dep, dep_id))),
@@ -164,8 +163,18 @@ impl<'a> Analyzer<'a> {
     ret_val
   }
 
-  pub fn push_variable_scope(&mut self) -> VariableScopeId {
-    self.scoping.variable.push(VariableScope::new())
+  pub fn set_variable_scope_depth(&mut self, scope: ScopeId) {
+    let depth = self.scoping.variable.current_depth();
+    if let Some(&existing) = self.module_info().scopes_depth.get(&scope) {
+      debug_assert_eq!(existing, depth, "Scope: {scope:?} already exists");
+    }
+    self.module_info_mut().scopes_depth.insert(scope, depth);
+  }
+
+  pub fn push_variable_scope(&mut self, scope: ScopeId) -> VariableScopeId {
+    let id = self.scoping.variable.push(VariableScope::new());
+    self.set_variable_scope_depth(scope);
+    id
   }
 
   pub fn pop_variable_scope(&mut self) -> VariableScopeId {

crates/tree_shaker/src/scope/scope_tree.rs

@@ -13,15 +13,11 @@ pub struct ScopeTree<I: Idx, T> {
 
 impl<I: Idx, T> Default for ScopeTree<I, T> {
   fn default() -> Self {
-    Self::new()
+    Self { nodes: IndexVec::new(), stack: vec![] }
   }
 }
 
 impl<I: Idx, T> ScopeTree<I, T> {
-  pub fn new() -> Self {
-    ScopeTree { nodes: IndexVec::new(), stack: vec![] }
-  }
-
   pub fn current_id(&self) -> I {
     *self.stack.last().unwrap()
   }

crates/tree_shaker/src/scope/variable_scope.rs

@@ -327,27 +327,29 @@ impl<'a> Analyzer<'a> {
     self.init_on_scope(variable_scope, symbol, value, init_node);
   }
 
+  fn symbol_scope_id(&mut self, symbol: SymbolId) -> VariableScopeId {
+    let scope_id = self.semantic().scoping().symbol_scope_id(symbol);
+    let depth = self.module_info().scopes_depth[&scope_id];
+    self.scoping.variable.stack[depth]
+  }
+
   /// `None` for TDZ
   pub fn read_symbol(&mut self, symbol: SymbolId) -> Option<Entity<'a>> {
-    for depth in (0..self.scoping.variable.stack.len()).rev() {
-      let id = self.scoping.variable.stack[depth];
-      if let Some(value) = self.read_on_scope(id, symbol) {
-        return value;
-      }
+    let scope = self.symbol_scope_id(symbol);
+    if let Some(value) = self.read_on_scope(scope, symbol) {
+      value
+    } else {
+      self.mark_unresolved_reference(symbol);
+      Some(self.factory.unknown)
     }
-    self.mark_unresolved_reference(symbol);
-    Some(self.factory.unknown)
   }
 
   pub fn write_symbol(&mut self, symbol: SymbolId, new_val: Entity<'a>) {
-    for depth in (0..self.scoping.variable.stack.len()).rev() {
-      let id = self.scoping.variable.stack[depth];
-      if self.write_on_scope(id, symbol, new_val) {
-        return;
-      }
+    let scope = self.symbol_scope_id(symbol);
+    if !self.write_on_scope(scope, symbol, new_val) {
+      self.consume(new_val);
+      self.mark_unresolved_reference(symbol);
     }
-    self.consume(new_val);
-    self.mark_unresolved_reference(symbol);
   }
 
   fn mark_unresolved_reference(&mut self, symbol: SymbolId) {

crates/tree_shaker/src/utils/callee_info.rs

@@ -81,6 +81,16 @@ impl<'a> CalleeInfo<'a> {
   pub fn into_node(self) -> AstKind2<'a> {
     self.node.into()
   }
+
+  pub fn scope_id(self) -> ScopeId {
+    match self.node {
+      CalleeNode::Function(node) => node.scope_id(),
+      CalleeNode::ArrowFunctionExpression(node) => node.scope_id(),
+      CalleeNode::ClassStatics(node) => node.scope_id(),
+      CalleeNode::ClassConstructor(node) => node.scope_id(),
+      _ => unreachable!(),
+    }
+  }
 }
 
 impl<'a> Analyzer<'a> {