generate modules inside anonymous prelude module (#10)

* generate modules inside anonymous prelude module

* proper order of prelude struct items

Author: zazedd

MIKMATCH.md

@@ -88,6 +88,20 @@ let do_something = function%mikmatch
   | _ -> ...
 ```
 
+You can also reference patterns from modules and vice versa:
+```ocaml
+module Patterns = struct
+  let%mikmatch hex = {| ['0'-'9' 'a'-'f']+ |}
+end
+
+let%mikmatch hex_with_prefix = {| "0x" Patterns.hex |}
+
+module MorePatterns = struct
+  let%mikmatch multiple_hexes = {| Patterns.hex+ |}
+  let%mikmatch multiple_prefixed_hexes = {| hex_with_prefix+ |}
+end
+```
+
 ### Variable capture
 ```ocaml
 let%mikmatch num = {| digit+ |}

src/ppx_mikmatch.ml

@@ -17,21 +17,23 @@
 open Ppxlib
 open Ast_builder.Default
 
+let pvar ~loc name = ppat_var ~loc { txt = name; loc }
+let evar ~loc name = pexp_ident ~loc { txt = Lident name; loc }
+let rec list_take n = function [] -> [] | x :: xs -> if n <= 0 then [] else x :: list_take (n - 1) xs
+let list_is_empty = function [] -> true | _ -> false
+
 let make_alias_binding ~loc ~var_name =
-  let warning_attr =
-    attribute ~loc ~name:{ txt = "warning"; loc } ~payload:(PStr [ { pstr_desc = Pstr_eval (estring ~loc "-32", []); pstr_loc = loc } ])
-  in
-  {
-    pvb_pat = ppat_var ~loc { txt = var_name; loc };
-    pvb_expr = pexp_ident ~loc { txt = Lident var_name; loc };
-    pvb_constraint = None;
-    pvb_attributes = [ warning_attr ];
-    pvb_loc = loc;
-  }
+  let pat = pvar ~loc var_name in
+  let expr = evar ~loc var_name in
+  match [%stri let[@warning "-32"] [%p pat] = [%e expr]] with { pstr_desc = Pstr_value (_, [ vb ]); _ } -> vb | _ -> assert false
+
+type binding_location =
+  | TopLevel
+  | InModule of string list (* module path *)
 
 let transformation =
   object (self)
-    inherit [value_binding list] Ast_traverse.fold_map as super
+    inherit [(binding_location * value_binding) list] Ast_traverse.fold_map as super
 
     method! structure_item item acc =
       match item.pstr_desc with
@@ -63,14 +65,14 @@ let transformation =
                   let type_name = td.ptype_name.txt in
                   let items, binding = Transformations.transform_type ~loc rec_flag type_name pattern_str td in
                   let alias = pstr_value ~loc Nonrecursive [ make_alias_binding ~loc ~var_name:type_name ] in
-                  (alias :: items_acc) @ items, binding :: bindings_acc
+                  (alias :: items_acc) @ items, (TopLevel, binding) :: bindings_acc
                 | _ -> items_acc, bindings_acc)
               ([], acc) type_decls
           in
 
           let wrapped = pstr_include ~loc:item.pstr_loc (include_infos ~loc:item.pstr_loc (pmod_structure ~loc:item.pstr_loc all_items)) in
           wrapped, all_bindings)
-      (* let%mikmatch x = {|some regex|}*)
+      (* let%mikmatch ... = {| ... |}*)
       | Pstr_extension (({ txt = "mikmatch"; _ }, PStr [ { pstr_desc = Pstr_value (rec_flag, vbs); _ } ]), _) ->
         let processed_vbs, collected_bindings =
           List.fold_left
@@ -80,19 +82,21 @@ let transformation =
               | Ppat_var { txt = var_name; _ }, Pexp_constant (Pconst_string (_, loc, _)) ->
                 let binding = Transformations.transform_let ~loc vb in
                 let alias = make_alias_binding ~loc ~var_name in
-                alias :: vbs_acc, binding :: bindings_acc
+                alias :: vbs_acc, (TopLevel, binding) :: bindings_acc
               (* destructuring - let%mikmatch {|/pattern/|} = expr *)
               | Ppat_constant (Pconst_string (pattern_str, _, _)), _ ->
                 let new_vb, new_bindings = Transformations.transform_destructuring_let ~loc:vb.pvb_loc pattern_str vb.pvb_expr in
-                new_vb :: vbs_acc, new_bindings @ bindings_acc
+                new_vb :: vbs_acc, List.map (fun b -> TopLevel, b) new_bindings @ bindings_acc
               | _ -> vbs_acc, bindings_acc
             end
             ([], acc) vbs
         in
 
         let new_item = { item with pstr_desc = Pstr_value (rec_flag, List.rev processed_vbs) } in
         new_item, collected_bindings
-      (* let x = expression (which might contain %mik/%pcre) *)
+      (* let ... = expression (which might contain %mikmatch)
+         e.g. let ... = {%mikmatch| ... |}
+      *)
       | Pstr_value (rec_flag, vbs) ->
         let processed_vbs, collected_bindings =
           List.fold_left
@@ -105,7 +109,7 @@ let transformation =
                 let alias =
                   match vb.pvb_pat.ppat_desc with Ppat_var { txt = var_name; loc } -> make_alias_binding ~loc ~var_name | _ -> new_vb
                 in
-                alias :: vbs_acc, binding :: bindings_acc
+                alias :: vbs_acc, (TopLevel, binding) :: bindings_acc
               | _ ->
                 let new_expr, new_bindings = self#expression vb.pvb_expr bindings_acc in
                 let new_vb = { vb with pvb_expr = new_expr } in
@@ -114,6 +118,33 @@ let transformation =
         in
         let new_item = { item with pstr_desc = Pstr_value (rec_flag, List.rev processed_vbs) } in
         new_item, collected_bindings
+      (* module M = struct ... end which might contain %mikmatch defns *)
+      | Pstr_module { pmb_name = { txt = Some mod_name; _ } as name; pmb_expr; pmb_attributes; pmb_loc } -> begin
+        match pmb_expr.pmod_desc with
+        | Pmod_structure mod_items ->
+          let mod_items', mod_bindings = self#structure mod_items [] in
+          if mod_bindings = [] then super#structure_item item acc
+          else (
+            let tagged_bindings =
+              List.map
+                (fun (loc, vb) -> match loc with InModule path -> InModule (mod_name :: path), vb | TopLevel -> InModule [ mod_name ], vb)
+                mod_bindings
+            in
+
+            (* include the module from the prelude struct, but only if we tagged bingings, then keep other original items *)
+            let include_item =
+              pstr_include ~loc:pmb_loc (include_infos ~loc:pmb_loc (pmod_ident ~loc:pmb_loc { txt = Lident mod_name; loc = pmb_loc }))
+            in
+            let new_items = if list_is_empty tagged_bindings then mod_items' else include_item :: mod_items' in
+
+            let alias_module =
+              pstr_module ~loc:pmb_loc { pmb_name = name; pmb_expr = pmod_structure ~loc:pmb_loc new_items; pmb_attributes; pmb_loc }
+            in
+            alias_module, tagged_bindings @ acc)
+        | _ ->
+          (* other module types, no transformation needed *)
+          super#structure_item item acc
+      end
       | _ -> super#structure_item item acc
 
     method! expression e_ext acc =
@@ -130,21 +161,21 @@ let transformation =
         begin match e.pexp_desc with
         | Pexp_function ([], _, Pfunction_cases (cases, _, _)) ->
           let cases, binding = Transformations.transform_cases ~loc cases in
-          [%expr fun _ppx_mikmatch_v -> [%e cases]], binding @ acc
+          [%expr fun _ppx_mikmatch_v -> [%e cases]], List.map (fun b -> TopLevel, b) binding @ acc
         | Pexp_match (e, cases) ->
           let cases, binding = Transformations.transform_cases ~loc cases in
           ( [%expr
               let _ppx_mikmatch_v = [%e e] in
               [%e cases]],
-            binding @ acc )
+            List.map (fun b -> TopLevel, b) binding @ acc )
         | Pexp_let (rec_flag, vbs, body) ->
           let processed_vbs, new_bindings =
             List.fold_left
               (fun (vbs_acc, bindings_acc) vb ->
                 match vb.pvb_pat.ppat_desc, vb.pvb_expr.pexp_desc with
                 | Ppat_constant (Pconst_string (pattern_str, _, _)), _ ->
                   let new_vb, new_bindings = Transformations.transform_destructuring_let ~loc:vb.pvb_loc pattern_str vb.pvb_expr in
-                  new_vb :: vbs_acc, new_bindings @ bindings_acc
+                  new_vb :: vbs_acc, List.map (fun b -> TopLevel, b) new_bindings @ bindings_acc
                 | _ ->
                   Util.error ~loc
                     "[%%pcre] and [%%mikmatch] only apply to match, function, global let declarations of strings, and let destructuring.")
@@ -157,9 +188,13 @@ let transformation =
         end
       (* match smth with | {%mikmatch|some regex|} -> ...*)
       | Pexp_match (matched_expr, cases) when has_ext_case cases ->
-        Transformations.transform_mixed_match ~loc:e_ext.pexp_loc ~matched_expr cases acc
+        let plain_acc = List.map snd acc in
+        let expr, bindings = Transformations.transform_mixed_match ~loc:e_ext.pexp_loc ~matched_expr cases plain_acc in
+        expr, List.map (fun b -> TopLevel, b) bindings
       | Pexp_function (params, constraint_, Pfunction_cases (cases, _, _)) when has_ext_case cases ->
-        let transformed, acc = Transformations.transform_mixed_match ~loc:e_ext.pexp_loc cases acc in
+        let plain_acc = List.map snd acc in
+        let transformed, bindings = Transformations.transform_mixed_match ~loc:e_ext.pexp_loc cases plain_acc in
+        let acc = List.map (fun b -> TopLevel, b) bindings in
         begin match params with
         | [] -> transformed, acc
         | _ -> { e_ext with pexp_desc = Pexp_function (params, constraint_, Pfunction_body transformed) }, acc
@@ -187,14 +222,89 @@ let impl str =
   match rev_bindings with
   | [] -> str
   | _ -> begin
-    let loc = match List.hd (List.rev rev_bindings) with { pvb_loc; _ } -> pvb_loc in
-    let struct_items =
-      rev_bindings
-      |> List.rev
-      |> List.fold_left
-           (fun acc binding -> acc @ [%str let[@warning "-32"] [%p binding.pvb_pat] = [%e binding.pvb_expr]])
-           [%str [%%i pstr_value ~loc Nonrecursive [ dispatch_function_binding ~loc ]]]
+    let loc = match List.hd (List.rev rev_bindings) with _, { pvb_loc; _ } -> pvb_loc in
+
+    let bindings = List.rev rev_bindings in
+
+    (* process all bindings in source order, build modules as needed *)
+    let rec emit_in_order remaining =
+      match remaining with
+      | [] -> []
+      | (TopLevel, vb) :: rest ->
+        (* emit top-level binding immediately *)
+        let items = [%str let[@warning "-32"] [%p vb.pvb_pat] = [%e vb.pvb_expr]] in
+        items @ emit_in_order rest
+      | (InModule path, vb) :: rest ->
+        let root = List.hd path in
+        (* collect all consecutive bindings for this root module *)
+        let same_root, different =
+          let rec collect acc = function
+            | ((InModule p, _) as b) :: rest when List.hd p = root -> collect (b :: acc) rest
+            | rest -> List.rev acc, rest
+          in
+          collect [ InModule path, vb ] rest
+        in
+        let mod_items = build_module_tree ~loc root same_root in
+        mod_items @ emit_in_order different
+    and build_module_tree ~loc root module_bindings =
+      let by_path = Hashtbl.create 16 in
+      List.iter
+        begin fun binding ->
+          match binding with
+          | InModule path, vb ->
+            let existing = try Hashtbl.find by_path path with Not_found -> [] in
+            Hashtbl.replace by_path path (vb :: existing)
+          | TopLevel, _ -> assert false
+        end
+        module_bindings;
+
+      let rec build_at_path current_path =
+        let direct_bindings = try Hashtbl.find by_path current_path |> List.rev with Not_found -> [] in
+        let direct_items = List.concat_map (fun vb -> [%str let[@warning "-32"] [%p vb.pvb_pat] = [%e vb.pvb_expr]]) direct_bindings in
+
+        let child_modules =
+          Hashtbl.fold
+            begin fun path _ acc ->
+              if
+                List.length path = List.length current_path + 1
+                && List.for_all2 ( = ) current_path (list_take (List.length current_path) path)
+              then List.nth path (List.length current_path) :: acc
+              else acc
+            end
+            by_path []
+          |> List.sort_uniq compare
+        in
+
+        (* build nested modules *)
+        let nested_items =
+          List.concat_map
+            (fun child_name ->
+              let child_path = current_path @ [ child_name ] in
+              let child_items = build_at_path child_path in
+              [
+                pstr_module ~loc
+                  {
+                    pmb_name = { txt = Some child_name; loc };
+                    pmb_expr = pmod_structure ~loc child_items;
+                    pmb_attributes = [];
+                    pmb_loc = loc;
+                  };
+              ])
+            child_modules
+        in
+
+        direct_items @ nested_items
+      in
+
+      let mod_body = build_at_path [ root ] in
+      [
+        pstr_module ~loc
+          { pmb_name = { txt = Some root; loc }; pmb_expr = pmod_structure ~loc mod_body; pmb_attributes = []; pmb_loc = loc };
+      ]
     in
+
+    let struct_items = [%str [%%i pstr_value ~loc Nonrecursive [ dispatch_function_binding ~loc ]]] @ emit_in_order bindings in
+
     let mod_expr = pmod_structure ~loc struct_items in
     [%str open [%m mod_expr]] @ str
   end