From da5210d89bf65948aa2fb469239751e4f473df7f Mon Sep 17 00:00:00 2001
From: alanechang <alanechang@janestreet.com>
Date: Thu, 25 Jan 2024 18:02:47 -0500
Subject: [PATCH] flambda-backend: Fix layout annotation encoding to work
 better with ppxlib  (#2234)

* add jane prefix

* add test
---
 parsing/jane_syntax.ml                        |  5 ++-
 .../tests/parsetree/source_jane_street.ml     |  3 ++
 testsuite/tests/parsetree/test.ml             | 44 ++++++++++++++++---
 3 files changed, 44 insertions(+), 8 deletions(-)

diff --git a/parsing/jane_syntax.ml b/parsing/jane_syntax.ml
index 2a9a9c2ceaa..c0ed11c6ed5 100644
--- a/parsing/jane_syntax.ml
+++ b/parsing/jane_syntax.ml
@@ -211,7 +211,7 @@ module Make_payload_protocol_of_stringable (Stringable : Stringable) :
     let structure_item_of_none =
       { pstr_desc =
           Pstr_attribute
-            { attr_name = Location.mknoloc "none";
+            { attr_name = Location.mknoloc "jane.none";
               attr_payload = PStr [];
               attr_loc = Location.none
             };
@@ -278,7 +278,8 @@ module Make_payload_protocol_of_stringable (Stringable : Stringable) :
         | _ -> raise Unexpected
 
       let is_none_structure_item = function
-        | { pstr_desc = Pstr_attribute { attr_name = { txt = "none" } } } ->
+        | { pstr_desc = Pstr_attribute { attr_name = { txt = "jane.none" } } }
+          ->
           true
         | _ -> false
 
diff --git a/testsuite/tests/parsetree/source_jane_street.ml b/testsuite/tests/parsetree/source_jane_street.ml
index c9e611e660d..8f9001ab3a0 100644
--- a/testsuite/tests/parsetree/source_jane_street.ml
+++ b/testsuite/tests/parsetree/source_jane_street.ml
@@ -9,6 +9,7 @@
 let f (type a : immediate) (x : a) = x;;
 let f (type (a : immediate)) (x : a) = x;;
 let f (type (a : immediate) (b : immediate)) (x : a) = x;;
+let f (type (a : immediate) (b : immediate) c) (x : a) = x;;
 
 let f y (type a : immediate) (x : a) = x;;
 let f y (type (a : immediate)) (x : a) = x;;
@@ -17,6 +18,8 @@ let f y (type (a : immediate) (b : immediate)) (x : a) = x;;
 let f y (type a : immediate) = y;;
 let f y (type (a : immediate)) = y;;
 let f y (type (a : immediate) (b : immediate)) = y;;
+let f y (type (a : immediate) (b : immediate) c) = y;;
+
 
 (* Just newtypes, no value parameters *)
 let f (type a : immediate) (type b : immediate)
diff --git a/testsuite/tests/parsetree/test.ml b/testsuite/tests/parsetree/test.ml
index 84b31bbc841..ffe13415234 100644
--- a/testsuite/tests/parsetree/test.ml
+++ b/testsuite/tests/parsetree/test.ml
@@ -61,7 +61,7 @@ let test parse_fun pprint print map filename ~extra_checks =
   | ast ->
       let str = to_string pprint ast in
       begin
-        match extra_checks str with
+        match extra_checks (to_string print ast) str with
         | Ok () -> ()
         | Error reason ->
             Printf.printf "%s: FAIL, %s\n" filename reason;
@@ -112,7 +112,7 @@ let rec process path ~extra_checks =
       path
       ~extra_checks
 
-let process ?(extra_checks = fun _ -> Ok ()) text = process text ~extra_checks
+let process ?(extra_checks = fun _ _ -> Ok ()) text = process text ~extra_checks
 
 (* Produce an error if any attribute/extension node does not start with the
    text prefix.
@@ -128,7 +128,7 @@ let process ?(extra_checks = fun _ -> Ok ()) text = process text ~extra_checks
   We've chosen to keep those constructs out of the test file in preference
   to updating this logic to properly handle them (which is hard).
 *)
-let check_all_attributes_and_extensions_start_with text ~prefix =
+let check_all_printed_attributes_and_extensions_start_with text ~prefix =
   let check introduction_string =
     String.split_on_char '[' text
     |> List.for_all (fun s ->
@@ -146,14 +146,46 @@ let check_all_attributes_and_extensions_start_with text ~prefix =
            prefix)
 ;;
 
+let check_all_ast_attributes_and_extensions_start_with raw_parsetree_str ~prefixes =
+  (* Sadly can't use Ast_mapper here because it decodes Jane Syntax by default and
+     we will need quite a bit of code duplication for it to work for this use case. *)
+  let check introduction_string =
+    Misc.Stdlib.String.split_on_string ~split_on:(introduction_string ^ " \"")
+      raw_parsetree_str
+    |> List.tl
+    |> List.for_all (fun s ->
+          List.exists
+            (fun prefix -> String.starts_with s ~prefix)
+            prefixes)
+  in
+  if check "extension" && check "attribute"
+  then Ok ()
+  else
+      Error
+        (Printf.sprintf
+          "Printast produced an extension node or attribute that doesn't \
+           begin with one of [%s]"
+           (String.concat ", " prefixes))
+;;
+
 let () =
   process "source.ml";
   Language_extension.enable_maximal ();
-  process "source_jane_street.ml" ~extra_checks:(fun text ->
-  (* Check that printing Jane Street language extensions produces no more
+  process "source_jane_street.ml" ~extra_checks:(fun raw_parsetree_str text ->
+  (* Additionally check that:
+
+     1. Jane Street language extensions only use "extension." and "jane." prefixed
+     attributes and exntensions for its parsetree encoding. This is important for
+     ppx support.
+
+     2. Printing Jane Street language extensions produces no more
      attributes or extension nodes than the input program, all of whose
      attributes begin with "test". This ensures that Jane Syntax attributes
      aren't printed.
    *)
-    check_all_attributes_and_extensions_start_with text ~prefix:"test");
+    Result.bind
+      (check_all_ast_attributes_and_extensions_start_with raw_parsetree_str
+        ~prefixes:["extension."; "jane."; "test."])
+      (fun () -> check_all_printed_attributes_and_extensions_start_with text
+                    ~prefix:"test"));
 ;;