Rewrite some expressions like String.prototype[concat|replace].[call|…

…apply]("literal", ...args) (#93)

* Rewrite some String.prototype[concat|replace|].call expressions

* Add a some more v8 tests

* Clean up

* suggested test case

* Use Expr provided method

* Do not use a generated ident for literals (#95)

docker/v8test.Dockerfile

@@ -14,7 +14,7 @@ COPY ./scripts/crawler.js /build/package/scripts/crawler.js
 RUN npm i --prefix package
 
 # rewrite v8 mjsunit test files
-RUN node package/scripts/crawler.js --override v8/test/mjsunit/ '^(str|arr|arg|num|rege|mod|glob|obj|val|whit|this|throw|try|unbox|pro|call|code|comp|func|for|field).*'
+RUN node package/scripts/crawler.js --override v8/test/mjsunit/ '^(str|arr|arg|num|rege|mod|glob|obj|val|whit|this|throw|try|unbox|pro|call|code|comp|func|for|field|substr|unicode).*'
 
 # launch mjsunit tests
 CMD ["/build/v8/tools/run-tests.py", "--outdir=out/x64.release", "mjsunit"]

src/tests/string_method_test.rs

@@ -109,6 +109,26 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn test_prototype_concat_call() -> Result<(), String> {
+        let original_code = "{const a = String.prototype.concat.call(1,a(),3)}".to_string();
+        let js_file = "test.js".to_string();
+        let rewritten = rewrite_js(original_code, js_file).map_err(|e| e.to_string())?;
+        assert_that(&rewritten.code)
+            .contains("let __datadog_test_0, __datadog_test_1;
+    const a = (__datadog_test_0 = String.prototype.concat, __datadog_test_1 = a(), _ddiast.stringConcat(__datadog_test_0.call(1, __datadog_test_1, 3), __datadog_test_0, 1, __datadog_test_1, 3));");
+        Ok(())
+    }
+
+    #[test]
+    fn test_prototype_concat_call_literals() -> Result<(), String> {
+        let original_code = "{const a = String.prototype.concat.call(1,2,3)}".to_string();
+        let js_file = "test.js".to_string();
+        let rewritten = rewrite_js(original_code, js_file).map_err(|e| e.to_string())?;
+        assert_that(&rewritten.code).contains("const a = String.prototype.concat.call(1,2,3)");
+        Ok(())
+    }
+
     #[test]
     fn test_ident_trim() -> Result<(), String> {
         let original_code = "{const a = b.trim();}".to_string();

src/transform/call_expr_transform.rs

@@ -145,7 +145,10 @@ fn replace_prototype_call_or_apply(
     ident_provider: &mut dyn IdentProvider,
 ) -> Option<ResultExpr> {
     let prototype_call_option = FunctionPrototypeTransform::get_expression_parts_from_call_or_apply(
-        call, member, ident_name,
+        call,
+        member,
+        ident_name,
+        csi_methods,
     );
 
     match prototype_call_option {
@@ -252,9 +255,11 @@ fn replace_call_expr_if_csi_method_with_member(
         let mut call_replacement = call.clone();
 
         // __datadog_token_$i = a
-        let ident_replacement =
+        let ident_replacement_option =
             ident_provider.get_temporal_ident_used_in_assignation(expr, &mut assignations, &span);
 
+        let ident_replacement = ident_replacement_option.map_or_else(|| expr.clone(), Expr::Ident);
+
         let ident_callee = match member_expr_opt {
             Some(member_expr) => {
                 // __datadog_token_$i2 = member
@@ -269,7 +274,7 @@ fn replace_call_expr_if_csi_method_with_member(
                 // __datadog_token_$i.substring
                 let member_expr = MemberExpr {
                     span,
-                    obj: Box::new(Expr::Ident(ident_replacement.clone())),
+                    obj: Box::new(ident_replacement.clone()),
                     prop: MemberProp::Ident(ident_name.clone()),
                 };
 
@@ -283,12 +288,12 @@ fn replace_call_expr_if_csi_method_with_member(
             }
         };
 
-        arguments.push(Expr::Ident(ident_replacement.clone()));
+        arguments.push(ident_replacement.clone());
 
         // change callee to __datadog_token_$i2.call
         call_replacement.callee = Callee::Expr(Box::new(Expr::Member(MemberExpr {
             span,
-            obj: Box::new(Expr::Ident(ident_callee)),
+            obj: Box::new(ident_callee.map_or_else(|| expr.clone(), Expr::Ident)),
             prop: MemberProp::Ident(IdentName::new(JsWord::from("call"), span)),
         })));
 
@@ -308,7 +313,7 @@ fn replace_call_expr_if_csi_method_with_member(
             0,
             ExprOrSpread {
                 spread: None,
-                expr: Box::new(Expr::Ident(ident_replacement)),
+                expr: Box::new(ident_replacement),
             },
         );
 

src/transform/function_prototype_transform.rs

@@ -8,6 +8,8 @@ use swc_ecma_visit::swc_ecma_ast::{
     CallExpr, Callee, Expr, ExprOrSpread, Ident, MemberExpr, MemberProp,
 };
 
+use crate::visitor::csi_methods::CsiMethods;
+
 pub const PROTOTYPE: &str = "prototype";
 pub const CALL_METHOD_NAME: &str = "call";
 pub const APPLY_METHOD_NAME: &str = "apply";
@@ -35,6 +37,7 @@ impl FunctionPrototypeTransform {
         call: &CallExpr,
         member: &MemberExpr,
         ident_name: &IdentName,
+        csi_methods: &CsiMethods,
     ) -> Option<(Expr, Ident, CallExpr)> {
         if !Self::is_call_or_apply(ident_name) {
             return None;
@@ -47,11 +50,6 @@ impl FunctionPrototypeTransform {
                 return None;
             }
 
-            let this_expr = &call.args[0].expr;
-            if this_expr.is_lit() {
-                return None;
-            }
-
             let mut filtered_args = vec![];
             if !filter_call_args(
                 &call.args,
@@ -62,6 +60,14 @@ impl FunctionPrototypeTransform {
             }
 
             let method_ident = path_parts[0].clone();
+            let this_expr = &call.args[0].expr;
+
+            if this_expr.is_lit()
+                && (!csi_methods.method_allows_literal_callers(&method_ident.sym)
+                    || all_args_are_literal(&filtered_args))
+            {
+                return None;
+            }
 
             let new_callee = MemberExpr {
                 obj: this_expr.clone(),
@@ -130,3 +136,12 @@ fn get_prototype_member_path(member: &MemberExpr, parts: &mut Vec<Ident>) -> boo
     }
     !parts.is_empty()
 }
+
+fn all_args_are_literal(args: &[ExprOrSpread]) -> bool {
+    args.iter()
+        .all(|elem| elem.expr.is_lit() || is_undefined_or_null(&elem.expr))
+}
+
+fn is_undefined_or_null(expr: &Expr) -> bool {
+    expr.is_ident_ref_to("undefined") || expr.is_ident_ref_to("null")
+}

src/transform/operand_handler.rs

@@ -27,12 +27,14 @@ pub trait OperandHandler {
             Expr::Ident(_) => {
                 if ident_mode == IdentMode::Replace {
                     operand.map_with_mut(|op| {
-                        Expr::Ident(ident_provider.get_ident_used_in_assignation(
+                        let ident = ident_provider.get_ident_used_in_assignation(
                             &op,
                             assignations,
                             arguments,
                             span,
-                        ))
+                        );
+
+                        ident.map_or(op, Expr::Ident)
                     })
                 } else {
                     arguments.push(operand.clone())
@@ -70,12 +72,10 @@ pub trait OperandHandler {
         ident_provider: &mut dyn IdentProvider,
     ) {
         operand.map_with_mut(|op| {
-            Expr::Ident(ident_provider.get_ident_used_in_assignation(
-                &op,
-                assignations,
-                arguments,
-                span,
-            ))
+            let ident =
+                ident_provider.get_ident_used_in_assignation(&op, assignations, arguments, span);
+
+            ident.map_or(op, Expr::Ident)
         })
     }
 

src/visitor/ident_provider.rs

@@ -18,11 +18,14 @@ pub trait IdentProvider {
         assignations: &mut Vec<Expr>,
         arguments: &mut Vec<Expr>,
         span: &Span,
-    ) -> Ident {
+    ) -> Option<Ident> {
         let id = self.get_temporal_ident_used_in_assignation(operand, assignations, span);
 
         // store ident as argument
-        arguments.push(Expr::Ident(id.clone()));
+        arguments.push(
+            id.as_ref()
+                .map_or_else(|| operand.clone(), |ident| Expr::Ident(ident.clone())),
+        );
 
         id
     }
@@ -32,7 +35,11 @@ pub trait IdentProvider {
         operand: &Expr,
         assignations: &mut Vec<Expr>,
         span: &Span,
-    ) -> Ident {
+    ) -> Option<Ident> {
+        if operand.is_lit() {
+            return None;
+        }
+
         let next_ident = self.next_ident();
         let (assign, id) = self.create_assign_expression(next_ident, operand, span);
 
@@ -41,7 +48,7 @@ pub trait IdentProvider {
 
         assignations.push(Expr::Assign(assign));
 
-        id
+        Some(id)
     }
 
     fn create_assign_expression(

test/string_method.spec.js

@@ -163,6 +163,11 @@ _ddiast.stringSubstring(__datadog_test_1.call(__datadog_test_0, 2), __datadog_te
       rewriteAndExpectNoTransformation(js)
     })
 
+    it('does not modify String.prototype.substring direct call', () => {
+      const js = 'String.prototype.substring.call("abc", 0, a);'
+      rewriteAndExpectNoTransformation(js)
+    })
+
     it('does modify member.prop.substring call', () => {
       const js = 'a.b.c.substring(1);'
       rewriteAndExpect(
@@ -200,6 +205,61 @@ _ddiast.stringSubstring(__datadog_test_1.call(__datadog_test_0, 2), __datadog_te
     })
   })
 
+  describe('concat (method that allows literals)', () => {
+    it('does not modify String.prototype.concat call if all args are literals', () => {
+      const js = 'String.prototype.concat.call("hello", "world");'
+      rewriteAndExpectNoTransformation(js)
+    })
+
+    it('does modify String.prototype.concat call if args are literals and null', () => {
+      const js = 'String.prototype.concat.call("hello", 2, undefined, null);'
+      rewriteAndExpectNoTransformation(js)
+    })
+
+    it('does modify String.prototype.concat call if some ident', () => {
+      const js = 'String.prototype.concat.call("hello", a, "world");'
+      rewriteAndExpect(
+        js,
+        `{
+  let __datadog_test_0, __datadog_test_1;
+(__datadog_test_0 = String.prototype.concat, __datadog_test_1 = a, _ddiast.concat(\
+__datadog_test_0.call("hello", __datadog_test_1, "world"), __datadog_test_0, "hello", \
+__datadog_test_1, "world"));
+      }`
+      )
+    })
+
+    it('does not modify String.prototype.concat apply if all args are literals', () => {
+      const js = 'String.prototype.concat.apply("hello", ["world", null, "moon"]);'
+      rewriteAndExpectNoTransformation(js)
+    })
+
+    it('does modify String.prototype.concat apply if this is not a literal', () => {
+      const js = 'String.prototype.concat.apply(a, ["world", null]);'
+      rewriteAndExpect(
+        js,
+        `{
+  let __datadog_test_0, __datadog_test_1;
+(__datadog_test_0 = a, __datadog_test_1 = String.prototype.concat, _ddiast.concat(__datadog_test_1.call(\
+__datadog_test_0, "world", null), __datadog_test_1, __datadog_test_0, "world", null));
+      }`
+      )
+    })
+
+    it('does modify String.prototype.concat apply if an argument is not a literal', () => {
+      const js = 'String.prototype.concat.apply("hello", ["world", a]);'
+      rewriteAndExpect(
+        js,
+        `{
+  let __datadog_test_0, __datadog_test_1;
+(__datadog_test_0 = String.prototype.concat, __datadog_test_1 = a, \
+_ddiast.concat(__datadog_test_0.call("hello", "world", __datadog_test_1), __datadog_test_0, \
+"hello", "world", __datadog_test_1));
+      }`
+      )
+    })
+  })
+
   const methodAllowingLiterals = ['concat', 'replace']
 
   itEach(
@@ -239,9 +299,9 @@ _ddiast.stringSubstring(__datadog_test_1.call(__datadog_test_0, 2), __datadog_te
             const js = builder.build(`return 'a'.${method}(${args});`)
             rewriteAndExpectAndExpectEval(
               js,
-              builder.build(`let __datadog_test_0, __datadog_test_1;
-        return (__datadog_test_0 = 'a', __datadog_test_1 = __datadog_test_0.${method}, _ddiast.${method}(\
-__datadog_test_1.call(__datadog_test_0${argsWithComma}), __datadog_test_1, __datadog_test_0${argsWithComma}));`)
+              builder.build(`let __datadog_test_0;
+        return (__datadog_test_0 = 'a'.${method}, _ddiast.${method}(\
+__datadog_test_0.call('a'${argsWithComma}), __datadog_test_0, 'a'${argsWithComma}));`)
             )
           })
         } else {
@@ -284,10 +344,23 @@ __datadog_test_1.call(__datadog_test_0${argsWithComma}), __datadog_test_1, __dat
           )
         })
 
-        it(`does not modify literal String.prototype.${value}.call`, () => {
-          const js = `String.prototype.${method}.call("hello"${argsWithComma});`
-          rewriteAndExpectNoTransformation(js)
-        })
+        if (methodAllowingLiterals.indexOf(method) !== -1) {
+          it(`does modify literal String.prototype.${value}.call`, () => {
+            const builder = fn().args({ a: 'heLLo' })
+            const js = builder.build(`String.prototype.${method}.call(a${argsWithComma});`)
+            rewriteAndExpectAndExpectEval(
+              js,
+              builder.build(`let __datadog_test_0, __datadog_test_1;
+        (__datadog_test_0 = a, __datadog_test_1 = String.prototype.${method}, _ddiast.${method}(\
+__datadog_test_1.call(__datadog_test_0${argsWithComma}), __datadog_test_1, __datadog_test_0${argsWithComma}));`)
+            )
+          })
+        } else {
+          it(`does not modify literal String.prototype.${value}.call`, () => {
+            const js = `String.prototype.${method}.call("hello"${argsWithComma});`
+            rewriteAndExpectNoTransformation(js)
+          })
+        }
 
         it(`does modify String.prototype.${method}.call`, () => {
           const builder = fn().args('heLLo')