Preserve raw extension payloads through ast0

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

Author: cknitt

compiler/ml/ast_mapper_from0.ml

@@ -102,6 +102,41 @@ let map_constant ~loc ~is_template = function
   | Pconst_string (s, q) -> Pconst_string (s, q)
   | Pconst_float (s, suffix) -> Pconst_float (s, suffix)
 
+let is_raw_source_extension = function
+  | "raw" | "ffi" | "re" -> true
+  | _ -> false
+
+let map_raw_source_payload sub = function
+  | PStr
+      [
+        {
+          pstr_desc =
+            Pstr_eval
+              ( {
+                  pexp_desc = Pexp_constant (Pconst_string (s, delim));
+                  pexp_loc;
+                  pexp_attributes;
+                },
+                eval_attributes );
+          pstr_loc;
+        };
+      ] ->
+    let expression =
+      Ast_helper.Exp.constant
+        ~loc:(sub.location sub pexp_loc)
+        ~attrs:(sub.attributes sub pexp_attributes)
+        (Pt.Pconst_string (s, delim))
+    in
+    Some
+      (Pt.PStr
+         [
+           Ast_helper.Str.eval
+             ~loc:(sub.location sub pstr_loc)
+             ~attrs:(sub.attributes sub eval_attributes)
+             expression;
+         ])
+  | _ -> None
+
 let for_of_attr_name = "_res.for_of"
 let for_await_of_attr_name = "_res.for_await_of"
 
@@ -1080,7 +1115,16 @@ let default_mapper =
           pc_rhs = this.expr this pc_rhs;
         });
     location = (fun _this l -> l);
-    extension = (fun this (s, e) -> (map_loc this s, this.payload this e));
+    extension =
+      (fun this (s, payload) ->
+        let payload =
+          if is_raw_source_extension s.txt then
+            match map_raw_source_payload this payload with
+            | Some payload -> payload
+            | None -> this.payload this payload
+          else this.payload this payload
+        in
+        (map_loc this s, payload));
     attribute = (fun this (s, e) -> (map_loc this s, this.payload this e));
     attributes = (fun this l -> List.map (this.attribute this) l);
     payload =

tests/ounit_tests/ounit_ast_mapper0_tests.ml

@@ -255,6 +255,45 @@ let test_string_literals_roundtrip_through_ast0 _ =
   assert_string_expr ~expected:{|{"answer":42}|} ~delim:(Some "json")
     (map_expr0 (map_expr_to0 json_expr))
 
+let assert_raw_extension_payload ~name ~expected expression =
+  match expression.Parsetree.pexp_desc with
+  | Pexp_extension
+      ( {txt},
+        PStr
+          [
+            {
+              pstr_desc =
+                Pstr_eval
+                  ( {pexp_desc = Pexp_constant (Pconst_string (actual, delim))},
+                    _ );
+            };
+          ] ) ->
+    OUnit.assert_equal name txt;
+    OUnit.assert_equal ~printer:(Printf.sprintf "%S") expected actual;
+    OUnit.assert_equal ~printer:Ext_obj.dump (Some "js") delim
+  | _ -> assert_failure "Expected a raw extension string payload"
+
+let test_raw_extension_payloads_roundtrip_through_ast0 _ =
+  let encoded = {|'\\n'|} in
+  List.iter
+    (fun name ->
+      let payload =
+        Parsetree0.PStr
+          [
+            Ast_helper0.Str.eval ~loc
+              (Ast_helper0.Exp.constant ~loc
+                 (Parsetree0.Pconst_string (encoded, Some "js")));
+          ]
+      in
+      let expression0 =
+        Ast_helper0.Exp.extension ~loc (Location.mknoloc name, payload)
+      in
+      let expression = map_expr0 expression0 in
+      assert_raw_extension_payload ~name ~expected:encoded expression;
+      assert_raw_extension_payload ~name ~expected:encoded
+        (map_expr0 (map_expr_to0 expression)))
+    ["raw"; "ffi"; "re"]
+
 (* Function-node attributes such as [@this] must stay node attributes across
    the v0 bridge: the built-in PPX reads decorators from [pexp_attributes],
    so a round trip that moves them into [p_attrs] silently disables them. *)
@@ -314,6 +353,8 @@ let suites =
          >:: test_ast0_strings_convert_to_internal_representation;
          "string_literals_roundtrip_through_ast0"
          >:: test_string_literals_roundtrip_through_ast0;
+         "raw_extension_payloads_roundtrip_through_ast0"
+         >:: test_raw_extension_payloads_roundtrip_through_ast0;
          "malformed_internal_record_rest_attr_fails"
          >:: test_malformed_internal_record_rest_attr_fails;
          "record_rest_roundtrips_through_ast0"

tests/tests/src/string_literal_normalization_test.mjs

@@ -15,14 +15,29 @@ function constantSwitch() {
   return 1;
 }
 
+const rawBridgeProgramValue = '\\n';
+;
+
+let rawBridgeExpression = '\\n';
+
+let rawBridgeFunction = (() => '\\n');
+
+let rawBridgeRegex = /\\n/;
+
 Mocha.describe("String_literal_normalization_test", () => {
-  Mocha.test("ordinary escapes have one semantic representation", () => Test_utils.eq("File \"string_literal_normalization_test.res\", line 20, characters 69-76", escaped, "abc"));
-  Mocha.test("surrogate-pair escapes have one semantic representation", () => Test_utils.eq("File \"string_literal_normalization_test.res\", line 23, characters 7-14", surrogatePair, "😀"));
+  Mocha.test("ordinary escapes have one semantic representation", () => Test_utils.eq("File \"string_literal_normalization_test.res\", line 28, characters 69-76", escaped, "abc"));
+  Mocha.test("surrogate-pair escapes have one semantic representation", () => Test_utils.eq("File \"string_literal_normalization_test.res\", line 31, characters 7-14", surrogatePair, "😀"));
   Mocha.test("ordinary literals participate in constant folding", () => {
-    Test_utils.eq("File \"string_literal_normalization_test.res\", line 27, characters 7-14", concatenated, "ab");
-    Test_utils.eq("File \"string_literal_normalization_test.res\", line 28, characters 7-14", constantSwitch(), 1);
+    Test_utils.eq("File \"string_literal_normalization_test.res\", line 35, characters 7-14", concatenated, "ab");
+    Test_utils.eq("File \"string_literal_normalization_test.res\", line 36, characters 7-14", constantSwitch(), 1);
+  });
+  Mocha.test("template segments survive the ast0 bridge", () => Test_utils.eq("File \"string_literal_normalization_test.res\", line 39, characters 61-68", interpolated, "abc"));
+  Mocha.test("raw extension payloads preserve source spelling through ast0", () => {
+    Test_utils.eq("File \"string_literal_normalization_test.res\", line 42, characters 7-14", rawBridgeExpression, "\\n");
+    Test_utils.eq("File \"string_literal_normalization_test.res\", line 43, characters 7-14", rawBridgeFunction(), "\\n");
+    Test_utils.eq("File \"string_literal_normalization_test.res\", line 44, characters 7-14", rawBridgeProgramValue, "\\n");
+    Test_utils.eq("File \"string_literal_normalization_test.res\", line 45, characters 7-14", rawBridgeRegex.test("\\n"), true);
   });
-  Mocha.test("template segments survive the ast0 bridge", () => Test_utils.eq("File \"string_literal_normalization_test.res\", line 31, characters 61-68", interpolated, "abc"));
 });
 
 export {
@@ -31,5 +46,8 @@ export {
   concatenated,
   interpolated,
   constantSwitch,
+  rawBridgeExpression,
+  rawBridgeFunction,
+  rawBridgeRegex,
 }
 /*  Not a pure module */

tests/tests/src/string_literal_normalization_test.res

@@ -16,6 +16,14 @@ let constantSwitch = () =>
   | _ => 4
   }
 
+%%raw("const rawBridgeProgramValue = '\\n';")
+
+@val external rawBridgeProgramValue: string = "rawBridgeProgramValue"
+
+let rawBridgeExpression: string = %raw("'\\n'")
+let rawBridgeFunction: unit => string = %ffi("() => '\\n'")
+let rawBridgeRegex = /\\n/
+
 describe(__MODULE__, () => {
   test("ordinary escapes have one semantic representation", () => eq(__LOC__, escaped, "abc"))
 
@@ -29,4 +37,11 @@ describe(__MODULE__, () => {
   })
 
   test("template segments survive the ast0 bridge", () => eq(__LOC__, interpolated, "abc"))
+
+  test("raw extension payloads preserve source spelling through ast0", () => {
+    eq(__LOC__, rawBridgeExpression, "\\n")
+    eq(__LOC__, rawBridgeFunction(), "\\n")
+    eq(__LOC__, rawBridgeProgramValue, "\\n")
+    eq(__LOC__, rawBridgeRegex->RegExp.test("\\n"), true)
+  })
 })