From 01cc2ceac0253df5b3d5c5bddf30c2dc91d97faa Mon Sep 17 00:00:00 2001
From: Boshen <1430279+Boshen@users.noreply.github.com>
Date: Mon, 2 Sep 2024 16:22:26 +0000
Subject: [PATCH] feat(semantic)!: add `ScopeTree:get_child_ids` API behind a
 runtime flag (#5403)

The runtime API for enabling `get_child_ids` is `SemanticBuilder::with_scope_tree_child_ids`.

I've tested this PR in Rolldown to ensure correctness.
---
 crates/oxc_semantic/src/builder.rs            |  6 ++++
 crates/oxc_semantic/src/scope.rs              | 30 +++++++++++++++++++
 .../oxc_semantic/tests/integration/scopes.rs  | 18 +++++++++++
 .../tests/integration/util/mod.rs             |  9 ++++++
 4 files changed, 63 insertions(+)

diff --git a/crates/oxc_semantic/src/builder.rs b/crates/oxc_semantic/src/builder.rs
index 3987d51a18e95..a1715dcb1eb07 100644
--- a/crates/oxc_semantic/src/builder.rs
+++ b/crates/oxc_semantic/src/builder.rs
@@ -186,6 +186,12 @@ impl<'a> SemanticBuilder<'a> {
         self
     }
 
+    #[must_use]
+    pub fn with_scope_tree_child_ids(mut self, yes: bool) -> Self {
+        self.scope.build_child_ids = yes;
+        self
+    }
+
     /// Get the built module record from `build_module_record`
     pub fn module_record(&self) -> Arc<ModuleRecord> {
         Arc::clone(&self.module_record)
diff --git a/crates/oxc_semantic/src/scope.rs b/crates/oxc_semantic/src/scope.rs
index 5616612189f09..56498c12f3290 100644
--- a/crates/oxc_semantic/src/scope.rs
+++ b/crates/oxc_semantic/src/scope.rs
@@ -28,13 +28,23 @@ pub type UnresolvedReferences = FxHashMap<CompactStr, Vec<ReferenceId>>;
 pub struct ScopeTree {
     /// Maps a scope to the parent scope it belongs in.
     parent_ids: IndexVec<ScopeId, Option<ScopeId>>,
+
+    /// Maps a scope to direct children scopes.
+    child_ids: IndexVec<ScopeId, Vec<ScopeId>>,
+
+    /// Runtime flag for constructing child_ids.
+    pub(crate) build_child_ids: bool,
+
     /// Maps a scope to its node id.
     node_ids: IndexVec<ScopeId, AstNodeId>,
+
     flags: IndexVec<ScopeId, ScopeFlags>,
+
     /// Symbol bindings in a scope.
     ///
     /// A binding is a mapping from an identifier name to its [`SymbolId`]
     bindings: IndexVec<ScopeId, Bindings>,
+
     pub(crate) root_unresolved_references: UnresolvedReferences,
 }
 
@@ -154,6 +164,11 @@ impl ScopeTree {
 
     pub fn set_parent_id(&mut self, scope_id: ScopeId, parent_id: Option<ScopeId>) {
         self.parent_ids[scope_id] = parent_id;
+        if self.build_child_ids {
+            if let Some(parent_id) = parent_id {
+                self.child_ids[parent_id].push(scope_id);
+            }
+        }
     }
 
     /// Get a variable binding by name that was declared in the top-level scope
@@ -232,6 +247,12 @@ impl ScopeTree {
         &mut self.bindings[scope_id]
     }
 
+    /// Get the child scopes of a scope
+    #[inline]
+    pub fn get_child_ids(&self, scope_id: ScopeId) -> &[ScopeId] {
+        &self.child_ids[scope_id]
+    }
+
     /// Create a scope.
     #[inline]
     pub fn add_scope(
@@ -244,6 +265,12 @@ impl ScopeTree {
         self.flags.push(flags);
         self.bindings.push(Bindings::default());
         self.node_ids.push(node_id);
+        if self.build_child_ids {
+            self.child_ids.push(vec![]);
+            if let Some(parent_id) = parent_id {
+                self.child_ids[parent_id].push(scope_id);
+            }
+        }
         scope_id
     }
 
@@ -265,5 +292,8 @@ impl ScopeTree {
         self.flags.reserve(additional);
         self.bindings.reserve(additional);
         self.node_ids.reserve(additional);
+        if self.build_child_ids {
+            self.child_ids.reserve(additional);
+        }
     }
 }
diff --git a/crates/oxc_semantic/tests/integration/scopes.rs b/crates/oxc_semantic/tests/integration/scopes.rs
index c7d4ffc7c0974..cd12d4fe9c03f 100644
--- a/crates/oxc_semantic/tests/integration/scopes.rs
+++ b/crates/oxc_semantic/tests/integration/scopes.rs
@@ -220,3 +220,21 @@ fn var_hoisting() {
     .is_in_scope(ScopeFlags::Top)
     .test();
 }
+
+#[test]
+fn get_child_ids() {
+    let test = SemanticTester::js(
+        "
+            function foo() {
+            }
+        ",
+    )
+    .with_scope_tree_child_ids(true);
+    let semantic = test.build();
+    let (_symbols, scopes) = semantic.into_symbol_table_and_scope_tree();
+
+    let child_scope_ids = scopes.get_child_ids(scopes.root_scope_id());
+    assert_eq!(child_scope_ids.len(), 1);
+    let child_scope_ids = scopes.get_child_ids(child_scope_ids[0]);
+    assert!(child_scope_ids.is_empty());
+}
diff --git a/crates/oxc_semantic/tests/integration/util/mod.rs b/crates/oxc_semantic/tests/integration/util/mod.rs
index 38d13714d5597..5d42f17c88df9 100644
--- a/crates/oxc_semantic/tests/integration/util/mod.rs
+++ b/crates/oxc_semantic/tests/integration/util/mod.rs
@@ -25,6 +25,8 @@ pub struct SemanticTester<'a> {
     ///
     /// [`ControlFlowGraph`]: oxc_cfg::ControlFlowGraph
     cfg: bool,
+    /// Build a child ids for scope tree?
+    scope_tree_child_ids: bool,
     /// Expect semantic analysis to produce errors.
     ///
     /// Default is `false`.
@@ -63,6 +65,7 @@ impl<'a> SemanticTester<'a> {
             source_type,
             source_text,
             cfg: false,
+            scope_tree_child_ids: false,
             expect_errors: false,
         }
     }
@@ -95,6 +98,11 @@ impl<'a> SemanticTester<'a> {
         self
     }
 
+    pub fn with_scope_tree_child_ids(mut self, yes: bool) -> Self {
+        self.scope_tree_child_ids = yes;
+        self
+    }
+
     /// The program being tested is expected to produce errors during semantic analysis.
     ///
     /// By default, programs are expected to be error-free.
@@ -163,6 +171,7 @@ impl<'a> SemanticTester<'a> {
             .with_check_syntax_error(true)
             .with_trivias(parse.trivias)
             .with_cfg(self.cfg)
+            .with_scope_tree_child_ids(self.scope_tree_child_ids)
             .build(program)
     }