Add [@@@expand_inline] and support for floating attribute context free transformations

CHANGES.md

@@ -1,6 +1,9 @@
 unreleased
 ----------
 
+- Add the `[@@@expand_inline]` transformation and support for floating attribute context
+  free transformations. (#560, @jaymody)
+
 - Add a `-raise-embedded-errors` flag to the driver. Setting this flag raises the first
   `ocaml.error` embedded in the final AST.
 

doc/writing-ppxs.mld

@@ -467,6 +467,24 @@ This mechanism is implemented for derivers implemented in [ppxlib] and is conven
 
 Inline derivers will generate a [.corrected] version of the file that Dune can use to promote your file. For more information on how to use this feature to remove a dependency on [ppxlib] and a specific PPX from your project, refer to {{:https://ocaml.org/docs/metaprogramming#dropping-ppxs-dependency-with-derivinginline}this guide}.
 
+In addition to [[@@deriving_inline]], there is also [[@@@expand_inline <structure payload>]] and [[@@@expand_inline: <signature payload>]]. These can be use to inline code generated by other context free transformations (not just derivers):
+
+{@ocaml[
+  [@@@expand_inline let _ = [%add_suffix "foo"]]
+
+  let _ = "foo_suffixed"
+
+  [@@@end]
+
+  module type S = sig
+    [@@@expand_inline: type foo = [%pair_of string]]
+
+    type foo = string * string
+
+    [@@@end]
+  end
+]}
+
 {1 Integration with Dune}
 
 If your PPX is written as a Dune project, you'll need to specify the [kind]

src/attribute.ml

@@ -454,6 +454,12 @@ module Floating = struct
       payload = Payload_parser (pattern, fun ~attr_loc:_ ~name_loc:_ -> k);
     }
 
+  let convert_attr_res t attr =
+    let open Result in
+    if Name.Pattern.matches t.name attr.attr_name.txt then
+      convert t.payload attr >>| fun value -> Some value
+    else Ok None
+
   let convert_res ts x =
     let open Result in
     match ts with

src/attribute.mli

@@ -220,6 +220,11 @@ module Floating : sig
     ('a, 'b) t list -> 'a -> ('b option, Location.Error.t NonEmptyList.t) result
 
   val convert : ('a, 'b) t list -> 'a -> 'b option
+
+  val convert_attr_res :
+    ('a, 'b) t ->
+    attribute ->
+    ('b option, Location.Error.t NonEmptyList.t) result
 end
 
 val explicitly_drop : Ast_traverse0.iter

src/common.ml

@@ -276,3 +276,12 @@ module With_errors = struct
 
   let combine_errors list = (List.map list ~f:fst, List.concat_map list ~f:snd)
 end
+
+let valid_string_constant_delimiter string =
+  let rec attempt_string_constant_delimiter n =
+    let delimiter = String.make n 'x' in
+    if String.is_substring string ~substring:("|" ^ delimiter ^ "}") then
+      attempt_string_constant_delimiter (n + 1)
+    else delimiter
+  in
+  attempt_string_constant_delimiter 0

src/common.mli

@@ -103,3 +103,7 @@ module With_errors : sig
 
   val combine_errors : 'a t list -> 'a list t
 end
+
+val valid_string_constant_delimiter : string -> string
+(** [valid_string_constant_delimiter x] finds a delimiter [y] such that
+    [Pconst_string (x, loc, Some y)] is valid. *)

src/context_free.ml

@@ -43,6 +43,21 @@ module Rule = struct
       List.partition l ~f:(fun (T t) -> not t.expect)
   end
 
+  module Attr_floating_inline = struct
+    type ('a, 'b) unpacked = {
+      attribute : ('a, 'b) Attribute.Floating.t;
+      expand_items : bool;
+      expand : ctxt:Expansion_context.Deriver.t -> 'b -> 'a list;
+    }
+
+    type 'a t = T : ('a, _) unpacked -> 'a t
+
+    let attr_name (T t) = Attribute.Floating.name t.attribute
+
+    let split_normal_and_expand l =
+      List.partition l ~f:(fun (T t) -> not t.expand_items)
+  end
+
   module Special_function = struct
     type t = {
       name : string;
@@ -84,6 +99,8 @@ module Rule = struct
           (structure_item, class_type_declaration) Attr_group_inline.t t
       | Attr_sig_class_type_decl :
           (signature_item, class_type_declaration) Attr_group_inline.t t
+      | Attr_str_floating : structure_item Attr_floating_inline.t t
+      | Attr_sig_floating : signature_item Attr_floating_inline.t t
 
     type (_, _) equality = Eq : ('a, 'a) equality | Ne : (_, _) equality
 
@@ -103,6 +120,8 @@ module Rule = struct
       | Attr_sig_module_type_decl, Attr_sig_module_type_decl -> Eq
       | Attr_str_class_type_decl, Attr_str_class_type_decl -> Eq
       | Attr_sig_class_type_decl, Attr_sig_class_type_decl -> Eq
+      | Attr_str_floating, Attr_str_floating -> Eq
+      | Attr_sig_floating, Attr_sig_floating -> Eq
       | _ -> Ne
   end
 
@@ -122,6 +141,11 @@ module Rule = struct
     (ctxt:Expansion_context.Deriver.t -> 'b -> 'c -> 'a list) ->
     t
 
+  type ('item, 'parsed_payload) attr_floating_inline =
+    ('item, 'parsed_payload) Attribute.Floating.t ->
+    (ctxt:Expansion_context.Deriver.t -> 'parsed_payload -> 'item list) ->
+    t
+
   let rec filter : type a. a Field.t -> t list -> a list =
    fun field l ->
     match l with
@@ -200,6 +224,18 @@ module Rule = struct
 
   let attr_sig_class_type_decl_expect attribute expand =
     T (Attr_sig_class_type_decl, T { attribute; expand; expect = true })
+
+  let attr_str_floating_expect attribute expand =
+    T (Attr_str_floating, T { attribute; expand; expand_items = false })
+
+  let attr_sig_floating_expect attribute expand =
+    T (Attr_sig_floating, T { attribute; expand; expand_items = false })
+
+  let attr_str_floating_expect_and_expand attribute expand =
+    T (Attr_str_floating, T { attribute; expand; expand_items = true })
+
+  let attr_sig_floating_expect_and_expand attribute expand =
+    T (Attr_sig_floating, T { attribute; expand; expand_items = true })
 end
 
 module Generated_code_hook = struct
@@ -402,6 +438,12 @@ let sort_attr_inline l =
         (Rule.Attr_inline.attr_name a)
         (Rule.Attr_inline.attr_name b))
 
+let sort_attr_floating_inline l =
+  List.sort l ~cmp:(fun a b ->
+      String.compare
+        (Rule.Attr_floating_inline.attr_name a)
+        (Rule.Attr_floating_inline.attr_name b))
+
 let context_free_attribute_modification ~loc =
   Error
     ( Location.Error.createf ~loc
@@ -462,6 +504,27 @@ let handle_attr_inline attrs ~convert_exn ~item ~expanded_item ~loc ~base_ctxt
             let error_item = [ convert_exn exn ] in
             return (error_item :: acc)))
 
+let handle_attr_floating_inline attrs ~item ~loc ~base_ctxt ~embed_errors
+    ~convert_exn =
+  List.fold_left attrs ~init:(return [])
+    ~f:(fun acc (Rule.Attr_floating_inline.T a) ->
+      acc >>= fun acc ->
+      Attribute.Floating.convert_attr_res a.attribute item
+      |> of_result ~default:None
+      >>= function
+      | None -> return acc
+      | Some value -> (
+          let ctxt =
+            Expansion_context.Deriver.make ~derived_item_loc:loc ~inline:true
+              ~base:base_ctxt ()
+          in
+          try
+            let expect_items = a.expand ~ctxt value in
+            return (expect_items :: acc)
+          with exn when embed_errors ->
+            let error_item = [ convert_exn exn ] in
+            return (error_item :: acc)))
+
 module Expect_mismatch_handler = struct
   type t = {
     f : 'a. 'a Attribute.Floating.Context.t -> Location.t -> 'a list -> unit;
@@ -542,10 +605,22 @@ class map_top_down ?(expect_mismatch_handler = Expect_mismatch_handler.nop)
     |> sort_attr_group_inline |> Rule.Attr_group_inline.split_normal_and_expect
   in
 
+  let attr_str_floating_expect, attr_str_floating_expect_and_expand =
+    Rule.filter Attr_str_floating rules
+    |> sort_attr_floating_inline
+    |> Rule.Attr_floating_inline.split_normal_and_expand
+  in
+  let attr_sig_floating_expect, attr_sig_floating_expect_and_expand =
+    Rule.filter Attr_sig_floating rules
+    |> sort_attr_floating_inline
+    |> Rule.Attr_floating_inline.split_normal_and_expand
+  in
+
   let map_node = map_node ~hook ~embed_errors in
   let map_nodes = map_nodes ~hook ~embed_errors in
   let handle_attr_group_inline = handle_attr_group_inline ~embed_errors in
   let handle_attr_inline = handle_attr_inline ~embed_errors in
+  let handle_attr_floating_inline = handle_attr_floating_inline ~embed_errors in
 
   object (self)
     inherit Ast_traverse.map_with_expansion_context_and_errors as super
@@ -724,6 +799,7 @@ class map_top_down ?(expect_mismatch_handler = Expect_mismatch_handler.nop)
     (* TODO: try to factorize #structure and #signature without meta-programming *)
     (*$*)
     method! structure base_ctxt st =
+      let convert_exn = exn_to_stri in
       let rec with_extra_items item ~extra_items ~expect_items ~rest
           ~in_generated_code =
         loop (rev_concat extra_items) ~in_generated_code:true
@@ -769,9 +845,34 @@ class map_top_down ?(expect_mismatch_handler = Expect_mismatch_handler.nop)
                       Generated_code_hook.replace hook Structure_item
                         item.pstr_loc (Many items);
                     loop rest ~in_generated_code >>| fun rest -> items @ rest)
+            | Pstr_attribute at ->
+                handle_attr_floating_inline attr_str_floating_expect ~item:at
+                  ~loc ~base_ctxt ~convert_exn
+                >>= fun expect_items ->
+                handle_attr_floating_inline attr_str_floating_expect_and_expand
+                  ~item:at ~loc ~base_ctxt ~convert_exn
+                >>= fun expect_items_unexpanded ->
+                List.map expect_items_unexpanded ~f:(self#structure base_ctxt)
+                |> combine_errors
+                >>= fun expect_items_expanded ->
+                (* Shouldn't matter if we use [rev_concat] or [List.concat] here, there
+                   should be only one (outer) list among [expect_items] and
+                   [expect_items_expanded] unless a single floating attribute is somehow
+                   registered twice. *)
+                (match rev_concat (expect_items @ expect_items_expanded) with
+                | [] -> return ()
+                | expected ->
+                    Code_matcher.match_structure_res rest
+                      ~pos:item.pstr_loc.loc_end ~expected
+                      ~mismatch_handler:
+                        (expect_mismatch_handler.f Structure_item)
+                    |> of_result ~default:())
+                >>= fun () ->
+                super#structure_item base_ctxt item >>= fun expanded_item ->
+                loop rest ~in_generated_code >>| fun expanded_rest ->
+                expanded_item :: expanded_rest
             | _ -> (
                 super#structure_item base_ctxt item >>= fun expanded_item ->
-                let convert_exn = exn_to_stri in
                 match (item.pstr_desc, expanded_item.pstr_desc) with
                 | Pstr_type (rf, tds), Pstr_type (exp_rf, exp_tds) ->
                     (* No context-free rule can rewrite rec flags atm, this
@@ -833,6 +934,7 @@ class map_top_down ?(expect_mismatch_handler = Expect_mismatch_handler.nop)
 
     (*$ str_to_sig _last_text_block *)
     method! signature base_ctxt sg =
+      let convert_exn = exn_to_sigi in
       let rec with_extra_items item ~extra_items ~expect_items ~rest
           ~in_generated_code =
         loop (rev_concat extra_items) ~in_generated_code:true
@@ -878,9 +980,34 @@ class map_top_down ?(expect_mismatch_handler = Expect_mismatch_handler.nop)
                       Generated_code_hook.replace hook Signature_item
                         item.psig_loc (Many items);
                     loop rest ~in_generated_code >>| fun rest -> items @ rest)
+            | Psig_attribute at ->
+                handle_attr_floating_inline attr_sig_floating_expect ~item:at
+                  ~loc ~base_ctxt ~convert_exn
+                >>= fun expect_items ->
+                handle_attr_floating_inline attr_sig_floating_expect_and_expand
+                  ~item:at ~loc ~base_ctxt ~convert_exn
+                >>= fun expect_items_unexpanded ->
+                List.map expect_items_unexpanded ~f:(self#signature base_ctxt)
+                |> combine_errors
+                >>= fun expect_items_expanded ->
+                (* Shouldn't matter if we use [rev_concat] or [List.concat] here, there
+                   should be only one (outer) list among [expect_items] and
+                   [expect_items_expanded] unless a single floating attribute is somehow
+                   registered twice. *)
+                (match rev_concat (expect_items @ expect_items_expanded) with
+                | [] -> return ()
+                | expected ->
+                    Code_matcher.match_signature_res rest
+                      ~pos:item.psig_loc.loc_end ~expected
+                      ~mismatch_handler:
+                        (expect_mismatch_handler.f Signature_item)
+                    |> of_result ~default:())
+                >>= fun () ->
+                super#signature_item base_ctxt item >>= fun expanded_item ->
+                loop rest ~in_generated_code >>| fun expanded_rest ->
+                expanded_item :: expanded_rest
             | _ -> (
                 super#signature_item base_ctxt item >>= fun expanded_item ->
-                let convert_exn = exn_to_sigi in
                 match (item.psig_desc, expanded_item.psig_desc) with
                 | Psig_type (rf, tds), Psig_type (exp_rf, exp_tds) ->
                     (* No context-free rule can rewrite rec flags atm, this

src/context_free.mli

@@ -129,6 +129,20 @@ module Rule : sig
 
   val attr_sig_class_type_decl_expect :
     (signature_item, class_type_declaration, _) attr_group_inline
+
+  type ('item, 'parsed_payload) attr_floating_inline =
+    ('item, 'parsed_payload) Attribute.Floating.t ->
+    (ctxt:Expansion_context.Deriver.t -> 'parsed_payload -> 'item list) ->
+    t
+
+  val attr_str_floating_expect : (structure_item, _) attr_floating_inline
+  val attr_sig_floating_expect : (signature_item, _) attr_floating_inline
+
+  val attr_str_floating_expect_and_expand :
+    (structure_item, _) attr_floating_inline
+
+  val attr_sig_floating_expect_and_expand :
+    (signature_item, _) attr_floating_inline
 end
 
 (**/**)

src/driver.ml

@@ -1581,3 +1581,19 @@ let enable_checks () =
 let enable_location_check () = perform_locations_check := true
 let disable_location_check () = perform_locations_check := false
 let map_structure st = map_structure st
+
+let () =
+  register_transformation "expand_inline"
+    ~rules:
+      [
+        Context_free.Rule.attr_str_floating_expect_and_expand
+          (Attribute.Floating.declare "expand_inline" Structure_item
+             Ast_pattern.(pstr __)
+             Fn.id)
+          (fun ~ctxt:_ items -> Utils.prettify_odoc_attributes#structure items);
+        Context_free.Rule.attr_sig_floating_expect_and_expand
+          (Attribute.Floating.declare "expand_inline" Signature_item
+             Ast_pattern.(psig __)
+             Fn.id)
+          (fun ~ctxt:_ items -> Utils.prettify_odoc_attributes#signature items);
+      ]

src/utils.ml

@@ -1,5 +1,53 @@
 open Import
 
+(* We want to make sure we only convert actual odoc comments [(** ... *)] and not
+   actual instances of [@@ocaml.doc "..."]. When parsed, both get translated as
+   an attribute.
+
+   To differentiate them, we can take advantage of the fact that the location
+   attached to the attribute node for (** ... *) comments is equal to the location of
+   the string itself, while for [@@ocaml.doc "..."] they are different.
+
+   The same is true for [@@@ocaml.text]. *)
+let get_odoc_contents_if_comment = function
+  | {
+      attr_loc;
+      attr_name = { txt = "doc" | "ocaml.doc" | "text" | "ocaml.text"; _ };
+      attr_payload =
+        PStr
+          [
+            {
+              pstr_desc =
+                Pstr_eval
+                  ( {
+                      pexp_desc = Pexp_constant (Pconst_string (text, loc, _));
+                      _;
+                    },
+                    _ );
+              _;
+            };
+          ];
+    }
+    when Location.compare attr_loc loc = 0 ->
+      Some text
+  | _ -> None
+
+let prettify_odoc_attributes =
+  object
+    inherit Ast_traverse.map as super
+
+    method! attribute attr =
+      let attr = super#attribute attr in
+      match get_odoc_contents_if_comment attr with
+      | Some txt ->
+          let open Ast_builder.Default in
+          let loc = Location.none in
+          let delim = Some (Common.valid_string_constant_delimiter txt) in
+          let expr = pexp_constant ~loc (Pconst_string (txt, loc, delim)) in
+          { attr with attr_payload = PStr [ pstr_eval ~loc expr [] ] }
+      | None -> attr
+  end
+
 let with_output fn ~binary ~f =
   match fn with
   | None | Some "-" ->