perf: cf_scope as a stack

Author: kermanx

.vscode/settings.json

@@ -1,6 +1,6 @@
 {
   "rust-analyzer.procMacro.ignored": { "napi-derive": ["napi"] },
-  "rust-analyzer.cargo.features": ["flame"],
+  "rust-analyzer.cargo.features": "all",
 
   "[rust]": {
     "editor.formatOnSave": true,

crates/jsshaker/src/analyzer/exhaustive.rs

@@ -83,7 +83,7 @@ impl<'a> Analyzer<'a> {
     runner: Rc<dyn Fn(&mut Analyzer<'a>) -> Entity<'a> + 'a>,
   ) -> Entity<'a> {
     self.push_cf_scope(
-      CfScopeKind::Exhaustive(self.allocator.alloc(ExhaustiveData {
+      CfScopeKind::Exhaustive(Box::new(ExhaustiveData {
         clean: true,
         temp_deps: drain.then(|| allocator::HashSet::new_in(self.allocator)),
         register_deps: register.then(|| allocator::HashSet::new_in(self.allocator)),
@@ -114,9 +114,11 @@ impl<'a> Analyzer<'a> {
         break;
       }
     }
-    let id = self.pop_cf_scope();
-    let data = self.scoping.cf.get_mut(id).exhaustive_data_mut().unwrap();
-    if let Some(register_deps) = data.register_deps.take() {
+    let scope = self.pop_cf_scope();
+    let CfScopeKind::Exhaustive(data) = scope.kind else {
+      unreachable!();
+    };
+    if let Some(register_deps) = data.register_deps {
       self.register_exhaustive_callbacks(drain, runner, register_deps);
     }
     first_ret.unwrap()

crates/jsshaker/src/analyzer/rw_tracking.rs

@@ -52,8 +52,8 @@ impl<'a> Analyzer<'a> {
     let target_depth = self.find_first_different_cf_scope(scope);
     let mut registered = false;
     let mut tracker_dep = None;
-    for depth in (target_depth..self.scoping.cf.stack.len()).rev() {
-      let scope = self.scoping.cf.get_mut_from_depth(depth);
+    for depth in (target_depth..self.scoping.cf.stack_len()).rev() {
+      let scope = self.scoping.cf.data_at_mut(depth);
       if let Some(data) = scope.exhaustive_data_mut() {
         if data.clean
           && let Some(temp_deps) = data.temp_deps.as_mut()
@@ -89,8 +89,8 @@ impl<'a> Analyzer<'a> {
     let mut exhaustive = false;
     let mut must_mark = true;
     let mut has_fn_scope = false;
-    for depth in scope_depth..self.scoping.cf.stack.len() {
-      let scope = self.scoping.cf.get_mut_from_depth(depth);
+    for depth in scope_depth..self.scoping.cf.stack_len() {
+      let scope = self.scoping.cf.data_at_mut(depth);
       has_fn_scope |= scope.kind.is_function();
       if let Some(data) = scope.exhaustive_data_mut() {
         exhaustive = true;
@@ -109,8 +109,8 @@ impl<'a> Analyzer<'a> {
     }
     if has_fn_scope {
       let mut non_det = false;
-      for depth in (scope_depth..self.scoping.cf.stack.len()).rev() {
-        let scope = self.scoping.cf.get_mut_from_depth(depth);
+      for depth in (scope_depth..self.scoping.cf.stack_len()).rev() {
+        let scope = self.scoping.cf.data_at_mut(depth);
         non_det |= scope.non_det();
         if let Some(data) = scope.fn_cache_tracking_data_mut() {
           data.track_write(target, cacheable.map(|e| (non_det, e)));

crates/jsshaker/src/builtins/react/dependencies.rs

@@ -140,10 +140,8 @@ pub fn check_dependencies<'a>(
   };
 
   if require_rerun {
-    for depth in 0..analyzer.scoping.cf.stack.len() {
-      if let Some(exhaustive_data) =
-        analyzer.scoping.cf.get_mut_from_depth(depth).exhaustive_data_mut()
-      {
+    for scope in analyzer.scoping.cf.iter_stack_mut() {
+      if let Some(exhaustive_data) = scope.exhaustive_data_mut() {
         exhaustive_data.clean = false;
         break;
       }

crates/jsshaker/src/nodes/stmt/if_statement.rs

@@ -19,7 +19,7 @@ impl<'a> Analyzer<'a> {
 
     let mut both_exit = true;
     let mut exit_target_inner = 0;
-    let mut exit_target_outer = self.scoping.cf.stack.len();
+    let mut exit_target_outer = self.scoping.cf.stack_len();
     let mut acc_dep_1 = None;
     let mut acc_dep_2 = None;
 
@@ -34,7 +34,7 @@ impl<'a> Analyzer<'a> {
         node.alternate.is_some(),
       );
       self.exec_statement(&node.consequent);
-      let conditional_scope = self.pop_cf_scope_and_get_mut();
+      let mut conditional_scope = self.pop_cf_scope();
       if let CfScopeKind::ExitBlocker(Some(stopped_exit)) = &conditional_scope.kind {
         exit_target_inner = exit_target_inner.max(*stopped_exit);
         exit_target_outer = exit_target_outer.min(*stopped_exit);
@@ -55,7 +55,7 @@ impl<'a> Analyzer<'a> {
       );
       if let Some(alternate) = &node.alternate {
         self.exec_statement(alternate);
-        let conditional_scope = self.pop_cf_scope_and_get_mut();
+        let mut conditional_scope = self.pop_cf_scope();
         if let CfScopeKind::ExitBlocker(Some(stopped_exit)) = &conditional_scope.kind {
           exit_target_inner = exit_target_inner.max(*stopped_exit);
           exit_target_outer = exit_target_outer.min(*stopped_exit);
@@ -72,12 +72,12 @@ impl<'a> Analyzer<'a> {
     let acc_dep = Some(self.dep((acc_dep_1, acc_dep_2)));
     if both_exit {
       if let Some(acc_dep) =
-        self.exit_to_impl(exit_target_inner, self.scoping.cf.stack.len(), true, acc_dep)
+        self.exit_to_impl(exit_target_inner, self.scoping.cf.stack_len(), true, acc_dep)
       {
         self.exit_to_impl(exit_target_outer, exit_target_inner, false, acc_dep);
       }
     } else {
-      self.exit_to_impl(exit_target_outer, self.scoping.cf.stack.len(), false, acc_dep);
+      self.exit_to_impl(exit_target_outer, self.scoping.cf.stack_len(), false, acc_dep);
     }
   }
 }

crates/jsshaker/src/scope/cf_scope.rs

@@ -2,13 +2,13 @@ use oxc::{ast::ast::LabeledStatement, span::Atom};
 
 use crate::{
   analyzer::{Analyzer, exhaustive::ExhaustiveData},
-  define_box_bump_idx,
+  define_stacked_tree_idx,
   dep::{Dep, DepCollector, DepVec},
   utils::ast::AstKind2,
   value::cache::FnCacheTrackingData,
 };
 
-define_box_bump_idx! {
+define_stacked_tree_idx! {
   pub struct CfScopeId;
 }
 
@@ -17,14 +17,14 @@ pub enum CfScopeKind<'a> {
   Root,
   Module,
   Labeled(&'a LabeledStatement<'a>),
-  Function(&'a mut FnCacheTrackingData<'a>),
+  Function(Box<FnCacheTrackingData<'a>>),
   LoopBreak,
   LoopContinue,
   Switch,
 
   Dependent,
   NonDet,
-  Exhaustive(&'a mut ExhaustiveData<'a>),
+  Exhaustive(Box<ExhaustiveData<'a>>),
   ExitBlocker(Option<usize>),
 }
 
@@ -157,8 +157,8 @@ impl<'a> Analyzer<'a> {
   pub fn get_exec_dep(&mut self, target_depth: usize) -> (DepVec<'a>, bool) {
     let mut deps = self.factory.vec();
     let mut non_det = false;
-    for id in target_depth..self.scoping.cf.stack.len() {
-      let scope = self.scoping.cf.get_mut_from_depth(id);
+    for id in target_depth..self.scoping.cf.stack_len() {
+      let scope = self.scoping.cf.data_at_mut(id);
       if let Some(dep) = scope.deps.collect(self.factory) {
         deps.push(dep);
       }
@@ -168,11 +168,11 @@ impl<'a> Analyzer<'a> {
   }
 
   pub fn exit_to(&mut self, target_depth: usize) {
-    self.exit_to_impl(target_depth, self.scoping.cf.stack.len(), true, None);
+    self.exit_to_impl(target_depth, self.scoping.cf.stack_len(), true, None);
   }
 
   pub fn exit_to_not_must(&mut self, target_depth: usize) {
-    self.exit_to_impl(target_depth, self.scoping.cf.stack.len(), false, None);
+    self.exit_to_impl(target_depth, self.scoping.cf.stack_len(), false, None);
   }
 
   /// `None` => Interrupted by if branch
@@ -185,8 +185,7 @@ impl<'a> Analyzer<'a> {
     mut acc_dep: Option<Dep<'a>>,
   ) -> Option<Option<Dep<'a>>> {
     for depth in (target_depth..from_depth).rev() {
-      let id = self.scoping.cf.stack[depth];
-      let cf_scope = self.scoping.cf.get_mut(id);
+      let cf_scope = self.scoping.cf.data_at_mut(depth);
 
       if cf_scope.must_exited() {
         return Some(Some(self.factory.no_dep));
@@ -310,8 +309,7 @@ impl<'a> Analyzer<'a> {
   pub fn global_effect(&mut self) {
     let mut deps = vec![];
     let mut first_stage = true;
-    for scope in self.scoping.cf.stack.iter().rev() {
-      let scope = &mut self.scoping.cf.nodes.get_mut(*scope).data;
+    for scope in self.scoping.cf.iter_stack_mut().rev() {
       if first_stage {
         match scope.deoptimize_state {
           DeoptimizeState::Never => {

crates/jsshaker/src/scope/mod.rs

@@ -6,8 +6,6 @@ mod stacked_tree;
 mod linked_tree;
 pub mod variable_scope;
 
-use std::mem;
-
 use call_scope::CallScope;
 use cf_scope::CfScope;
 pub use cf_scope::{CfScopeId, CfScopeKind};
@@ -28,7 +26,7 @@ use crate::{
 pub struct Scoping<'a> {
   pub call: Vec<CallScope<'a>>,
   pub variable: LinkedTree<'a, VariableScopeId, VariableScope<'a>>,
-  pub cf: StackedTree<'a, CfScopeId, CfScope<'a>>,
+  pub cf: StackedTree<CfScopeId, CfScope<'a>>,
   pub root_cf_scope: CfScopeId,
   pub try_catch_depth: Option<usize>,
   pub current_callsite: AstKind2<'a>,
@@ -40,8 +38,8 @@ impl<'a> Scoping<'a> {
     let mut variable = LinkedTree::new_in(factory.allocator);
     let root_variable_scope =
       variable.push(VariableScope::new_in_with_this(factory.allocator, factory.unknown));
-    let mut cf = StackedTree::new_in(factory.allocator);
-    let root_cf_scope = cf.push(CfScope::new(CfScopeKind::Root, factory.vec(), false));
+    let cf = StackedTree::new(CfScope::new(CfScopeKind::Root, factory.vec(), false));
+    let root_cf_scope = cf.root;
     factory.root_cf_scope = Some(root_cf_scope);
     Scoping {
       call: vec![CallScope::new_in(
@@ -80,11 +78,11 @@ impl<'a> Analyzer<'a> {
   }
 
   pub fn cf_scope(&self) -> &CfScope<'a> {
-    self.scoping.cf.get_current()
+    self.scoping.cf.current_data()
   }
 
   pub fn cf_scope_mut(&mut self) -> &mut CfScope<'a> {
-    self.scoping.cf.get_current_mut()
+    self.scoping.cf.current_data_mut()
   }
 
   pub fn variable_scope(&self) -> &VariableScope<'a> {
@@ -111,9 +109,7 @@ impl<'a> Analyzer<'a> {
     let old_variable_scope_stack = self.replace_variable_scope(info.func.lexical_scope);
     let body_variable_scope = self.push_variable_scope();
     let cf_scope_depth = self.push_cf_scope_with_deps(
-      CfScopeKind::Function(
-        self.allocator.alloc(FnCacheTrackingData::new_in(self.allocator, info)),
-      ),
+      CfScopeKind::Function(Box::new(FnCacheTrackingData::new_in(self.allocator, info))),
       self.factory.vec1(info.call_dep),
       false,
     );
@@ -133,17 +129,14 @@ impl<'a> Analyzer<'a> {
   pub fn pop_call_scope(&mut self) -> (Entity<'a>, FnCacheTrackingData<'a>) {
     let scope = self.scoping.call.pop().unwrap();
     let ret_val = scope.ret_val(self);
-    let cf_scope_id = self.pop_cf_scope();
-    let cf_scope = self.scoping.cf.get_mut(cf_scope_id);
-    let CfScopeKind::Function(tracking_data) = &mut cf_scope.kind else {
+    let cf_scope = self.pop_cf_scope();
+    let CfScopeKind::Function(tracking_data) = cf_scope.kind else {
       unreachable!();
     };
-    let tracking_data = mem::take(*tracking_data);
-
     self.pop_variable_scope();
     self.replace_variable_scope(scope.old_variable_scope);
     self.set_current_module(scope.old_module);
-    (ret_val, tracking_data)
+    (ret_val, *tracking_data)
   }
 
   pub fn push_variable_scope(&mut self) -> VariableScopeId {
@@ -180,23 +173,18 @@ impl<'a> Analyzer<'a> {
     );
   }
 
-  pub fn pop_cf_scope(&mut self) -> CfScopeId {
+  pub fn pop_cf_scope(&mut self) -> CfScope<'a> {
     self.scoping.cf.pop()
   }
 
   pub fn pop_multiple_cf_scopes(&mut self, count: usize) -> Option<Dep<'a>> {
     let mut exec_deps = self.factory.vec();
     for _ in 0..count {
-      let id = self.scoping.cf.stack.pop().unwrap();
-      if let Some(dep) = self.scoping.cf.get_mut(id).deps.collect(self.factory) {
+      let mut scope = self.pop_cf_scope();
+      if let Some(dep) = scope.deps.collect(self.factory) {
         exec_deps.push(dep);
       }
     }
     if exec_deps.is_empty() { None } else { Some(self.dep(exec_deps)) }
   }
-
-  pub fn pop_cf_scope_and_get_mut(&mut self) -> &mut CfScope<'a> {
-    let id = self.pop_cf_scope();
-    self.scoping.cf.get_mut(id)
-  }
 }