Normalize PPX-rewritten template segments

Signed-off-by: Christoph Knittel <ck@cca.io>

Author: cknitt

compiler/ml/ast_mapper_from0.ml

@@ -122,6 +122,11 @@ let semantic_string semantic =
 let source_string ~loc source =
   Pt.Pconst_string {source; semantic = decode_js_string ~loc source}
 
+let template_source_from0 = function
+  | source, Some ("js" | "*j") -> source
+  | semantic, _ ->
+    String_literal.encode_js_template (normalize_ppx_semantic_string semantic)
+
 let map_constant ~loc = function
   | Pconst_integer (s, suffix) -> Pt.Pconst_integer (s, suffix)
   | Pconst_char semantic ->
@@ -600,9 +605,10 @@ module E = struct
       let inner = sub.expr sub {e with pexp_attributes = inner_attrs0} in
       await ~loc ~attrs:(sub.attributes sub await_attrs0) inner
     | Pexp_ident x -> ident ~loc ~attrs (map_loc sub x)
-    | Pexp_constant (Pconst_string (source, Some "js"))
-      when has_template_attr attrs ->
+    | Pexp_constant (Pconst_string (text, delimiter))
+      when has_template_attr attrs && delimiter <> Some "json" ->
       let attrs = remove_template_attr attrs in
+      let source = template_source_from0 (text, delimiter) in
       template ~loc ~attrs [{txt = source; loc}] []
     | Pexp_constant x ->
       let template = has_template_attr attrs in
@@ -751,26 +757,31 @@ module E = struct
           "`json` literals do not support interpolation"
       in
       let rec collect sources values = function
+        | [{pexp_desc = Pexp_constant (Pconst_string (_, Some "json"))}] ->
+          reject_json_interpolation ()
         | [
             {
-              pexp_desc = Pexp_constant (Pconst_string (txt, Some "js"));
+              pexp_desc = Pexp_constant (Pconst_string (text, delimiter));
               pexp_loc;
             };
           ] ->
+          let txt = template_source_from0 (text, delimiter) in
           Some
             ( List.rev
                 ({Location.txt; loc = sub.location sub pexp_loc} :: sources),
               List.rev values )
-        | [{pexp_desc = Pexp_constant (Pconst_string (_, Some "json"))}] ->
+        | {pexp_desc = Pexp_constant (Pconst_string (_, Some "json"))}
+          :: _value :: _rest ->
           reject_json_interpolation ()
-        | {pexp_desc = Pexp_constant (Pconst_string (txt, Some "js")); pexp_loc}
+        | {
+            pexp_desc = Pexp_constant (Pconst_string (text, delimiter));
+            pexp_loc;
+          }
           :: value :: rest ->
+          let txt = template_source_from0 (text, delimiter) in
           collect
             ({Location.txt; loc = sub.location sub pexp_loc} :: sources)
             (value :: values) rest
-        | {pexp_desc = Pexp_constant (Pconst_string (_, Some "json"))}
-          :: _value :: _rest ->
-          reject_json_interpolation ()
         | _ -> None
       in
       begin match collect [] [] parts with

tests/ounit_tests/ounit_ast_mapper0_tests.ml

@@ -444,6 +444,39 @@ let test_ppx_rewritten_tagged_template_segments _ =
     OUnit.assert_equal ~printer:(Printf.sprintf "%S") {e|\${value}\`\\|e} txt
   | _ -> assert_failure "Expected a rewritten tagged template after roundtrip"
 
+let test_ppx_rewritten_template_segments _ =
+  let template_attr = attr "res.template" (Parsetree0.PStr []) in
+  let semantic = "${value}`\\" in
+  let uninterpolated =
+    Ast_helper0.Exp.constant ~loc ~attrs:[template_attr]
+      (Ast_helper0.Const.string semantic)
+  in
+  assert_template_expr ~expected:{e|\${value}\`\\|e} (map_expr0 uninterpolated);
+  let segment semantic =
+    Ast_helper0.Exp.constant ~loc ~attrs:[template_attr]
+      (Ast_helper0.Const.string semantic)
+  in
+  let concat lhs rhs =
+    Ast_helper0.Exp.apply ~loc ~attrs:[template_attr]
+      (Ast_helper0.Exp.ident ~loc (Location.mknoloc (Longident.Lident "^")))
+      [(Asttypes.Noloc.Nolabel, lhs); (Asttypes.Noloc.Nolabel, rhs)]
+  in
+  let value =
+    Ast_helper0.Exp.ident ~loc (Location.mknoloc (Longident.Lident "value"))
+  in
+  let interpolated =
+    concat (concat (segment "${head}") value) (segment "`\\")
+  in
+  match (map_expr0 interpolated).pexp_desc with
+  | Pexp_template
+      {
+        source_segments = [{txt = head}; {txt = tail}];
+        values = [{pexp_desc = Pexp_ident {txt = Longident.Lident "value"}}];
+      } ->
+    OUnit.assert_equal ~printer:(Printf.sprintf "%S") {e|\${head}|e} head;
+    OUnit.assert_equal ~printer:(Printf.sprintf "%S") {e|\`\\|e} tail
+  | _ -> assert_failure "Expected a rewritten template after roundtrip"
+
 let test_interpolated_templates_roundtrip_through_ast0 _ =
   let head_loc = source_loc 1 8 in
   let tail_loc = source_loc 12 16 in
@@ -646,6 +679,8 @@ let suites =
          >:: test_tagged_templates_roundtrip_through_ast0;
          "ppx_rewritten_tagged_template_segments"
          >:: test_ppx_rewritten_tagged_template_segments;
+         "ppx_rewritten_template_segments"
+         >:: test_ppx_rewritten_template_segments;
          "interpolated_templates_roundtrip_through_ast0"
          >:: test_interpolated_templates_roundtrip_through_ast0;
          "ast0_json_interpolation_is_rejected"