fix: Correctly handle computed object assignments (#18)

Author: kamilogorek

src/swc.rs

@@ -340,6 +340,7 @@ fn prefix_getters_setters(kind: ast::MethodKind, scope_name: &mut ScopeName) {
 /// This is only possible if the expression is an identifier or a member expression.
 fn infer_name_from_expr(mut expr: &ast::Expr) -> Option<ScopeName> {
     let mut scope_name = ScopeName::new();
+
     loop {
         match expr {
             ast::Expr::Ident(ident) => {
@@ -356,6 +357,11 @@ fn infer_name_from_expr(mut expr: &ast::Expr) -> Option<ScopeName> {
                         .push_front(NameComponent::ident(ident.clone()));
                     scope_name.components.push_front(NameComponent::interp("."));
                 }
+
+                if let Some(computed_prop) = member.prop.as_computed() {
+                    push_computed_prop_name(computed_prop, &mut scope_name)
+                }
+
                 expr = &member.obj;
             }
 
@@ -371,6 +377,35 @@ fn infer_name_from_expr(mut expr: &ast::Expr) -> Option<ScopeName> {
     }
 }
 
+fn push_computed_prop_name(prop_name: &ast::ComputedPropName, scope_name: &mut ScopeName) {
+    if let Some(literal) = prop_name.expr.as_lit() {
+        let component = NameComponent::interp(format!("[{}]", lit_as_string(literal)));
+        scope_name.components.push_front(component);
+    } else if let Some(ident) = prop_name.expr.as_ident() {
+        scope_name.components.push_front(NameComponent::interp("]"));
+        scope_name
+            .components
+            .push_front(NameComponent::ident(ident.clone()));
+        scope_name.components.push_front(NameComponent::interp("["));
+    } else {
+        scope_name
+            .components
+            .push_front(NameComponent::interp("[<computed>]"));
+    }
+}
+
+fn lit_as_string(lit: &ast::Lit) -> String {
+    match lit {
+        ast::Lit::Str(v) => format!("\"{}\"", v.value),
+        ast::Lit::Num(v) => v.value.to_string(),
+        ast::Lit::BigInt(v) => format!("{}n", v.value),
+        ast::Lit::Bool(v) => v.value.to_string(),
+        ast::Lit::Regex(v) => format!("/{}/{}", v.exp, v.flags),
+        ast::Lit::Null(_) => String::from("null"),
+        ast::Lit::JSXText(v) => v.value.to_string(),
+    }
+}
+
 fn prop_name_to_component(prop: &ast::PropName) -> NameComponent {
     match prop {
         ast::PropName::Ident(ref i) => NameComponent::ident(i.clone()),

tests/extract.rs

@@ -218,7 +218,7 @@ fn extract_anon_obj_literal() {
 #[test]
 fn extract_empty_function() {
     let src = r#"
-         (function () {
+        (function () {
           return () => {};
         })()
         "#;
@@ -254,3 +254,23 @@ fn extract_nested_iife_objects() {
     ];
     assert_eq!(scopes, expected);
 }
+
+#[test]
+fn extract_computed_properties() {
+    let src = r#"         
+        Klass.prototype[42] = () => {}
+        Klass.prototype["method"] = () => {}
+        Klass.prototype[method] = () => {}
+        Klass.prototype[1 + 1] = () => {};
+        "#;
+    let scopes = extract_scope_names(src).unwrap();
+    let scopes = scope_strs(scopes);
+
+    let expected = [
+        Some("Klass.prototype[42]".into()),
+        Some("Klass.prototype[\"method\"]".into()),
+        Some("Klass.prototype[method]".into()),
+        Some("Klass.prototype[<computed>]".into()),
+    ];
+    assert_eq!(scopes, expected);
+}