Transform arrow functions without block

src/tests/arrow_func_tests.rs

@@ -0,0 +1,75 @@
+/**
+* Unless explicitly stated otherwise all files in this repository are licensed under the Apache-2.0 License.
+* This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2022 Datadog, Inc.
+**/
+
+#[cfg(test)]
+mod tests {
+
+    use crate::{tests::rewrite_js, transform::transform_status::Status};
+    use speculoos::{assert_that, string::StrAssertions};
+
+    #[test]
+    fn test_arrow_simple_not_modified() -> Result<(), String> {
+        let original_code = "{const a = (arg) => arg ? true : false}".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
+                .transform_status
+                .is_some_and(|status| status.status == Status::NotModified),
+        );
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_arrow_paren_not_modified() -> Result<(), String> {
+        let original_code = "{const a = (arg) => (arg ? true : false)}".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
+                .transform_status
+                .is_some_and(|status| status.status == Status::NotModified),
+        );
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_arrow_block_not_modified() -> Result<(), String> {
+        let original_code = "{const a = (arg) => { return arg ? true : false }}".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
+                .transform_status
+                .is_some_and(|status| status.status == Status::NotModified),
+        );
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_arrow_modified() -> Result<(), String> {
+        let original_code = "{const a = (arg) => arg ? `hello ${arg}` : '' }".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
+                .transform_status
+                .is_some_and(|status| status.status == Status::Modified),
+        );
+        assert_that(&rewritten.code).contains("const a = (arg)=>{
+        let __datadog_test_0;
+        return arg ? (__datadog_test_0 = arg, _ddiast.tplOperator(`hello ${__datadog_test_0}`, __datadog_test_0)) : '';
+    };");
+
+        Ok(())
+    }
+}

src/tests/mod.rs

@@ -11,6 +11,7 @@ use crate::{
 use anyhow::Error;
 use std::path::PathBuf;
 
+mod arrow_func_tests;
 mod binary_assignation_test;
 mod binary_expression_test;
 mod literal_test;

src/tests/template_literal_test.rs

@@ -96,4 +96,49 @@ mod tests {
     var f = (__datadog_test_0 = yield 'yielded', _ddiast.tplOperator(`foo${__datadog_test_0}bar`, __datadog_test_0));");
         Ok(())
     }
+
+    #[test]
+    fn test_template_literal_with_arrow() -> Result<(), String> {
+        let original_code = "function names(arg) {
+            const flag = arg;
+            const addPrefix = (value) => (flag ? value : `\"my_prefix_${value}\"`);
+            const result = `${addPrefix('NAME_0')}`;
+            return result;
+    }"
+        .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;
+    const flag = arg;
+    const addPrefix = (value)=>{
+        let __datadog_test_0;",
+        );
+        Ok(())
+    }
+
+    #[test]
+    fn test_template_literal_with_paren() -> Result<(), String> {
+        let original_code = "function names(arg) {
+            const flag = arg;
+            const addPrefix = (value) => (flag ? value : `\"my_prefix_${value}\"`);
+            let a, b
+            const result = (a = 'NAME_0', b = 'NAME_1', `${addPrefix(a)}`);
+            return result;
+    }"
+        .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;
+    const flag = arg;
+    const addPrefix = (value)=>{
+        let __datadog_test_0;",
+        );
+        Ok(())
+    }
 }

src/transform/arrow_transform.rs

@@ -0,0 +1,33 @@
+/**
+ * Unless explicitly stated otherwise all files in this repository are licensed under the Apache-2.0 License.
+ * This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2022 Datadog, Inc.
+ **/
+use swc_common::{SyntaxContext, DUMMY_SP};
+use swc_ecma_ast::*;
+
+use super::transform_status::TransformResult;
+
+pub struct ArrowTransform {}
+
+impl ArrowTransform {
+    pub fn to_dd_arrow_expr(arrow: &ArrowExpr) -> TransformResult<Expr> {
+        if arrow.body.is_expr() {
+            let mut arrow_block = arrow.clone();
+
+            let return_stmt = ReturnStmt {
+                span: DUMMY_SP,
+                arg: Some(Box::new(*arrow_block.body.expr().unwrap())),
+            };
+
+            arrow_block.body = Box::new(BlockStmtOrExpr::BlockStmt(BlockStmt {
+                span: DUMMY_SP,
+                ctxt: SyntaxContext::empty(),
+                stmts: vec![Stmt::Return(return_stmt)],
+            }));
+
+            return TransformResult::modified(Expr::Arrow(arrow_block));
+        }
+
+        TransformResult::not_modified()
+    }
+}

src/transform/mod.rs

@@ -2,6 +2,7 @@
 * Unless explicitly stated otherwise all files in this repository are licensed under the Apache-2.0 License.
 * This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2022 Datadog, Inc.
 **/
+pub(crate) mod arrow_transform;
 pub(crate) mod assign_add_transform;
 pub(crate) mod binary_add_transform;
 pub(crate) mod call_expr_transform;

src/visitor/operation_transform_visitor.rs

@@ -9,6 +9,7 @@ use swc_ecma_visit::{Visit, VisitMut, VisitMutWith};
 use crate::{
     telemetry::Telemetry,
     transform::{
+        arrow_transform::ArrowTransform,
         assign_add_transform::AssignAddTransform,
         binary_add_transform::BinaryAddTransform,
         call_expr_transform::CallExprTransform,
@@ -162,6 +163,14 @@ impl VisitMut for OperationTransformVisitor<'_> {
                 }
             }
 
+            Expr::Arrow(arrow) => {
+                let transform_result = ArrowTransform::to_dd_arrow_expr(arrow);
+                if transform_result.is_modified() {
+                    expr.map_with_mut(|e| transform_result.expr.unwrap_or(e));
+                    // do not update status yet
+                }
+            }
+
             _ => {
                 expr.visit_mut_children_with(self);
             }

test/resources/issue-101.js

@@ -0,0 +1,112 @@
+'use strict'
+
+function names(arg) {
+  const flag = arg
+  const addPrefix = (value) => (flag ? `"${value}"` : `"my_prefix.${value}"`)
+  const result = `
+      ${addPrefix('NAME_0')}
+      ${addPrefix('NAME_1')}
+      ${addPrefix('NAME_2')}
+      ${addPrefix('NAME_3')}
+    `
+  return result
+}
+
+function namesBlock(arg) {
+  const flag = arg
+  const addPrefix = (value) => {
+    return flag ? `"${value}"` : `"my_prefix.${value}"`
+  }
+  const result = `
+      ${addPrefix('NAME_0')}
+      ${addPrefix('NAME_1')}
+      ${addPrefix('NAME_2')}
+      ${addPrefix('NAME_3')}
+    `
+  return result
+}
+
+function namesNested(arg) {
+  const flag = arg
+  const addSufix = (value) => `${!value}_suffix`
+  const addPrefix = (value) => (flag ? `"${value}"` : `"my_prefix.${addSufix(value)}"`)
+  const result = `
+      ${addPrefix('NAME_0')}
+      ${addPrefix('NAME_1')}
+      ${addPrefix('NAME_2')}
+      ${addPrefix('NAME_3')}
+    `
+  return result
+}
+
+function namesDoubleResult(arg) {
+  const flag = arg
+  const addPrefix = (value) => (flag ? `"${value}"` : `"my_prefix.${value}"`)
+  const resultFake = `
+      ${addPrefix('NAME_0')}
+      ${addPrefix('NAME_1')}
+      ${addPrefix('NAME_2')}
+      ${addPrefix('NAME_3')}
+    `
+  const result = `
+      ${addPrefix('NAME_0')}
+      ${addPrefix('NAME_1')}
+      ${addPrefix('NAME_2')}
+      ${addPrefix('NAME_3')}
+    `
+  return result
+}
+
+function namesNoFlag() {
+  const addPrefix = (value) => `"my_prefix.${value}"`
+  const result = `
+      ${addPrefix('NAME_0')}
+      ${addPrefix('NAME_1')}
+      ${addPrefix('NAME_2')}
+      ${addPrefix('NAME_3')}
+    `
+  return result
+}
+
+function namesNoFlagDoubleResult() {
+  const addPrefix = (value) => `"my_prefix.${value}"`
+  const resultFake = `
+      ${addPrefix('NAME_0')}
+      ${addPrefix('NAME_1')}
+      ${addPrefix('NAME_2')}
+      ${addPrefix('NAME_3')}
+    `
+  const result = `
+      ${addPrefix('NAME_0')}
+      ${addPrefix('NAME_1')}
+      ${addPrefix('NAME_2')}
+      ${addPrefix('NAME_3')}
+    `
+  return result
+}
+
+function paren(arg) {
+  const flag = arg
+  let a, b, c
+  const addPrefix = (value) => (flag ? `"${value}"` : `"my_prefix.${value}"`)
+  const result =
+    ((a = 'NAME_0'),
+    (b = 'NAME_1'),
+    (c = 'NAME_2'),
+    `
+    ${addPrefix(a)}
+    ${addPrefix(b)}
+    ${addPrefix(c)}
+  `)
+  return result
+}
+
+module.exports = {
+  names,
+  namesBlock,
+  namesNested,
+  namesDoubleResult,
+  namesNoFlag,
+  namesNoFlagDoubleResult,
+  paren,
+}

test/template_literal.spec.js

@@ -381,5 +381,31 @@ __datadog_test_0 + __datadog_test_1, __datadog_test_0, __datadog_test_1)), _ddia
       rewriteAndCompare(testFunction, [1, 'b'])
       rewriteAndCompare(testFunction, ['a', 2])
     })
+
+    it('issue 101', () => {
+      // eslint-disable-next-line no-unused-vars
+      const _ddiast = {
+        plusOperator: (res) => res,
+        tplOperator: (res) => res
+      }
+
+      const js = readFileSync(path.join(__dirname, 'resources/issue-101.js')).toString()
+
+      // eslint-disable-next-line no-eval
+      const issue101 = eval(js)
+
+      const rewritten = rewriteAst(js)
+
+      // eslint-disable-next-line no-eval
+      const rewrittenIssue101 = eval(rewritten)
+
+      expect(issue101.names(false)).to.be.equal(rewrittenIssue101.names(false))
+      expect(issue101.namesBlock(false)).to.be.equal(rewrittenIssue101.namesBlock(false))
+      expect(issue101.namesNested(false)).to.be.equal(rewrittenIssue101.namesNested(false))
+      expect(issue101.namesDoubleResult(false)).to.be.equal(rewrittenIssue101.namesDoubleResult(false))
+      expect(issue101.namesNoFlag(false)).to.be.equal(rewrittenIssue101.namesNoFlag(false))
+      expect(issue101.namesNoFlagDoubleResult(false)).to.be.equal(rewrittenIssue101.namesNoFlagDoubleResult(false))
+      expect(issue101.paren(false)).to.be.equal(rewrittenIssue101.paren(false))
+    })
   })
 })