unikernel/duniverse/ppxlib/traverse/dune

@@ -0,0 +1,9 @@
+(library
+ (name ppxlib_traverse)
+ (public_name ppxlib.traverse)
+ (kind ppx_deriver)
+ (flags
+  (:standard -safe-string))
+ (libraries ppxlib ppxlib_ast ppxlib_traverse_builtins stdppx stdlib-shims)
+ (preprocess
+  (pps ppxlib_metaquot)))

unikernel/duniverse/ppxlib/traverse/ppxlib_traverse.ml

@@ -0,0 +1,631 @@
+open Stdppx
+open Ppxlib
+open Ast_builder.Default
+
+let alphabet =
+  Array.init
+    (Char.code 'z' - Char.code 'a' + 1)
+    ~f:(fun i -> String.make 1 (Char.chr (i + Char.code 'a')))
+
+let vars_of_list ~get_loc l =
+  List.mapi l ~f:(fun i x -> { txt = alphabet.(i); loc = get_loc x })
+
+let evar_of_var { txt; loc } = evar ~loc txt
+let pvar_of_var { txt; loc } = pvar ~loc txt
+let tvar_of_var { txt; loc } = ptyp_var ~loc txt
+let evars_of_vars = List.map ~f:evar_of_var
+let pvars_of_vars = List.map ~f:pvar_of_var
+let tvars_of_vars = List.map ~f:tvar_of_var
+
+let fst_expr ~loc expr =
+  match expr.pexp_desc with
+  | Pexp_tuple [ fst; _ ] -> fst
+  | _ -> [%expr Stdlib.fst [%e expr]]
+
+let snd_expr ~loc expr =
+  match expr.pexp_desc with
+  | Pexp_tuple [ _; snd ] -> snd
+  | _ -> [%expr Stdlib.snd [%e expr]]
+
+let methods_of_class_exn = function
+  | {
+      pstr_desc =
+        Pstr_class
+          [
+            {
+              pci_expr = { pcl_desc = Pcl_structure { pcstr_fields = l; _ }; _ };
+              _;
+            };
+          ];
+      _;
+    } ->
+      l
+  | _ -> assert false
+
+module Backends = struct
+  class reconstructors =
+    object
+      method record ~loc flds = pexp_record ~loc flds None
+
+      method construct ~loc id args =
+        pexp_construct ~loc id
+          (match args with [] -> None | _ -> Some (pexp_tuple ~loc args))
+
+      method tuple ~loc es = pexp_tuple ~loc es
+    end
+
+  class type what = object
+    method name : string
+    inherit reconstructors
+
+    method class_params :
+      loc:Location.t -> (core_type * (variance * injectivity)) list
+
+    method virtual_methods : loc:Location.t -> class_field list
+    method apply : loc:Location.t -> expression -> expression list -> expression
+    method abstract : loc:Location.t -> pattern -> expression -> expression
+
+    (* Basic combinator type *)
+    method typ : loc:Location.t -> core_type -> core_type
+    method any : loc:Location.t -> expression
+
+    method combine :
+      loc:Location.t ->
+      (string loc * expression) list ->
+      reconstruct:expression ->
+      expression
+  end
+
+  let mapper : what =
+    object
+      method name = "map"
+      inherit reconstructors
+      method class_params ~loc:_ = []
+      method virtual_methods ~loc:_ = []
+      method apply ~loc expr args = eapply ~loc expr args
+      method abstract ~loc patt expr = pexp_fun ~loc Nolabel None patt expr
+      method typ ~loc ty = ptyp_arrow ~loc Nolabel ty ty
+      method any ~loc = [%expr fun x -> x]
+
+      method combine ~loc combinators ~reconstruct =
+        List.fold_right combinators ~init:reconstruct ~f:(fun (v, expr) acc ->
+            pexp_let ~loc Nonrecursive
+              [ value_binding ~loc ~pat:(pvar_of_var v) ~expr ]
+              acc)
+    end
+
+  let iterator : what =
+    object
+      method name = "iter"
+      inherit reconstructors
+      method class_params ~loc:_ = []
+      method virtual_methods ~loc:_ = []
+      method apply ~loc expr args = eapply ~loc expr args
+      method abstract ~loc patt expr = pexp_fun ~loc Nolabel None patt expr
+      method typ ~loc ty = [%type: [%t ty] -> unit]
+      method any ~loc = [%expr fun _ -> ()]
+
+      method combine ~loc combinators ~reconstruct:_ =
+        match List.rev combinators with
+        | [] -> [%expr ()]
+        | (_, expr) :: rest ->
+            List.fold_left rest ~init:expr ~f:(fun acc (_v, expr) ->
+                pexp_sequence ~loc expr acc)
+    end
+
+  let folder : what =
+    object
+      method name = "fold"
+      inherit reconstructors
+
+      method class_params ~loc =
+        [ (ptyp_var ~loc "acc", (NoVariance, NoInjectivity)) ]
+
+      method virtual_methods ~loc:_ = []
+      method apply ~loc expr args = eapply ~loc expr (args @ [ evar ~loc "acc" ])
+
+      method abstract ~loc patt expr =
+        eabstract ~loc [ patt; pvar ~loc "acc" ] expr
+
+      method typ ~loc ty = [%type: [%t ty] -> 'acc -> 'acc]
+      method any ~loc = [%expr fun _ acc -> acc]
+
+      method combine ~loc combinators ~reconstruct:_ =
+        match combinators with
+        | [ (_, expr) ] -> expr
+        | _ ->
+            List.fold_right combinators ~init:[%expr acc]
+              ~f:(fun (_v, expr) acc ->
+                [%expr
+                  let acc = [%e expr] in
+                  [%e acc]])
+    end
+
+  let fold_mapper : what =
+    object
+      method name = "fold_map"
+      inherit reconstructors
+
+      method class_params ~loc =
+        [ (ptyp_var ~loc "acc", (NoVariance, NoInjectivity)) ]
+
+      method virtual_methods ~loc:_ = []
+      method apply ~loc expr args = eapply ~loc expr (args @ [ evar ~loc "acc" ])
+
+      method abstract ~loc patt expr =
+        eabstract ~loc [ patt; pvar ~loc "acc" ] expr
+
+      method typ ~loc ty = [%type: [%t ty] -> 'acc -> [%t ty] * 'acc]
+      method any ~loc = [%expr fun x acc -> (x, acc)]
+
+      method combine ~loc combinators ~reconstruct =
+        List.fold_right combinators
+          ~init:[%expr [%e reconstruct], acc]
+          ~f:(fun (v, expr) acc ->
+            [%expr
+              let [%p pvar_of_var v], acc = [%e expr] in
+              [%e acc]])
+    end
+
+  exception Found
+
+  let uses_var var =
+    let iter =
+      object
+        inherit Ast_traverse.iter as super
+
+        method! expression_desc =
+          function
+          | Pexp_ident { txt = Lident id; _ } when String.equal id var ->
+              raise_notrace Found
+          | e -> super#expression_desc e
+      end
+    in
+    fun e ->
+      try
+        iter#expression e;
+        false
+      with Found -> true
+
+  let mapper_with_context : what =
+    let uses_ctx = uses_var "ctx" in
+    object
+      method name = "map_with_context"
+      inherit reconstructors
+
+      method class_params ~loc =
+        [ (ptyp_var ~loc "ctx", (NoVariance, NoInjectivity)) ]
+
+      method virtual_methods ~loc:_ = []
+      method apply ~loc expr args = eapply ~loc expr (evar ~loc "ctx" :: args)
+
+      method abstract ~loc patt expr =
+        if uses_ctx expr then eabstract ~loc [ pvar ~loc "ctx"; patt ] expr
+        else eabstract ~loc [ pvar ~loc "_ctx"; patt ] expr
+
+      method typ ~loc ty = [%type: 'ctx -> [%t ty] -> [%t ty]]
+      method any ~loc = [%expr fun _ctx x -> x]
+
+      method combine ~loc combinators ~reconstruct =
+        List.fold_right combinators ~init:reconstruct ~f:(fun (v, expr) acc ->
+            [%expr
+              let [%p pvar_of_var v] = [%e expr] in
+              [%e acc]])
+    end
+
+  let string_of_lid id = String.concat ~sep:"." (Longident.flatten_exn id)
+
+  let lifter : what =
+    object
+      method name = "lift"
+
+      method class_params ~loc =
+        [ (ptyp_var ~loc "res", (NoVariance, NoInjectivity)) ]
+
+      method virtual_methods ~loc =
+        methods_of_class_exn
+          [%stri
+            class virtual blah =
+              object
+                method virtual record : (string * 'res) list -> 'res
+                method virtual constr : string -> 'res list -> 'res
+                method virtual tuple : 'res list -> 'res
+                method virtual other : 'a. 'a -> 'res
+              end]
+
+      method apply ~loc expr args = eapply ~loc expr args
+      method abstract ~loc patt expr = pexp_fun ~loc Nolabel None patt expr
+      method typ ~loc ty = [%type: [%t ty] -> 'res]
+      method any ~loc = [%expr self#other]
+
+      method combine ~loc combinators ~reconstruct =
+        List.fold_right combinators ~init:reconstruct ~f:(fun (v, expr) acc ->
+            pexp_let ~loc Nonrecursive
+              [ value_binding ~loc ~pat:(pvar_of_var v) ~expr ]
+              acc)
+
+      method record ~loc flds =
+        let flds =
+          elist ~loc
+            (List.map flds ~f:(fun (lab, e) ->
+                 pexp_tuple
+                   ~loc:{ lab.loc with loc_end = e.pexp_loc.loc_end }
+                   [ estring ~loc:lab.loc (string_of_lid lab.txt); e ]))
+        in
+        [%expr self#record [%e flds]]
+
+      method construct ~loc id args =
+        let args = elist ~loc args in
+        [%expr
+          self#constr [%e estring ~loc:id.loc (string_of_lid id.txt)] [%e args]]
+
+      method tuple ~loc es = [%expr self#tuple [%e elist ~loc es]]
+    end
+
+  let lift_mapper_with_context : what =
+    let uses_ctx = uses_var "ctx" in
+    object
+      method name = "lift_map_with_context"
+
+      method class_params ~loc =
+        [
+          (ptyp_var ~loc "ctx", (NoVariance, NoInjectivity));
+          (ptyp_var ~loc "res", (NoVariance, NoInjectivity));
+        ]
+
+      method virtual_methods ~loc =
+        methods_of_class_exn
+          [%stri
+            class virtual blah =
+              object
+                method virtual record : 'ctx -> (string * 'res) list -> 'res
+                method virtual constr : 'ctx -> string -> 'res list -> 'res
+                method virtual tuple : 'ctx -> 'res list -> 'res
+                method virtual other : 'a. 'ctx -> 'a -> 'res
+              end]
+
+      method apply ~loc expr args = eapply ~loc expr (evar ~loc "ctx" :: args)
+
+      method abstract ~loc patt expr =
+        let ctx_pat =
+          if uses_ctx expr then pvar ~loc "ctx" else pvar ~loc "_ctx"
+        in
+        eabstract ~loc [ ctx_pat; patt ] expr
+
+      method typ ~loc ty = [%type: 'ctx -> [%t ty] -> [%t ty] * 'res]
+      method any ~loc = [%expr fun ctx x -> (x, self#other ctx x)]
+
+      method combine ~loc combinators ~reconstruct =
+        List.fold_right combinators ~init:reconstruct ~f:(fun (v, expr) acc ->
+            pexp_let ~loc Nonrecursive
+              [ value_binding ~loc ~pat:(pvar_of_var v) ~expr ]
+              acc)
+
+      method record ~loc flds =
+        let record =
+          pexp_record ~loc
+            (List.map flds ~f:(fun (lab, e) -> (lab, fst_expr ~loc e)))
+            None
+        in
+        let flds =
+          elist ~loc
+            (List.map flds ~f:(fun (lab, e) ->
+                 pexp_tuple
+                   ~loc:{ lab.loc with loc_end = e.pexp_loc.loc_end }
+                   [
+                     estring ~loc:lab.loc (string_of_lid lab.txt);
+                     snd_expr ~loc e;
+                   ]))
+        in
+        [%expr [%e record], self#record ctx [%e flds]]
+
+      method construct ~loc id args =
+        let constr =
+          pexp_construct ~loc id
+            (pexp_tuple_opt ~loc (List.map args ~f:(fst_expr ~loc)))
+        in
+        let res =
+          let args = elist ~loc (List.map args ~f:(snd_expr ~loc)) in
+          [%expr
+            self#constr ctx
+              [%e estring ~loc:id.loc (string_of_lid id.txt)]
+              [%e args]]
+        in
+        [%expr [%e constr], [%e res]]
+
+      method tuple ~loc es =
+        let tuple = pexp_tuple ~loc (List.map es ~f:(fst_expr ~loc)) in
+        let res =
+          [%expr
+            self#tuple ctx [%e elist ~loc (List.map es ~f:(snd_expr ~loc))]]
+        in
+        [%expr [%e tuple], [%e res]]
+    end
+
+  let all =
+    [
+      mapper;
+      iterator;
+      folder;
+      fold_mapper;
+      mapper_with_context;
+      lifter;
+      lift_mapper_with_context;
+    ]
+end
+
+type what = Backends.what
+
+let mapper_type ~(what : what) ~loc type_name params =
+  let vars = vars_of_list params ~get_loc:(fun t -> t.ptyp_loc) in
+  let params = tvars_of_vars vars in
+  let ty = ptyp_constr ~loc type_name params in
+  let ty =
+    List.fold_right params ~init:(what#typ ~loc ty) ~f:(fun param ty ->
+        let loc = param.ptyp_loc in
+        ptyp_arrow ~loc Nolabel (what#typ ~loc param) ty)
+  in
+  ptyp_poly ~loc vars ty
+
+let constrained_mapper ~(what : what) ?(is_gadt = false) mapper td =
+  let vars = vars_of_list td.ptype_params ~get_loc:(fun (t, _) -> t.ptyp_loc) in
+  let make_type params =
+    let loc = td.ptype_loc in
+    let ty = ptyp_constr ~loc (Loc.map td.ptype_name ~f:lident) params in
+    List.fold_right params ~init:(what#typ ~loc:td.ptype_loc ty)
+      ~f:(fun param ty ->
+        let loc = param.ptyp_loc in
+        ptyp_arrow ~loc Nolabel (what#typ ~loc param) ty)
+  in
+  let typ =
+    let loc = td.ptype_loc in
+    ptyp_poly ~loc vars (make_type (tvars_of_vars vars))
+  in
+  let mapper =
+    if false || is_gadt then
+      let typs =
+        List.map vars ~f:(fun v ->
+            ptyp_constr ~loc:v.loc (Loc.map v ~f:lident) [])
+      in
+      List.fold_right vars
+        ~init:(pexp_constraint ~loc:mapper.pexp_loc mapper (make_type typs))
+        ~f:(fun v e -> pexp_newtype ~loc:v.loc v e)
+    else mapper
+  in
+  pexp_poly ~loc:mapper.pexp_loc mapper (Some typ)
+
+let mangle_type_name lid =
+  let rec mangled_parts lid ~suffix =
+    match lid with
+    | Lident s -> String.lowercase_ascii s :: suffix
+    | Ldot (lid, s) ->
+        mangled_parts lid ~suffix:("__" :: String.lowercase_ascii s :: suffix)
+    | Lapply (a, b) ->
+        mangled_parts a ~suffix:("_'" :: mangled_parts b ~suffix:("'" :: suffix))
+  in
+  mangled_parts lid ~suffix:[] |> String.concat ~sep:""
+
+let rec type_expr_mapper ~(what : what) te =
+  let loc = te.ptyp_loc in
+  match te.ptyp_desc with
+  | Ptyp_var s -> evar ~loc ("_" ^ s)
+  | Ptyp_tuple tes ->
+      let vars = vars_of_list tes ~get_loc:(fun t -> t.ptyp_loc) in
+      let deconstruct = ppat_tuple ~loc (pvars_of_vars vars) in
+      let reconstruct = what#tuple ~loc (evars_of_vars vars) in
+      let mappers = map_variables ~what vars tes in
+      what#abstract ~loc deconstruct (what#combine ~loc mappers ~reconstruct)
+  | Ptyp_constr (path, params) -> (
+      let map =
+        pexp_send ~loc (evar ~loc "self") (Loc.map path ~f:mangle_type_name)
+      in
+      match params with
+      | [] -> map
+      | _ ->
+          eapply ~loc map
+            (List.map params ~f:(fun te -> type_expr_mapper ~what te)))
+  | _ -> what#any ~loc
+
+and map_variables ~(what : what) vars tes =
+  List.map2 tes vars ~f:(fun te var ->
+      ( var,
+        what#apply ~loc:te.ptyp_loc
+          (type_expr_mapper ~what te)
+          [ evar_of_var var ] ))
+
+let gen_record' ~(what : what) ~loc lds =
+  let vars = List.map lds ~f:(fun ld -> ld.pld_name) in
+  let deconstruct =
+    ppat_record ~loc
+      (List.map vars ~f:(fun v -> (Loc.map v ~f:lident, pvar_of_var v)))
+      Closed
+  in
+  let reconstruct =
+    what#record ~loc
+      (List.map vars ~f:(fun v -> (Loc.map v ~f:lident, evar_of_var v)))
+  in
+  let mappers =
+    map_variables ~what vars (List.map lds ~f:(fun ld -> ld.pld_type))
+  in
+  (deconstruct, reconstruct, mappers)
+
+let gen_record ~(what : what) ~loc lds =
+  let deconstruct, reconstruct, mappers = gen_record' ~what lds ~loc in
+  what#abstract ~loc deconstruct (what#combine ~loc mappers ~reconstruct)
+
+let is_constant_constructor cd =
+  match cd.pcd_args with Pcstr_tuple [] -> true | _ -> false
+
+let erase_type_variables =
+  object
+    inherit Ast_traverse.map as super
+
+    method! core_type_desc =
+      function Ptyp_var _ -> Ptyp_any | x -> super#core_type_desc x
+  end
+
+let gen_variant ~(what : what) ~loc cds =
+  if
+    String.( <> ) what#name "lift"
+    && List.for_all cds ~f:is_constant_constructor
+  then what#any ~loc
+  else
+    let cases =
+      List.map cds ~f:(fun cd ->
+          let cstr = Loc.map cd.pcd_name ~f:lident in
+          let loc = cd.pcd_loc in
+          let args =
+            match cd.pcd_res with
+            | None -> cd.pcd_args
+            | Some _ ->
+                (* This is a big sur-approximation but it's enough for our only use of GADTs
+                   in ppx_custom_format *)
+                erase_type_variables#constructor_arguments cd.pcd_args
+          in
+          match args with
+          | Pcstr_tuple args ->
+              let vars = vars_of_list args ~get_loc:(fun t -> t.ptyp_loc) in
+              let deconstruct =
+                ppat_construct cstr ~loc
+                  (match vars with
+                  | [] -> None
+                  | _ -> Some (ppat_tuple ~loc (pvars_of_vars vars)))
+              in
+              let reconstruct = what#construct cstr ~loc (evars_of_vars vars) in
+              let mappers = map_variables ~what vars args in
+              case ~lhs:deconstruct
+                ~rhs:(what#combine ~loc mappers ~reconstruct)
+                ~guard:None
+          | Pcstr_record labels ->
+              let deconstruct, reconstruct, mappers =
+                gen_record' ~loc ~what labels
+              in
+              let deconstruct = ppat_construct ~loc cstr (Some deconstruct) in
+              let reconstruct = what#construct ~loc cstr [ reconstruct ] in
+              case ~lhs:deconstruct
+                ~rhs:(what#combine ~loc mappers ~reconstruct)
+                ~guard:None)
+    in
+    what#abstract ~loc (pvar ~loc "x") (pexp_match ~loc (evar ~loc "x") cases)
+
+let gen_mapper ~(what : what) td =
+  let body =
+    let loc = td.ptype_loc in
+    match td.ptype_kind with
+    | Ptype_open -> what#any ~loc
+    | Ptype_record lds -> gen_record ~what lds ~loc
+    | Ptype_variant cds -> gen_variant ~what cds ~loc
+    | Ptype_abstract -> (
+        match td.ptype_manifest with
+        | None -> what#any ~loc
+        | Some te -> type_expr_mapper ~what te)
+  in
+  List.fold_right td.ptype_params ~init:body ~f:(fun (ty, _) acc ->
+      let loc = ty.ptyp_loc in
+      match ty.ptyp_desc with
+      | Ptyp_var s -> pexp_fun ~loc Nolabel None (pvar ~loc ("_" ^ s)) acc
+      | _ -> pexp_fun ~loc Nolabel None (ppat_any ~loc) acc)
+
+let type_deps =
+  let collect =
+    object
+      inherit [int Longident.Map.t] Ast_traverse.fold as super
+
+      method! core_type t acc =
+        let acc =
+          match t.ptyp_desc with
+          | Ptyp_constr (id, vars) ->
+              Longident.Map.add id.txt (List.length vars) acc
+          | _ -> acc
+        in
+        super#core_type t acc
+    end
+  in
+  fun tds ->
+    let empty = Longident.Map.empty in
+    let map =
+      List.fold_left tds ~init:empty ~f:(fun map td ->
+          let map = collect#type_kind td.ptype_kind map in
+          match (td.ptype_kind, td.ptype_manifest) with
+          | Ptype_abstract, Some ty -> collect#core_type ty map
+          | _ -> map)
+    in
+    let map =
+      List.fold_left tds ~init:map ~f:(fun map td ->
+          Longident.Map.remove (Lident td.ptype_name.txt) map)
+    in
+    Longident.Map.bindings map
+
+let filter_virtual_methods ~methods ~virtual_methods =
+  let collect =
+    object
+      inherit [String.Set.t] Ast_traverse.fold as super
+
+      method! expression_desc x acc =
+        match x with
+        | Pexp_send
+            ( _,
+              { txt = ("tuple" | "record" | "constr" | "other") as s; loc = _ }
+            ) ->
+            String.Set.add s acc
+        | _ -> super#expression_desc x acc
+    end
+  in
+  let used = collect#list collect#class_field methods String.Set.empty in
+  List.filter virtual_methods ~f:(fun m ->
+      match m.pcf_desc with
+      | Pcf_method (s, _, _) -> String.Set.mem s.txt used
+      | _ -> false)
+
+let gen_class ~(what : what) ~loc tds =
+  let class_params = what#class_params ~loc in
+  let virtual_methods =
+    List.map (type_deps tds) ~f:(fun (id, arity) ->
+        pcf_method ~loc
+          ( { txt = mangle_type_name id; loc },
+            Public,
+            Cfk_virtual
+              (mapper_type ~what ~loc { txt = id; loc }
+                 (List.init ~len:arity ~f:(fun _ -> ptyp_any ~loc))) ))
+  in
+  let methods =
+    List.map tds ~f:(fun td ->
+        let loc = td.ptype_loc in
+        let mapper = gen_mapper ~what td in
+        let is_gadt =
+          match td.ptype_kind with
+          | Ptype_variant cds ->
+              List.exists cds ~f:(fun cd -> Option.is_some cd.pcd_res)
+          | _ -> false
+        in
+        let mapper = constrained_mapper ~what ~is_gadt mapper td in
+        pcf_method ~loc (td.ptype_name, Public, Cfk_concrete (Fresh, mapper)))
+  in
+  let virtual_methods =
+    filter_virtual_methods ~methods ~virtual_methods:(what#virtual_methods ~loc)
+    @ virtual_methods
+  in
+  let virt = if List.is_empty virtual_methods then Concrete else Virtual in
+  class_infos ~loc ~virt ~params:class_params ~name:{ loc; txt = what#name }
+    ~expr:
+      (pcl_structure ~loc
+         (class_structure
+            ~self:(ppat_var ~loc { txt = "self"; loc })
+            ~fields:(virtual_methods @ methods)))
+
+let gen_str ~(what : what) ~loc ~path:_ (rf, tds) =
+  (match rf with
+  | Nonrecursive ->
+      (* The method name would clash... *)
+      Location.raise_errorf ~loc "ppxlib_traverse doesn't support nonrec"
+  | Recursive -> ());
+  let cl = gen_class ~loc ~what tds in
+  [ pstr_class ~loc:cl.pci_loc [ cl ] ]
+
+let () =
+  let derivers =
+    List.map Backends.all ~f:(fun what ->
+        Deriving.add ("traverse_" ^ what#name)
+          ~str_type_decl:(Deriving.Generator.make_noarg (gen_str ~what)))
+  in
+  Deriving.add_alias "traverse" (List.rev derivers) |> Deriving.ignore