flambda-backend: Fix ast iteration/mapping for layout type declarations (#2145)

* Demonstrate bug in using ppxes with certain jane syntax elements

* Fix bug

Author: ncik-roberts

parsing/ast_iterator.ml

@@ -178,12 +178,19 @@ module T = struct
     | Ptyp_extension x -> sub.extension sub x
 
   let iter_type_declaration sub
-      {ptype_name; ptype_params; ptype_cstrs;
+     ({ptype_name; ptype_params; ptype_cstrs;
        ptype_kind;
        ptype_private = _;
        ptype_manifest;
        ptype_attributes;
-       ptype_loc} =
+       ptype_loc} as ty_decl) =
+    let ptype_attributes =
+      match Jane_syntax.Layouts.of_type_declaration ty_decl with
+      | Some (jkind, attrs) ->
+          iter_loc_txt sub sub.jkind_annotation jkind;
+          attrs
+      | None -> ptype_attributes
+    in
     iter_loc sub ptype_name;
     List.iter (iter_fst (sub.typ sub)) ptype_params;
     List.iter

parsing/ast_mapper.ml

@@ -217,22 +217,32 @@ module T = struct
     | Ptyp_extension x -> extension ~loc ~attrs (sub.extension sub x)
 
   let map_type_declaration sub
-      {ptype_name; ptype_params; ptype_cstrs;
+     ({ptype_name; ptype_params; ptype_cstrs;
        ptype_kind;
        ptype_private;
        ptype_manifest;
        ptype_attributes;
-       ptype_loc} =
+       ptype_loc} as tyd) =
     let loc = sub.location sub ptype_loc in
+    let jkind, ptype_attributes =
+      match Jane_syntax.Layouts.of_type_declaration tyd with
+      | None -> None, ptype_attributes
+      | Some (jkind, attributes) ->
+          let jkind = map_loc_txt sub sub.jkind_annotation jkind in
+          Some jkind, attributes
+    in
     let attrs = sub.attributes sub ptype_attributes in
-    Type.mk ~loc ~attrs (map_loc sub ptype_name)
+    Jane_syntax.Layouts.type_declaration_of ~loc ~attrs (map_loc sub ptype_name)
       ~params:(List.map (map_fst (sub.typ sub)) ptype_params)
       ~priv:ptype_private
       ~cstrs:(List.map
                 (map_tuple3 (sub.typ sub) (sub.typ sub) (sub.location sub))
                 ptype_cstrs)
       ~kind:(sub.type_kind sub ptype_kind)
-      ?manifest:(map_opt (sub.typ sub) ptype_manifest)
+      ~manifest:(map_opt (sub.typ sub) ptype_manifest)
+      ~jkind
+      ~docs:Docstrings.empty_docs
+      ~text:None
 
   let map_type_kind sub = function
     | Ptype_abstract -> Ptype_abstract

parsing/pprintast.ml

@@ -1742,10 +1742,18 @@ and type_def_list ctxt f (rf, exported, l) =
       else if exported then " ="
       else " :="
     in
-    pp f "@[<2>%s %a%a%s%s%a@]%a" kwd
+    let layout_annot, x =
+      match Jane_syntax.Layouts.of_type_declaration x with
+      | None -> "", x
+      | Some (jkind, remaining_attributes) ->
+          Printf.sprintf " : %s"
+            (Jane_asttypes.jkind_to_string jkind.txt),
+          { x with ptype_attributes = remaining_attributes }
+    in
+    pp f "@[<2>%s %a%a%s%s%s%a@]%a" kwd
       nonrec_flag rf
       (type_params ctxt) x.ptype_params
-      x.ptype_name.txt eq
+      x.ptype_name.txt layout_annot eq
       (type_declaration ctxt) x
       (item_attributes ctxt) x.ptype_attributes
   in

testsuite/tests/parsetree/ppx_no_op.ml

@@ -0,0 +1,8 @@
+open Ast_mapper
+
+(* This PPX rewriter does nothing. *)
+
+let () =
+  Language_extension.enable_maximal ();
+  Ast_mapper.register "no-op" (fun _ -> Ast_mapper.default_mapper);
+;;

testsuite/tests/parsetree/source_jane_street.ml

@@ -23,6 +23,14 @@ let f (type a : immediate) (type b : immediate)
       (type (c : immediate) (d : immediate))
   = ();;
 
+module type S_for_layouts = sig
+  type t : float64
+
+  type variant = A : ('a : immediate). 'a -> variant
+end;;
+
+type ('a : immediate) for_layouts = 'a;;
+
 (******************)
 (* Comprehensions *)
 
@@ -61,7 +69,7 @@ let f (type a : immediate) (type b : immediate)
 (* Local *)
 
 (* parameters *)
-let f (local_ x) ~(local_ y) ~z:(local_ z) ?foo:(local_ w = 1) = x + y + z + w;;
+let f (local_ x) ~(local_ y) ~z:(local_ z) ?foo:(local_ w = 1) () = x + y + z + w;;
 
 (* bindings *)
 let g () =
@@ -94,17 +102,22 @@ type 'a parameterized_record = {
 type fn = local_ int -> local_ int;;
 type nested_fn = (local_ int -> local_ int) -> local_ int;;
 type ('a, 'b) labeled_fn =
-  a:local_ 'a -> ?b:local_ b -> local_ 'a -> (int -> local_ 'b);;
+  a:local_ 'a -> ?b:local_ 'b -> local_ 'a -> (int -> local_ 'b);;
 
 (*******************)
 (* Include functor *)
 
+module F_struct (_ : sig end) = struct
+end
+
+module type F_sig = functor (_ : sig end) -> sig end
+
 module T = struct
-  include functor F
+  include functor F_struct
 end;;
 
 module type S = sig
-  include functor F
+  include functor F_sig
 end;;
 
 (********************)
@@ -115,10 +128,15 @@ let f x =
   | [::] -> [::]
   | ([:x:] [@test.attr1]) -> (([:x:])[@test.attr1])
   | ([:x;y:] [@test.attr2][@test.attr3]) ->
-      ([:x;y:] [@test.attr2][@test.attr3]);;
+      ([:x;y:] [@test.attr2][@test.attr3])
+  | _ -> assert false;;
 
 (******************)
 (* Labeled tuples *)
+let z, punned = 4, 5
+let x_must_be_even _ = assert false
+exception Odd
+
 let x = (~x:1, ~y:2)
 let x = ((~x:1, ~y:2) [@test.attr])
 let _ = ( ~x: 5, 2, ~z, ~(punned:int))

testsuite/tests/parsetree/test_ppx.compilers.reference

@@ -0,0 +1,19 @@
+(* TEST
+readonly_files = "source_jane_street.ml ppx_no_op.ml"
+include ocamlcommon
+* setup-ocamlc.byte-build-env
+** ocamlc.byte
+program = "${test_build_directory}/ppx_no_op.exe"
+all_modules = "ppx_no_op.ml"
+*** ocamlc.byte
+module = "source_jane_street.ml"
+flags = "-I ${test_build_directory} \
+         -w -26 \
+         -extension layouts \
+         -extension comprehensions \
+         -ppx ${program}"
+**** check-ocamlc.byte-output
+*)
+
+(* This test ensures that Jane Street syntax continues to be
+   handled properly by the compiler even after applying a PPX rewriter. *)