fix: define env (#1551)

* fix: define env

* test: update test cases for env_replace

* fix: typos

crates/mako/src/visitors/env_replacer.rs

@@ -49,42 +49,49 @@ impl VisitMut for EnvReplacer {
                 let mut member_visit_path = match prop {
                     MemberProp::Ident(Ident { sym, .. }) => sym.to_string(),
                     MemberProp::Computed(ComputedPropName {
-                        expr: expr_computed,
-                        ..
-                    }) => match expr_computed.as_ref() {
+                        expr: expr_compute, ..
+                    }) => match expr_compute.as_ref() {
                         Expr::Lit(Lit::Str(Str { value, .. })) => value.to_string(),
 
                         Expr::Lit(Lit::Num(Number { value, .. })) => value.to_string(),
-                        _ => return,
+                        _ => {
+                            obj.visit_mut_with(self);
+                            expr_compute.visit_mut_with(self);
+                            return;
+                        }
                     },
                     _ => return,
                 };
 
-                let mut current_member_obj = obj.as_ref();
+                let mut current_member_obj = obj.as_mut();
 
                 while let Expr::Member(MemberExpr { obj, prop, .. }) = current_member_obj {
                     match prop {
                         MemberProp::Ident(Ident { sym, .. }) => {
                             member_visit_path.push('.');
                             member_visit_path.push_str(sym.as_ref());
                         }
-                        MemberProp::Computed(ComputedPropName { expr, .. }) => {
-                            match expr.as_ref() {
-                                Expr::Lit(Lit::Str(Str { value, .. })) => {
-                                    member_visit_path.push('.');
-                                    member_visit_path.push_str(value.as_ref());
-                                }
-
-                                Expr::Lit(Lit::Num(Number { value, .. })) => {
-                                    member_visit_path.push('.');
-                                    member_visit_path.push_str(&value.to_string());
-                                }
-                                _ => return,
+                        MemberProp::Computed(ComputedPropName {
+                            expr: expr_compute, ..
+                        }) => match expr_compute.as_ref() {
+                            Expr::Lit(Lit::Str(Str { value, .. })) => {
+                                member_visit_path.push('.');
+                                member_visit_path.push_str(value.as_ref());
                             }
-                        }
+
+                            Expr::Lit(Lit::Num(Number { value, .. })) => {
+                                member_visit_path.push('.');
+                                member_visit_path.push_str(&value.to_string());
+                            }
+                            _ => {
+                                obj.visit_mut_with(self);
+                                expr_compute.visit_mut_with(self);
+                                return;
+                            }
+                        },
                         _ => return,
                     }
-                    current_member_obj = obj.as_ref();
+                    current_member_obj = obj.as_mut();
                 }
 
                 if let Expr::Ident(Ident { sym, .. }) = current_member_obj {
@@ -216,6 +223,7 @@ mod tests {
     use std::sync::Arc;
 
     use maplit::hashmap;
+    use regex::Regex;
     use serde_json::{json, Value};
     use swc_core::common::GLOBALS;
     use swc_core::ecma::visit::VisitMutWith;
@@ -260,7 +268,7 @@ mod tests {
                     "A".to_string() => json!(true)
                 }
             ),
-            r#"log(true);"#
+            "log(true);"
         );
     }
 
@@ -273,7 +281,7 @@ mod tests {
                     "A".to_string() => json!(1)
                 }
             ),
-            r#"log(1);"#
+            "log(1);"
         );
     }
 
@@ -286,7 +294,7 @@ mod tests {
                     "A".to_string() => json!("\"foo\"")
                 }
             ),
-            r#"log("foo");"#
+            "log(\"foo\");"
         );
     }
 
@@ -296,17 +304,10 @@ mod tests {
             run(
                 r#"log(A)"#,
                 hashmap! {
-                    "A".to_string() => json!([1, true, "\"foo\""])
+                    "A".to_string() => json!([1,true,"\"foo\""])
                 }
             ),
-            r#"
-log([
-    1,
-    true,
-    "foo"
-]);
-            "#
-            .trim()
+            "log([1,true,\"foo\"]);".trim()
         );
     }
 
@@ -319,9 +320,7 @@ log([
                     "A".to_string() => json!("{\"v\": 1}")
                 }
             ),
-            r#"log(({
-    "v": 1
-}));"#
+            "log(({\"v\": 1}));"
         );
     }
 
@@ -334,7 +333,7 @@ log([
                     "x.y".to_string() => json!(true)
                 }
             ),
-            r#"log(true);"#
+            "log(true);"
         );
     }
 
@@ -347,7 +346,7 @@ log([
                     "process.env.A".to_string() => json!(true)
                 }
             ),
-            r#"log(true);"#
+            "log(true);"
         );
     }
 
@@ -360,7 +359,7 @@ log([
                     "A.B".to_string() => json!(1)
                 }
             ),
-            r#"log(1);"#
+            "log(1);"
         );
     }
 
@@ -373,20 +372,98 @@ log([
                     "A.1".to_string() => json!(1)
                 }
             ),
-            r#"log(1);"#
+            "log(1);"
         );
     }
 
     #[test]
-    fn test_complicated() {
+    fn test_computed_after_ident() {
         assert_eq!(
             run(
                 r#"log(A.v["v"])"#,
                 hashmap! {
                     "A.v.v".to_string() => json!(1)
                 }
             ),
-            r#"log(1);"#
+            "log(1);"
+        );
+    }
+
+    #[test]
+    fn test_computed_as_member_key() {
+        assert_eq!(
+            run(
+                r#"log(A[v])"#,
+                hashmap! {
+                    "v".to_string() => json!(1)
+                }
+            ),
+            "log(A[1]);"
+        );
+    }
+
+    #[test]
+    fn test_complited_computed_as_member_key() {
+        assert_eq!(
+            run(
+                r#"log(A[v.v])"#,
+                hashmap! {
+                    "v.v".to_string() => json!(1)
+                }
+            ),
+            "log(A[1]);"
+        );
+    }
+
+    #[test]
+    fn test_computed_nested() {
+        assert_eq!(
+            run(
+                r#"log(A[v].B[v][v].C[v][v][v])"#,
+                hashmap! {
+                    "v".to_string() => json!(1)
+                }
+            ),
+            "log(A[1].B[1][1].C[1][1][1]);"
+        );
+    }
+
+    #[test]
+    fn test_should_not_replace_existed() {
+        assert_eq!(
+            run(
+                r#"let v = 2;log(A[v])"#,
+                hashmap! {
+                    "v".to_string() => json!(1)
+                }
+            ),
+            "let v = 2;log(A[v]);"
+        );
+    }
+
+    #[test]
+    fn test_should_not_replace_not_defined() {
+        assert_eq!(
+            run(
+                r#"log(A[B].v)"#,
+                hashmap! {
+                    "A.B".to_string() => json!("{\"v\": 1}")
+                }
+            ),
+            "log(A[B].v);"
+        );
+    }
+
+    #[test]
+    fn test_should_not_replace_as_member_ident() {
+        assert_eq!(
+            run(
+                r#"log(A.v)"#,
+                hashmap! {
+                    "v".to_string() => json!(1)
+                }
+            ),
+            "log(A.v);"
         );
     }
 
@@ -398,6 +475,10 @@ log([
             let mut visitor = EnvReplacer::new(envs, ast.unresolved_mark);
             ast.ast.visit_mut_with(&mut visitor);
         });
-        test_utils.js_ast_to_code()
+        let code = test_utils.js_ast_to_code();
+        Regex::new(r"\s*\n\s*")
+            .unwrap()
+            .replace_all(&code, "")
+            .to_string()
     }
 }