Traverse record declarations and generate typed field accessors

rewrite.ml

@@ -0,0 +1,50 @@
+open Asttypes
+open Parsetree
+open Ast_helper
+
+let ident ~loc name = Location.mkloc (Longident.Lident name) loc
+let value ~loc name = Exp.ident ~loc (ident ~loc name)
+let variable ~loc name = Pat.var ~loc (Location.mkloc name loc)
+let lambda ~loc label pattern body = Exp.fun_ ~loc label None pattern body
+
+let is_marker (name, _) = name.txt = "fieldglass" || name.txt = "lens"
+
+let configured declaration =
+  match List.filter is_marker declaration.ptype_attributes with
+  | [] -> false
+  | [(_, PStr [{ pstr_desc = Pstr_eval
+      ({ pexp_desc = Pexp_ident { txt = Longident.Lident "generate"; _ }; _ }, _); _ }])] -> true
+  | _ -> Location.raise_errorf ~loc:declaration.ptype_loc
+      "fieldglass: expected [@@fieldglass generate]"
+
+let generate declaration =
+  let loc = declaration.ptype_loc in
+  let fields = match declaration.ptype_kind with
+    | Ptype_record fields -> fields
+    | _ -> Location.raise_errorf ~loc "fieldglass: expected a record type"
+  in
+  let source = value ~loc "__fg_source" in
+  let source_pattern = Pat.constraint_ ~loc (variable ~loc "__fg_source")
+      (Typ.constr ~loc (ident ~loc declaration.ptype_name.txt)
+         (List.map (fun _ -> Typ.any ~loc ()) declaration.ptype_params)) in
+  List.map (fun field ->
+    let name = field.pld_name.txt in
+    let getter = lambda ~loc Nolabel source_pattern
+        (Exp.field ~loc source (ident ~loc name)) in
+    Str.value ~loc Nonrecursive [Vb.mk ~loc (variable ~loc name) getter]) fields
+
+let structure mapper items =
+  List.concat (List.map (fun item ->
+    let item = Ast_mapper.default_mapper.structure_item mapper item in
+    match item.pstr_desc with
+    | Pstr_type (recursive, declarations) ->
+        let generated = List.concat (List.map (fun declaration ->
+          if configured declaration then generate declaration else []) declarations) in
+        let declarations = List.map (fun declaration ->
+          { declaration with ptype_attributes =
+              List.filter (fun attribute -> not (is_marker attribute)) declaration.ptype_attributes }) declarations in
+        { item with pstr_desc = Pstr_type (recursive, declarations) } :: generated
+    | _ -> [item]) items)
+
+let () = Ast_mapper.register "fieldglass"
+    (fun _ -> { Ast_mapper.default_mapper with structure })