From de403a297de57454ef52dca6e3712f17f5077722 Mon Sep 17 00:00:00 2001
From: Boshen <boshenc@gmail.com>
Date: Mon, 12 Jun 2023 21:24:11 +0800
Subject: [PATCH] feat(hir): add enter_scope / leave_scope callbacks

---
 crates/oxc_hir/src/hir.rs   | 16 ++++++++++++++++
 crates/oxc_hir/src/visit.rs | 33 +++++++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+)

diff --git a/crates/oxc_hir/src/hir.rs b/crates/oxc_hir/src/hir.rs
index 45a861719..c0ae58f08 100644
--- a/crates/oxc_hir/src/hir.rs
+++ b/crates/oxc_hir/src/hir.rs
@@ -1131,6 +1131,14 @@ pub enum ForStatementInit<'a> {
     Expression(Expression<'a>),
 }
 
+impl<'a> ForStatementInit<'a> {
+    /// LexicalDeclaration[In, Yield, Await] :
+    ///   LetOrConst BindingList[?In, ?Yield, ?Await] ;
+    pub fn is_lexical_declaration(&self) -> bool {
+        matches!(self, Self::VariableDeclaration(decl) if decl.kind.is_lexical())
+    }
+}
+
 /// For-In Statement
 #[derive(Debug, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize), serde(tag = "type"))]
@@ -1161,6 +1169,14 @@ pub enum ForStatementLeft<'a> {
     AssignmentTarget(AssignmentTarget<'a>),
 }
 
+impl<'a> ForStatementLeft<'a> {
+    /// LexicalDeclaration[In, Yield, Await] :
+    ///   LetOrConst BindingList[?In, ?Yield, ?Await] ;
+    pub fn is_lexical_declaration(&self) -> bool {
+        matches!(self, Self::VariableDeclaration(decl) if decl.kind.is_lexical())
+    }
+}
+
 /// Continue Statement
 #[derive(Debug, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize), serde(tag = "type"))]
diff --git a/crates/oxc_hir/src/visit.rs b/crates/oxc_hir/src/visit.rs
index 2d1c5b68a..021344cec 100644
--- a/crates/oxc_hir/src/visit.rs
+++ b/crates/oxc_hir/src/visit.rs
@@ -6,6 +6,7 @@
 
 use oxc_allocator::Vec;
 use oxc_span::Span;
+use oxc_syntax::scope::ScopeFlags;
 
 #[allow(clippy::wildcard_imports)]
 use crate::{hir::*, hir_kind::HirKind};
@@ -14,6 +15,8 @@ use crate::{hir::*, hir_kind::HirKind};
 pub trait Visit<'a>: Sized {
     fn enter_node(&mut self, _kind: HirKind<'a>) {}
     fn leave_node(&mut self, _kind: HirKind<'a>) {}
+    fn enter_scope(&mut self, _flags: ScopeFlags) {}
+    fn leave_scope(&mut self) {}
 
     fn visit_program(&mut self, program: &'a Program<'a>) {
         let kind = HirKind::Program(program);
@@ -65,7 +68,9 @@ pub trait Visit<'a>: Sized {
     fn visit_block_statement(&mut self, stmt: &'a BlockStatement<'a>) {
         let kind = HirKind::BlockStatement(stmt);
         self.enter_node(kind);
+        self.enter_scope(ScopeFlags::empty());
         self.visit_statements(&stmt.body);
+        self.leave_scope();
         self.leave_node(kind);
     }
 
@@ -113,6 +118,11 @@ pub trait Visit<'a>: Sized {
     fn visit_for_statement(&mut self, stmt: &'a ForStatement<'a>) {
         let kind = HirKind::ForStatement(stmt);
         self.enter_node(kind);
+        let is_lexical_declaration =
+            stmt.init.as_ref().is_some_and(ForStatementInit::is_lexical_declaration);
+        if is_lexical_declaration {
+            self.enter_scope(ScopeFlags::empty());
+        }
         if let Some(init) = &stmt.init {
             self.visit_for_statement_init(init);
         }
@@ -125,6 +135,9 @@ pub trait Visit<'a>: Sized {
         if let Some(stmt) = &stmt.body {
             self.visit_statement(stmt);
         }
+        if is_lexical_declaration {
+            self.leave_scope();
+        }
         self.leave_node(kind);
     }
 
@@ -143,22 +156,36 @@ pub trait Visit<'a>: Sized {
     fn visit_for_in_statement(&mut self, stmt: &'a ForInStatement<'a>) {
         let kind = HirKind::ForInStatement(stmt);
         self.enter_node(kind);
+        let is_lexical_declaration = stmt.left.is_lexical_declaration();
+        if is_lexical_declaration {
+            self.enter_scope(ScopeFlags::empty());
+        }
         self.visit_for_statement_left(&stmt.left);
         self.visit_expression(&stmt.right);
         if let Some(stmt) = &stmt.body {
             self.visit_statement(stmt);
         }
+        if is_lexical_declaration {
+            self.leave_scope();
+        }
         self.leave_node(kind);
     }
 
     fn visit_for_of_statement(&mut self, stmt: &'a ForOfStatement<'a>) {
         let kind = HirKind::ForOfStatement(stmt);
         self.enter_node(kind);
+        let is_lexical_declaration = stmt.left.is_lexical_declaration();
+        if is_lexical_declaration {
+            self.enter_scope(ScopeFlags::empty());
+        }
         self.visit_for_statement_left(&stmt.left);
         self.visit_expression(&stmt.right);
         if let Some(stmt) = &stmt.body {
             self.visit_statement(stmt);
         }
+        if is_lexical_declaration {
+            self.leave_scope();
+        }
         self.leave_node(kind);
     }
 
@@ -246,10 +273,12 @@ pub trait Visit<'a>: Sized {
     fn visit_catch_clause(&mut self, clause: &'a CatchClause<'a>) {
         let kind = HirKind::CatchClause(clause);
         self.enter_node(kind);
+        self.enter_scope(ScopeFlags::empty());
         if let Some(param) = &clause.param {
             self.visit_binding_pattern(param);
         }
         self.visit_statements(&clause.body.body);
+        self.leave_scope();
         self.leave_node(kind);
     }
 
@@ -408,7 +437,9 @@ pub trait Visit<'a>: Sized {
     fn visit_static_block(&mut self, block: &'a StaticBlock<'a>) {
         let kind = HirKind::StaticBlock(block);
         self.enter_node(kind);
+        self.enter_scope(ScopeFlags::ClassStaticBlock);
         self.visit_statements(&block.body);
+        self.leave_scope();
         self.leave_node(kind);
     }
 
@@ -545,8 +576,10 @@ pub trait Visit<'a>: Sized {
     fn visit_arrow_expression(&mut self, expr: &'a ArrowExpression<'a>) {
         let kind = HirKind::ArrowExpression(expr);
         self.enter_node(kind);
+        self.enter_scope(ScopeFlags::Function);
         self.visit_formal_parameters(&expr.params);
         self.visit_function_body(&expr.body);
+        self.leave_scope();
         self.leave_node(kind);
     }