merge extractors

ocaml-ppx · Feb 23, 2023 · bb907f0 · bb907f0
1 parent 4ebeb19
commit bb907f0
Show file tree

Hide file tree

Showing 3 changed files with 57 additions and 80 deletions.
diff --git a/src/ppx_import.ml b/src/ppx_import.ml
@@ -648,89 +648,70 @@ let module_declaration_expand_intf ~ctxt modtype_decl =
   in
   Ppxlib.{psig_desc = Psig_modtype md_decl; psig_loc = loc}
 
-let type_declaration_expander ~ctxt payload =
-  let return_error e =
-    let loc = Ppxlib.Expansion_context.Extension.extension_point_loc ctxt in
-    let ext = Ppxlib.Location.error_extensionf ~loc "%s" e in
-    Ppxlib.Ast_builder.Default.pstr_extension ext [] ~loc
-  in
+type extracted_payload =
+  | Type_decl of Asttypes.rec_flag * Parsetree.type_declaration list
+  | Module_type_decl of Parsetree.module_type_declaration
+
+let type_extractor =
+  Ppxlib.Ast_pattern.(
+    pstr (pstr_type __ __ ^:: nil)
+    ||| psig (psig_type __ __ ^:: nil)
+    |> map2 ~f:(fun rec_flag type_decl -> Type_decl (rec_flag, type_decl)) )
+
+let module_type_extractor =
+  Ppxlib.Ast_pattern.(
+    psig (psig_modtype __ ^:: nil)
+    ||| pstr (pstr_modtype __ ^:: nil)
+    |> map1 ~f:(fun modtype -> Module_type_decl modtype) )
+
+let extractor = Ppxlib.Ast_pattern.(type_extractor ||| module_type_extractor)
+
+let expander ~(ctxt : Ppxlib.Expansion_context.Extension.t)
+    (payload : extracted_payload) =
   match payload with
-  | Parsetree.PStr [{pstr_desc = Pstr_type (rec_flag, type_decls); _}]
-   |Parsetree.PSig [{psig_desc = Psig_type (rec_flag, type_decls); _}] ->
+  | Type_decl (rec_flag, type_decls) ->
     type_declaration_expand ~ctxt rec_flag type_decls
-  | Parsetree.PStr [{pstr_desc = Pstr_modtype modtype_decl; _}]
-   |Parsetree.PSig [{psig_desc = Psig_modtype modtype_decl; _}] ->
+  | Module_type_decl modtype_decl ->
     module_declaration_expand ~ctxt modtype_decl
-  | Parsetree.PStr [{pstr_desc = _; _}] | Parsetree.PSig [{psig_desc = _; _}] ->
-    return_error
-      "[%%import] Expected a type declaration or a module type declaration"
-  | Parsetree.PStr (_ :: _) | Parsetree.PSig (_ :: _) ->
-    return_error
-      "[%%import] Expected exactly one item in the structure or signature, but \
-       found multiple items"
-  | Parsetree.PStr [] | Parsetree.PSig [] ->
-    return_error
-      "[%%import] Expected exactly one item in the structure or signature, but \
-       found none"
-  | Parsetree.PTyp _ ->
-    return_error
-      "[%%import] Type pattern (PTyp) is not supported, only type and module \
-       type declarations are allowed"
-  | Parsetree.PPat (_, _) ->
-    return_error
-      "[%%import] Pattern (PPat) is not supported, only type and module type \
-       declarations are allowed"
-
-let type_declaration_extension =
+
+let import_extension =
   Ppxlib.Extension.V3.declare "import" Ppxlib.Extension.Context.structure_item
-    Ppxlib.Ast_pattern.(__)
-    type_declaration_expander
-
-let type_declaration_expander_intf ~ctxt payload =
-  let return_error e =
-    let loc = Ppxlib.Expansion_context.Extension.extension_point_loc ctxt in
-    let ext = Ppxlib.Location.error_extensionf ~loc "%s" e in
-    Ppxlib.Ast_builder.Default.psig_extension ext [] ~loc
-  in
+    extractor expander
+
+let import_declaration_rule =
+  Ppxlib.Context_free.Rule.extension import_extension
+
+let type_extractor_intf =
+  Ppxlib.Ast_pattern.(
+    pstr (pstr_type __ __ ^:: nil)
+    ||| psig (psig_type __ __ ^:: nil)
+    |> map2 ~f:(fun rec_flag type_decl -> Type_decl (rec_flag, type_decl)) )
+
+let module_type_extractor_intf =
+  Ppxlib.Ast_pattern.(
+    psig (psig_modtype __ ^:: nil)
+    ||| pstr (pstr_modtype __ ^:: nil)
+    |> map1 ~f:(fun modtype -> Module_type_decl modtype) )
+
+let extractor_intf =
+  Ppxlib.Ast_pattern.(type_extractor_intf ||| module_type_extractor_intf)
+
+let expander_intf ~(ctxt : Ppxlib.Expansion_context.Extension.t)
+    (payload : extracted_payload) =
   match payload with
-  | Parsetree.PStr [{pstr_desc = Pstr_type (rec_flag, type_decls); _}]
-   |Parsetree.PSig [{psig_desc = Psig_type (rec_flag, type_decls); _}] ->
+  | Type_decl (rec_flag, type_decls) ->
     type_declaration_expand_intf ~ctxt rec_flag type_decls
-  | Parsetree.PStr [{pstr_desc = Pstr_modtype modtype_decl; _}]
-   |Parsetree.PSig [{psig_desc = Psig_modtype modtype_decl; _}] ->
+  | Module_type_decl modtype_decl ->
     module_declaration_expand_intf ~ctxt modtype_decl
-  | Parsetree.PStr [{pstr_desc = _; _}] | Parsetree.PSig [{psig_desc = _; _}] ->
-    return_error
-      "[%%import] Expected a type declaration or a module type declaration"
-  | Parsetree.PStr (_ :: _) | Parsetree.PSig (_ :: _) ->
-    return_error
-      "[%%import] Expected exactly one item in the structure or signature, but \
-       found multiple items"
-  | Parsetree.PStr [] | Parsetree.PSig [] ->
-    return_error
-      "[%%import] Expected exactly one item in the structure or signature, but \
-       found none"
-  | Parsetree.PTyp _ ->
-    return_error
-      "[%%import] Type pattern (PTyp) is not supported, only type and module \
-       type declarations are allowed"
-  | Parsetree.PPat (_, _) ->
-    return_error
-      "[%%import] Pattern (PPat) is not supported, only type and module type \
-       declarations are allowed"
-
-let type_declaration_extension_intf =
-  Ppxlib.Extension.V3.declare "import" Ppxlib.Extension.Context.signature_item
-    Ppxlib.Ast_pattern.(__)
-    type_declaration_expander_intf
 
-let type_declaration_rule =
-  Ppxlib.Context_free.Rule.extension type_declaration_extension
+let import_extension_intf =
+  Ppxlib.Extension.V3.declare "import" Ppxlib.Extension.Context.signature_item
+    extractor_intf expander_intf
 
-let type_declaration_rule_intf =
-  Ppxlib.Context_free.Rule.extension type_declaration_extension_intf
+let import_declaration_rule_intf =
+  Ppxlib.Context_free.Rule.extension import_extension_intf
 
 let () =
   Ppxlib.Driver.V2.register_transformation
-    ~rules:[type_declaration_rule; type_declaration_rule_intf]
+    ~rules:[import_declaration_rule; import_declaration_rule_intf]
     "ppx_import"
diff --git a/src_test/ppx_deriving/errors/run.t b/src_test/ppx_deriving/errors/run.t
@@ -116,8 +116,7 @@ It's been fixed for later versions in https://github.com/ocaml/ocaml/pull/8541
   1 | [%%import:
   2 | type b = int
   3 | type a = string]
-  Error: [%%import] Expected exactly one item in the structure or signature,
-         but found multiple items
+  Error: [] expected
 
 Ptyp
   $ cat >test.ml <<EOF
@@ -128,8 +127,7 @@ Ptyp
   File "test.ml", line 1, characters 0-18:
   1 | [%%import: string]
       ^^^^^^^^^^^^^^^^^^
-  Error: [%%import] Type pattern (PTyp) is not supported, only type and module
-         type declarations are allowed
+  Error: PStr expected
   [1]
 
 Inline module type declaration

diff --git a/src_test/ppx_deriving/errors_lte_407/run.t b/src_test/ppx_deriving/errors_lte_407/run.t
@@ -112,9 +112,7 @@ Ptyp
   File "test.ml", line 1, characters 0-18:
   1 | [%%import: string]
       ^^^^^^^^^^^^^^^^^^
-  Error: [%%import] Invalid extension usage. [%%import] only supports structure
-         items, signatures or type declarations, but a type pattern (PTyp) was
-         found.
+  Error: PStr expected
   [1]
 
 Inline module type declaration