[New Command] Expand PPX nodes (#1745)

CHANGES.md

@@ -6,6 +6,7 @@ unreleased
       what to append to the current unit name in the presence of wrapping (#1788)
     - Add `-unboxed-types` and `-no-unboxed-types` as ocaml ignored flags (#1795, fixes #1794)
     - destruct: Refinement in the presence of optional arguments (#1800 #1807, fixes #1770)
+    - Implement new expand-node command for expanding PPX annotations (#1745)
   + editor modes
     - vim: fix python-3.12 syntax warnings in merlin.py (#1798)
     - vim: Dead code / doc removal for previously deleted MerlinPhrase command (#1804)

src/analysis/ppx_expand.ml

@@ -0,0 +1,159 @@
+type ppx_kind =
+  | Expr of Parsetree.expression
+  | Sig_item of Parsetree.signature_item
+  | Str_item of Parsetree.structure_item
+
+let check_at_pos pos loc = Location_aux.compare_pos pos loc = 0
+
+let check_extension_node pos (expression : Parsetree.expression) =
+  match expression.pexp_desc with
+  | Pexp_extension (loc, _) ->
+      if check_at_pos pos loc.loc then Some expression.pexp_loc else None
+  | _ -> None
+
+let check_deriving_attr pos (attrs : Parsetree.attributes) =
+  let found_attr =
+    List.find_opt
+      (fun (attribute : Parsetree.attribute) ->
+        attribute.attr_name.txt = "deriving"
+        && check_at_pos pos attribute.attr_loc)
+      attrs
+  in
+  match found_attr with
+  | Some attribute -> Some attribute.attr_loc
+  | None -> None
+
+let check_structures pos (item : Parsetree.structure_item_desc) =
+  match item with
+  | Pstr_type (_, ty) ->
+      List.find_map
+        (fun (t : Parsetree.type_declaration) ->
+          check_deriving_attr pos t.ptype_attributes)
+        ty
+  | Pstr_exception tc -> check_deriving_attr pos tc.ptyexn_attributes
+  | Pstr_modtype mt -> check_deriving_attr pos mt.pmtd_attributes
+  | Pstr_typext tex -> check_deriving_attr pos tex.ptyext_attributes
+  | _ -> None
+
+let check_signatures pos (item : Parsetree.signature_item_desc) =
+  match item with
+  | Psig_type (_, ty) ->
+      List.find_map
+        (fun (t : Parsetree.type_declaration) ->
+          check_deriving_attr pos t.ptype_attributes)
+        ty
+  | Psig_exception tc -> check_deriving_attr pos tc.ptyexn_attributes
+  | Psig_modtype mt -> check_deriving_attr pos mt.pmtd_attributes
+  | Psig_typext tex -> check_deriving_attr pos tex.ptyext_attributes
+  | _ -> None
+
+let check_extension ~parsetree ~pos =
+  let kind = ref None in
+  let expr (self : Ast_iterator.iterator) (expr : Parsetree.expression) =
+    match check_extension_node pos expr with
+    | Some ext_loc -> kind := Some (Expr expr, ext_loc)
+    | None -> Ast_iterator.default_iterator.expr self expr
+  in
+  let signature_item (self : Ast_iterator.iterator)
+      (original_sg : Parsetree.signature_item) =
+    match check_signatures pos original_sg.psig_desc with
+    | Some attr_loc -> kind := Some (Sig_item original_sg, attr_loc)
+    | None -> Ast_iterator.default_iterator.signature_item self original_sg
+  in
+  let structure_item (self : Ast_iterator.iterator)
+      (original_str : Parsetree.structure_item) =
+    match check_structures pos original_str.pstr_desc with
+    | Some attr_loc -> kind := Some (Str_item original_str, attr_loc)
+    | None -> Ast_iterator.default_iterator.structure_item self original_str
+  in
+  let iterator =
+    { Ast_iterator.default_iterator with signature_item; structure_item; expr }
+  in
+  let () =
+    match parsetree with
+    | `Interface si -> iterator.signature iterator si
+    | `Implementation str -> iterator.structure iterator str
+  in
+  !kind
+
+let get_ppxed_source ~ppxed_parsetree ~pos ppx_kind_with_attr :
+    Query_protocol.ppxed_source =
+  let expression = ref None in
+  let signature = ref [] in
+  let structure = ref [] in
+  let () =
+    match ppx_kind_with_attr with
+    | Expr original_expr, _ -> (
+        let expr (self : Ast_iterator.iterator)
+            (new_expr : Parsetree.expression) =
+          match
+            Location_aux.included ~into:original_expr.pexp_loc new_expr.pexp_loc
+          with
+          | true -> expression := Some new_expr
+          | false -> Ast_iterator.default_iterator.expr self new_expr
+        in
+        let iterator = { Ast_iterator.default_iterator with expr } in
+        match ppxed_parsetree with
+        | `Interface si -> iterator.signature iterator si
+        | `Implementation str -> iterator.structure iterator str)
+    | Sig_item original_sg, _ -> (
+        let signature_item (self : Ast_iterator.iterator)
+            (new_sg : Parsetree.signature_item) =
+          let included =
+            Location_aux.included new_sg.psig_loc ~into:original_sg.psig_loc
+          in
+          match included && original_sg <> new_sg, new_sg.psig_loc.loc_ghost with
+          | true, _ -> signature := new_sg :: !signature
+          | false, false -> Ast_iterator.default_iterator.signature_item self new_sg
+          | false, true -> () (* We don't enter nested ppxes *)
+        in
+        let iterator = { Ast_iterator.default_iterator with signature_item } in
+        match ppxed_parsetree with
+        | `Interface si -> iterator.signature iterator si
+        | `Implementation str -> iterator.structure iterator str)
+    | Str_item original_str, _ -> (
+        let structure_item (self : Ast_iterator.iterator)
+            (new_str : Parsetree.structure_item) =
+          let included =
+            Location_aux.included new_str.pstr_loc ~into:original_str.pstr_loc
+          in
+          match included, new_str.pstr_loc.loc_ghost with
+          | true, _ -> 
+            (match check_structures pos new_str.pstr_desc with 
+            | None -> structure := new_str :: !structure
+            | Some _ -> ())
+          | false, false -> Ast_iterator.default_iterator.structure_item self new_str
+          | false, true -> ()
+        in
+        let iterator = { Ast_iterator.default_iterator with structure_item } in
+        match ppxed_parsetree with
+        | `Interface si -> iterator.signature iterator si
+        | `Implementation str -> iterator.structure iterator str)
+  in
+  match (ppx_kind_with_attr : ppx_kind * Warnings.loc) with
+  | Expr _, ext_loc ->
+      {
+        code = Pprintast.string_of_expression (Option.get !expression);
+        attr_start = ext_loc.loc_start;
+        attr_end = ext_loc.loc_end;
+      }
+  | Sig_item _, attr_loc ->
+      let exp =
+        Pprintast.signature Format.str_formatter (List.rev !signature);
+        Format.flush_str_formatter ()
+      in
+      {
+        code = exp;
+        attr_start = attr_loc.loc_start;
+        attr_end = attr_loc.loc_end;
+      }
+  | Str_item _, attr_loc ->
+      let exp =
+        Pprintast.structure Format.str_formatter (List.rev !structure);
+        Format.flush_str_formatter ()
+      in
+      {
+        code = exp;
+        attr_start = attr_loc.loc_start;
+        attr_end = attr_loc.loc_end;
+      }

src/analysis/ppx_expand.mli

@@ -0,0 +1,19 @@
+type ppx_kind =
+  | Expr of Parsetree.expression
+  | Sig_item of Parsetree.signature_item
+  | Str_item of Parsetree.structure_item
+
+val check_extension :
+  parsetree:
+    [ `Implementation of Parsetree.structure
+    | `Interface of Parsetree.signature ] ->
+  pos:Lexing.position ->
+  (ppx_kind * Warnings.loc) option
+
+val get_ppxed_source :
+  ppxed_parsetree:
+    [ `Implementation of Parsetree.structure
+    | `Interface of Parsetree.signature ] ->
+  pos:Lexing.position ->
+  ppx_kind * Warnings.loc ->
+  Query_protocol.ppxed_source

src/commands/new_commands.ml

@@ -242,6 +242,21 @@ Otherwise, Merlin looks for the documentation for the entity under the cursor (a
     end
   ;
 
+  command "expand-ppx"
+    ~doc: "Returns the generated code of a PPX."
+    ~spec: [
+      arg "-position" "<position> Position to expand"
+          (marg_position (fun pos _pos -> pos));
+    ]
+    ~default: `None
+    begin fun buffer pos ->
+      match pos with 
+      | `None -> failwith "-position <pos> is mandatory"
+      | #Msource.position as pos -> 
+        run buffer (Query_protocol.Expand_ppx pos)
+    end
+  ;
+
   command "enclosing"
     ~spec: [
       arg "-position" "<position> Position to complete"

src/commands/query_json.ml

@@ -114,6 +114,8 @@ let dump (type a) : a t -> json =
     ]
   | Syntax_document pos ->
     mk "syntax-document" [ ("position", mk_position pos) ]
+  | Expand_ppx pos ->
+    mk "ppx-expand" [ ("position", mk_position pos) ]
   | Locate (prefix, look_for, pos) ->
     mk "locate" [
       "prefix", (match prefix with
@@ -392,6 +394,19 @@ let json_of_response (type a) (query : a t) (response : a) : json =
         ("url", `String info.documentation);
       ]
     | `No_documentation -> `String "No documentation found")
+  | Expand_ppx _, resp -> 
+    let str = match resp with
+    | `Found ppx_info -> 
+      `Assoc
+      [
+        ("code", `String ppx_info.code);
+        ("deriver", `Assoc [
+          ("start", Lexing.json_of_position ppx_info.attr_start);
+          ("end", Lexing.json_of_position ppx_info.attr_end);
+        ])
+      ]
+    | `No_ppx -> `String "No PPX deriver/extension node found on this position"
+    in str
   | Locate_type _, resp -> json_of_locate resp
   | Locate _, resp -> json_of_locate resp
   | Jump _, resp ->

src/frontend/query_commands.ml

@@ -511,6 +511,18 @@ let dispatch pipeline (type a) : a Query_protocol.t -> a =
     | Some res -> `Found res
     | None -> `No_documentation)
 
+  | Expand_ppx pos -> (
+    let pos = Mpipeline.get_lexing_pos pipeline pos in
+    let parsetree = Mpipeline.reader_parsetree pipeline in
+    let ppxed_parsetree = Mpipeline.ppx_parsetree pipeline in
+    let ppx_kind_with_attr = Ppx_expand.check_extension ~parsetree ~pos in
+    match ppx_kind_with_attr with
+    | Some _ ->
+        `Found
+          (Ppx_expand.get_ppxed_source ~ppxed_parsetree ~pos
+              (Option.get ppx_kind_with_attr))
+    | None -> `No_ppx)
+
   | Locate (patho, ml_or_mli, pos) ->
     let typer = Mpipeline.typer_result pipeline in
     let local_defs = Mtyper.get_typedtree typer in

src/frontend/query_protocol.ml

@@ -103,6 +103,13 @@ type syntax_doc_result =
     documentation : string
 }
 
+type ppxed_source = 
+{
+  code : string;
+  attr_start : Lexing.position;
+  attr_end : Lexing.position;
+}
+
 type is_tail_position = [`No | `Tail_position | `Tail_call]
 
 type _ _bool = bool
@@ -152,6 +159,11 @@ type _ t =
     -> [ `Found of syntax_doc_result
        | `No_documentation
        ] t
+  | Expand_ppx
+    : Msource.position
+    -> [ `Found of ppxed_source
+       | `No_ppx
+       ] t
   | Locate_type
     : Msource.position
       -> [ `Found of string option * Lexing.position