From db0ab3c8ca45ce26195170fa0571e264d698a3a2 Mon Sep 17 00:00:00 2001
From: Andrey Popp <8mayday@gmail.com>
Date: Wed, 6 Nov 2024 21:49:57 +0400
Subject: [PATCH] ppx: derive signatures

---
 ppx/test/intf.t                  | 24 ++++++++++++
 ppx/tools/ppx_deriving_tools.ml  | 64 ++++++++++++++++++++++++++------
 ppx/tools/ppx_deriving_tools.mli |  7 +++-
 3 files changed, 83 insertions(+), 12 deletions(-)
 create mode 100644 ppx/test/intf.t

diff --git a/ppx/test/intf.t b/ppx/test/intf.t
new file mode 100644
index 0000000..132265b
--- /dev/null
+++ b/ppx/test/intf.t
@@ -0,0 +1,24 @@
+
+  $ echo "type ('a, 'b) either [@@deriving json]" | ../browser/ppx_deriving_json_js_test.exe -intf -
+  type ('a, 'b) either[@@deriving json]
+  include
+    sig
+      [@@@ocaml.warning "-32"]
+      val either_of_json :
+        (Js.Json.t -> 'a) -> (Js.Json.t -> 'b) -> Js.Json.t -> ('a, 'b) either
+      val either_to_json :
+        ('a -> Js.Json.t) -> ('b -> Js.Json.t) -> ('a, 'b) either -> Js.Json.t
+    end[@@ocaml.doc "@inline"][@@merlin.hide ]
+
+  $ echo "type ('a, 'b) either [@@deriving json]" | ../native/ppx_deriving_json_native_test.exe -intf -
+  type ('a, 'b) either[@@deriving json]
+  include
+    sig
+      [@@@ocaml.warning "-32"]
+      val either_of_json :
+        (Yojson.Basic.t -> 'a) ->
+          (Yojson.Basic.t -> 'b) -> Yojson.Basic.t -> ('a, 'b) either
+      val either_to_json :
+        ('a -> Yojson.Basic.t) ->
+          ('b -> Yojson.Basic.t) -> ('a, 'b) either -> Yojson.Basic.t
+    end[@@ocaml.doc "@inline"][@@merlin.hide ]
diff --git a/ppx/tools/ppx_deriving_tools.ml b/ppx/tools/ppx_deriving_tools.ml
index cea484f..89cb3dc 100644
--- a/ppx/tools/ppx_deriving_tools.ml
+++ b/ppx/tools/ppx_deriving_tools.ml
@@ -82,27 +82,38 @@ class virtual deriving =
     method virtual extension
         : loc:location -> path:label -> core_type -> expression
 
-    method virtual generator
+    method virtual str_type_decl
         : ctxt:Expansion_context.Deriver.t ->
           rec_flag * type_declaration list ->
           structure
+
+    method virtual sig_type_decl
+        : ctxt:Expansion_context.Deriver.t ->
+          rec_flag * type_declaration list ->
+          signature
   end
 
 let register ?deps deriving =
-  Deriving.add deriving#name
-    ~str_type_decl:
-      (Deriving.Generator.V2.make ?deps Deriving.Args.empty
-         deriving#generator)
-    ~extension:deriving#extension
+  let args = Deriving.Args.empty in
+  let str_type_decl = deriving#str_type_decl in
+  let sig_type_decl = deriving#sig_type_decl in
+  Deriving.add deriving#name ~extension:deriving#extension
+    ~str_type_decl:(Deriving.Generator.V2.make ?deps args str_type_decl)
+    ~sig_type_decl:(Deriving.Generator.V2.make ?deps args sig_type_decl)
 
 let register_combined ?deps name derivings =
-  let generator ~ctxt bindings =
+  let args = Deriving.Args.empty in
+  let str_type_decl ~ctxt bindings =
     List.fold_left derivings ~init:[] ~f:(fun str d ->
-        d#generator ~ctxt bindings @ str)
+        d#str_type_decl ~ctxt bindings @ str)
+  in
+  let sig_type_decl ~ctxt bindings =
+    List.fold_left derivings ~init:[] ~f:(fun str d ->
+        d#sig_type_decl ~ctxt bindings @ str)
   in
   Deriving.add name
-    ~str_type_decl:
-      (Deriving.Generator.V2.make ?deps Deriving.Args.empty generator)
+    ~str_type_decl:(Deriving.Generator.V2.make ?deps args str_type_decl)
+    ~sig_type_decl:(Deriving.Generator.V2.make ?deps args sig_type_decl)
 
 module Schema = struct
   let repr_row_field field =
@@ -162,6 +173,30 @@ module Schema = struct
                Location.raise_errorf ~loc
                  "this cannot be a type parameter"))
 
+  let derive_sig_type_decl ~derive_t ~derive_label ~ctxt (_rec_flag, tds)
+      =
+    let loc = Expansion_context.Deriver.derived_item_loc ctxt in
+    List.map tds ~f:(fun td ->
+        let name = td.ptype_name in
+        let type_ = derive_t ~loc name (gen_type_ascription td) in
+        let type_ =
+          List.fold_left (List.rev td.ptype_params) ~init:type_
+            ~f:(fun acc (t, _) ->
+              let loc = t.ptyp_loc in
+              let name =
+                match t.ptyp_desc with
+                | Ptyp_var txt -> { txt; loc }
+                | _ ->
+                    Location.raise_errorf ~loc
+                      "type variable is not a variable"
+              in
+              let t = derive_t ~loc name t in
+              ptyp_arrow ~loc Nolabel t acc)
+        in
+        psig_value ~loc
+          (value_description ~loc ~prim:[] ~name:(derive_label name)
+             ~type_))
+
   class virtual deriving1 =
     object (self)
       inherit deriving
@@ -276,7 +311,7 @@ module Schema = struct
           let loc = ty.ptyp_loc in
           as_fun ~loc (self#derive_of_core_type' ty)
 
-      method generator
+      method str_type_decl
           : ctxt:Expansion_context.Deriver.t ->
             rec_flag * type_declaration list ->
             structure =
@@ -289,6 +324,13 @@ module Schema = struct
             [@@@ocaml.warning "-39-11-27"]
 
             [%%i pstr_value ~loc Recursive bindings]]
+
+      method sig_type_decl
+          : ctxt:Expansion_context.Deriver.t ->
+            rec_flag * type_declaration list ->
+            signature =
+        derive_sig_type_decl ~derive_t:self#t
+          ~derive_label:self#derive_type_decl_label
     end
 end
 
diff --git a/ppx/tools/ppx_deriving_tools.mli b/ppx/tools/ppx_deriving_tools.mli
index 18ff21d..5ccf8a9 100644
--- a/ppx/tools/ppx_deriving_tools.mli
+++ b/ppx/tools/ppx_deriving_tools.mli
@@ -11,11 +11,16 @@ class virtual deriving : object
     loc:location -> path:label -> core_type -> expression
   (** a deriver can be applied to as type expression as extension node. *)
 
-  method virtual generator :
+  method virtual str_type_decl :
     ctxt:Expansion_context.Deriver.t ->
     rec_flag * type_declaration list ->
     structure
   (** or it can be attached to a type declaration. *)
+
+  method virtual sig_type_decl :
+    ctxt:Expansion_context.Deriver.t ->
+    rec_flag * type_declaration list ->
+    signature
 end
 
 val register : ?deps:Deriving.t list -> deriving -> Deriving.t