feat(es/ast): Add more utilities (#9054)

swc-project · Jun 15, 2024 · ab226dc · ab226dc
1 parent ea7191e
commit ab226dc
Show file tree

Hide file tree

Showing 11 changed files with 110 additions and 42 deletions.
diff --git a/crates/swc_ecma_ast/src/expr.rs b/crates/swc_ecma_ast/src/expr.rs
@@ -1,4 +1,6 @@
 #![allow(clippy::vec_box)]
+use std::mem::transmute;
+
 use is_macro::Is;
 use string_enum::StringEnum;
 use swc_atoms::Atom;
@@ -203,17 +205,60 @@ impl Expr {
         }
     }
 
+    /// Unwraps an expression with a given function.
+    ///
+    /// If the provided function returns [Some], the function is called again
+    /// with the returned value. If the provided functions returns [None],
+    /// the last expression is returned.
+    pub fn unwrap_with<'a, F>(&'a self, mut op: F) -> &'a Expr
+    where
+        F: FnMut(&'a Expr) -> Option<&'a Expr>,
+    {
+        let mut cur = self;
+        loop {
+            match op(cur) {
+                Some(next) => cur = next,
+                None => return cur,
+            }
+        }
+    }
+
+    /// Unwraps an expression with a given function.
+    ///
+    /// If the provided function returns [Some], the function is called again
+    /// with the returned value. If the provided functions returns [None],
+    /// the last expression is returned.
+    pub fn unwrap_mut_with<'a, F>(&'a mut self, mut op: F) -> &'a mut Expr
+    where
+        F: FnMut(&'a mut Expr) -> Option<&'a mut Expr>,
+    {
+        let mut cur = self;
+        loop {
+            match unsafe {
+                // Safety: Polonius is not yet stable
+                op(transmute::<&mut _, &mut _>(cur))
+            } {
+                Some(next) => cur = next,
+                None => {
+                    return cur;
+                }
+            }
+        }
+    }
+
     /// Normalize parenthesized expressions.
     ///
     /// This will normalize `(foo)`, `((foo))`, ... to `foo`.
     ///
     /// If `self` is not a parenthesized expression, it will be returned as is.
     pub fn unwrap_parens(&self) -> &Expr {
-        let mut cur = self;
-        while let Expr::Paren(ref expr) = cur {
-            cur = &expr.expr;
-        }
-        cur
+        self.unwrap_with(|e| {
+            if let Expr::Paren(expr) = e {
+                Some(&expr.expr)
+            } else {
+                None
+            }
+        })
     }
 
     /// Normalize parenthesized expressions.
@@ -222,11 +267,25 @@ impl Expr {
     ///
     /// If `self` is not a parenthesized expression, it will be returned as is.
     pub fn unwrap_parens_mut(&mut self) -> &mut Expr {
-        let mut cur = self;
-        while let Expr::Paren(ref mut expr) = cur {
-            cur = &mut expr.expr;
-        }
-        cur
+        self.unwrap_mut_with(|e| {
+            if let Expr::Paren(expr) = e {
+                Some(&mut expr.expr)
+            } else {
+                None
+            }
+        })
+    }
+
+    /// Normalize sequences and parenthesized expressions.
+    ///
+    /// This returns the last expression of a sequence expression or the
+    /// expression of a parenthesized expression.
+    pub fn unwrap_seqs_and_parens(&self) -> &Self {
+        self.unwrap_with(|expr| match expr {
+            Expr::Seq(SeqExpr { exprs, .. }) => exprs.last().map(|v| &**v),
+            Expr::Paren(ParenExpr { expr, .. }) => Some(expr),
+            _ => None,
+        })
     }
 
     /// Creates an expression from `exprs`. This will return first element if

diff --git a/crates/swc_ecma_ast/src/module_decl.rs b/crates/swc_ecma_ast/src/module_decl.rs
@@ -219,6 +219,31 @@ pub enum ImportSpecifier {
     Namespace(ImportStarAsSpecifier),
 }
 
+impl ImportSpecifier {
+    pub fn is_type_only(&self) -> bool {
+        match self {
+            ImportSpecifier::Named(named) => named.is_type_only,
+            ImportSpecifier::Default(..) | ImportSpecifier::Namespace(..) => false,
+        }
+    }
+
+    pub fn local(&self) -> &Ident {
+        match self {
+            ImportSpecifier::Named(named) => &named.local,
+            ImportSpecifier::Default(default) => &default.local,
+            ImportSpecifier::Namespace(ns) => &ns.local,
+        }
+    }
+
+    pub fn local_mut(&mut self) -> &mut Ident {
+        match self {
+            ImportSpecifier::Named(named) => &mut named.local,
+            ImportSpecifier::Default(default) => &mut default.local,
+            ImportSpecifier::Namespace(ns) => &mut ns.local,
+        }
+    }
+}
+
 /// e.g. `import foo from 'mod.js'`
 #[ast_node("ImportDefaultSpecifier")]
 #[derive(Eq, Hash, EqIgnoreSpan)]

diff --git a/crates/swc_ecma_lints/src/rules/constructor_super.rs b/crates/swc_ecma_lints/src/rules/constructor_super.rs
@@ -7,7 +7,6 @@ use swc_ecma_visit::{Visit, VisitWith};
 use crate::{
     config::{LintRuleReaction, RuleConfig},
     rule::{visitor_rule, Rule},
-    rules::utils::unwrap_seqs_and_parens,
 };
 
 const BAD_SUPER_MESSAGE: &str = "Unexpected 'super()' because 'super' is not a constructor";
@@ -81,7 +80,7 @@ impl ConstructorSuper {
 
     fn collect_class(&mut self, class: &Class) {
         self.class_meta.super_class = match &class.super_class {
-            Some(super_class) => match unwrap_seqs_and_parens(super_class.as_ref()) {
+            Some(super_class) => match super_class.unwrap_seqs_and_parens() {
                 Expr::Ident(_) | Expr::Class(_) => SuperClass::Valid,
                 _ => SuperClass::Invalid,
             },

diff --git a/crates/swc_ecma_lints/src/rules/no_compare_neg_zero.rs b/crates/swc_ecma_lints/src/rules/no_compare_neg_zero.rs
@@ -5,7 +5,6 @@ use swc_ecma_visit::{Visit, VisitWith};
 use crate::{
     config::{LintRuleReaction, RuleConfig},
     rule::{visitor_rule, Rule},
-    rules::utils::unwrap_seqs_and_parens,
 };
 
 pub fn no_compare_neg_zero(config: &RuleConfig<()>) -> Option<Box<dyn Rule>> {
@@ -46,10 +45,9 @@ impl NoCompareNegZero {
             op: op!(unary, "-"),
             arg,
             ..
-        }) = unwrap_seqs_and_parens(expr)
+        }) = expr.unwrap_seqs_and_parens()
         {
-            if let Expr::Lit(Lit::Num(Number { value, .. })) = unwrap_seqs_and_parens(arg.as_ref())
-            {
+            if let Expr::Lit(Lit::Num(Number { value, .. })) = arg.unwrap_seqs_and_parens() {
                 return *value == 0f64;
             }
         }

diff --git a/crates/swc_ecma_lints/src/rules/no_param_reassign.rs b/crates/swc_ecma_lints/src/rules/no_param_reassign.rs
@@ -13,7 +13,6 @@ use swc_ecma_visit::{noop_visit_type, Visit, VisitWith};
 use crate::{
     config::{LintRuleReaction, RuleConfig},
     rule::{visitor_rule, Rule},
-    rules::utils::unwrap_seqs_and_parens,
 };
 
 const INVALID_REGEX_MESSAGE: &str = "no-param-reassign: invalid regex pattern in allowPattern. Check syntax documentation https://docs.rs/regex/latest/regex/#syntax";
@@ -150,7 +149,7 @@ impl NoParamReassign {
             return;
         }
 
-        match unwrap_seqs_and_parens(member_expr.obj.as_ref()) {
+        match member_expr.obj.unwrap_seqs_and_parens() {
             Expr::Ident(ident) => {
                 if self.is_satisfying_function_param(ident) {
                     self.emit_report(ident.span, &ident.sym);
@@ -190,7 +189,7 @@ impl NoParamReassign {
     }
 
     fn check_expr(&self, expr: &Expr) {
-        match unwrap_seqs_and_parens(expr) {
+        match expr.unwrap_seqs_and_parens() {
             Expr::Ident(ident) => {
                 if self.is_satisfying_function_param(ident) {
                     self.emit_report(ident.span, &ident.sym);

diff --git a/crates/swc_ecma_lints/src/rules/no_throw_literal.rs b/crates/swc_ecma_lints/src/rules/no_throw_literal.rs
@@ -5,7 +5,6 @@ use swc_ecma_visit::{Visit, VisitWith};
 use crate::{
     config::{LintRuleReaction, RuleConfig},
     rule::{visitor_rule, Rule},
-    rules::utils::unwrap_seqs_and_parens,
 };
 
 const EXPECTED_AN_ERROR_OBJECT: &str = "Expected an error object to be thrown";
@@ -52,7 +51,7 @@ impl NoThrowLiteral {
 
     #[allow(clippy::only_used_in_recursion)]
     fn could_be_error(&self, expr: &Expr) -> bool {
-        match unwrap_seqs_and_parens(expr) {
+        match expr.unwrap_seqs_and_parens() {
             Expr::Ident(_)
             | Expr::New(_)
             | Expr::Call(_)
@@ -96,7 +95,7 @@ impl NoThrowLiteral {
     }
 
     fn check(&self, throw_stmt: &ThrowStmt) {
-        let arg = unwrap_seqs_and_parens(throw_stmt.arg.as_ref());
+        let arg = throw_stmt.arg.unwrap_seqs_and_parens();
 
         if !self.could_be_error(arg) {
             self.emit_report(throw_stmt.span, EXPECTED_AN_ERROR_OBJECT);

diff --git a/crates/swc_ecma_lints/src/rules/prefer_const.rs b/crates/swc_ecma_lints/src/rules/prefer_const.rs
@@ -6,7 +6,6 @@ use swc_ecma_visit::{Visit, VisitWith};
 use crate::{
     config::{LintRuleReaction, RuleConfig},
     rule::{visitor_rule, Rule},
-    rules::utils::unwrap_seqs_and_parens,
 };
 
 // todo: implement option destructuring: all | any
@@ -293,7 +292,7 @@ impl Visit for PreferConst {
     }
 
     fn visit_update_expr(&mut self, update_expr: &UpdateExpr) {
-        if let Expr::Ident(ident) = unwrap_seqs_and_parens(update_expr.arg.as_ref()) {
+        if let Expr::Ident(ident) = update_expr.arg.unwrap_seqs_and_parens() {
             self.consider_mutation_for_ident(ident, false);
         }
 

diff --git a/crates/swc_ecma_lints/src/rules/prefer_regex_literals.rs b/crates/swc_ecma_lints/src/rules/prefer_regex_literals.rs
@@ -6,7 +6,7 @@ use swc_ecma_visit::{noop_visit_type, Visit, VisitWith};
 use crate::{
     config::{LintRuleReaction, RuleConfig},
     rule::{visitor_rule, Rule},
-    rules::utils::{extract_arg_val, unwrap_seqs_and_parens, ArgValue},
+    rules::utils::{extract_arg_val, ArgValue},
 };
 
 const UNEXPECTED_REG_EXP_MESSAGE: &str =
@@ -90,14 +90,14 @@ impl PreferRegexLiterals {
         if let Some(ExprOrSpread { expr, .. }) = args.first() {
             self.first_arg = Some(extract_arg_val(
                 self.unresolved_ctxt,
-                unwrap_seqs_and_parens(expr.as_ref()),
+                expr.unwrap_seqs_and_parens(),
             ));
         }
 
         if let Some(ExprOrSpread { expr, .. }) = args.get(1) {
             self.second_arg = Some(extract_arg_val(
                 self.unresolved_ctxt,
-                unwrap_seqs_and_parens(expr.as_ref()),
+                expr.unwrap_seqs_and_parens(),
             ));
         }
     }

diff --git a/crates/swc_ecma_lints/src/rules/radix.rs b/crates/swc_ecma_lints/src/rules/radix.rs
@@ -10,7 +10,7 @@ use swc_ecma_visit::{noop_visit_type, Visit, VisitWith};
 use crate::{
     config::{LintRuleReaction, RuleConfig},
     rule::{visitor_rule, Rule},
-    rules::utils::{extract_arg_val, unwrap_seqs_and_parens, ArgValue},
+    rules::utils::{extract_arg_val, ArgValue},
 };
 
 const OBJ_NAMES: &[&str] = &["Number", "globalThis"];
@@ -140,7 +140,7 @@ impl Radix {
         match call_expr.args.get(1) {
             Some(ExprOrSpread { expr, .. }) => {
                 let expr = if self.unwrap_parens_and_seqs {
-                    unwrap_seqs_and_parens(expr.as_ref())
+                    expr.unwrap_seqs_and_parens()
                 } else {
                     expr.as_ref()
                 };

diff --git a/crates/swc_ecma_lints/src/rules/symbol_description.rs b/crates/swc_ecma_lints/src/rules/symbol_description.rs
@@ -6,7 +6,7 @@ use swc_ecma_visit::{Visit, VisitWith};
 use crate::{
     config::{LintRuleReaction, RuleConfig},
     rule::{visitor_rule, Rule},
-    rules::utils::{extract_arg_val, unwrap_seqs_and_parens, ArgValue},
+    rules::utils::{extract_arg_val, ArgValue},
 };
 
 const SYMBOL_EXPECTED_MESSAGE: &str = "Expected Symbol to have a description";
@@ -62,7 +62,7 @@ impl SymbolDescription {
     fn check(&self, span: Span, first_arg: Option<&ExprOrSpread>) {
         if let Some(ExprOrSpread { expr, .. }) = first_arg {
             if self.enforce_string_description {
-                match extract_arg_val(self.unresolved_ctxt, unwrap_seqs_and_parens(expr)) {
+                match extract_arg_val(self.unresolved_ctxt, expr.unwrap_seqs_and_parens()) {
                     ArgValue::Str(_) => {}
                     _ => {
                         self.emit_report(span, SYMBOL_STRING_DESCRIPTION_EXPECTED_MESSAGE);

diff --git a/crates/swc_ecma_lints/src/rules/utils.rs b/crates/swc_ecma_lints/src/rules/utils.rs
@@ -1,9 +1,7 @@
 use serde::{Deserialize, Serialize};
 use swc_atoms::Atom;
 use swc_common::SyntaxContext;
-use swc_ecma_ast::{
-    Expr, Lit, MemberExpr, MemberProp, Number, ParenExpr, Regex, SeqExpr, Str, TaggedTpl, Tpl,
-};
+use swc_ecma_ast::{Expr, Lit, MemberExpr, MemberProp, Number, Regex, Str, TaggedTpl, Tpl};
 
 #[derive(Debug, Clone, Copy, Serialize, Deserialize)]
 #[serde(rename_all = "lowercase")]
@@ -93,11 +91,3 @@ pub fn extract_arg_val(unresolved_ctxt: SyntaxContext, expr: &Expr) -> ArgValue
         _ => ArgValue::Other,
     }
 }
-
-pub fn unwrap_seqs_and_parens(expr: &Expr) -> &Expr {
-    match expr {
-        Expr::Seq(SeqExpr { exprs, .. }) => unwrap_seqs_and_parens(exprs.last().unwrap()),
-        Expr::Paren(ParenExpr { expr, .. }) => unwrap_seqs_and_parens(expr.as_ref()),
-        _ => expr,
-    }
-}