unikernel/duniverse/ocaml-cstruct/ppx/ppx_cstruct.ml

@@ -0,0 +1,662 @@
+(*
+ * Copyright (c) 2015 Nicolas Ojeda Bar <n.oje.bar@gmail.com>
+ *
+ * Permission to use, copy, modify, and distribute this software for any
+ * purpose with or without fee is hereby granted, provided that the above
+ * copyright notice and this permission notice appear in all copies.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
+ * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
+ * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
+ * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
+ * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
+ * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
+ * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
+ *)
+
+open Printf
+open Ppxlib
+open Ast_helper
+
+module Ast = struct
+  include Ast_builder.Default
+
+  let econstr ~loc tag xs =
+    pexp_construct ~loc (Loc.make ~loc (lident tag)) (
+      match xs with
+      | [] -> None
+      | _ -> Some (pexp_tuple ~loc xs)
+    )
+
+  let pconstr ~loc tag xs =
+    ppat_construct ~loc (Loc.make ~loc (lident tag)) (
+      match xs with
+      | [] -> None
+      | _ -> Some (ppat_tuple ~loc xs)
+    )
+
+  let tconstr ~loc tag xs =
+    ptyp_constr ~loc (Loc.make ~loc (lident tag)) xs
+end
+
+type mode = Big_endian | Little_endian | Host_endian | Bi_endian
+
+type prim =
+  | Char
+  | UInt8
+  | UInt16
+  | UInt32
+  | UInt64
+
+type ty =
+  | Prim of prim
+  | Buffer of prim * int
+
+type raw_field = {
+  name: string;
+  ty: ty;
+  definition_loc: Location.t;
+}
+
+type named_field = {
+  name: string;
+  ty: ty;
+  definition_loc: Location.t;
+  off: int;
+}
+
+type field =
+  | Named_field of named_field
+  | Ignored_field
+
+let field_is_ignored name =
+  String.get name 0 = '_'
+
+type t = {
+  name: string;
+  fields: field list;
+  len: int;
+  endian: mode;
+}
+
+let ty_of_string =
+  function
+  |"char_t"  |"char" -> Some Char
+  |"uint8_t" |"uint8" |"int8" |"int8_t"  -> Some UInt8
+  |"uint16_t"|"uint16"|"int16"|"int16_t" -> Some UInt16
+  |"uint32_t"|"uint32"|"int32"|"int32_t" -> Some UInt32
+  |"uint64_t"|"uint64"|"int64"|"int64_t" -> Some UInt64
+  |_ -> None
+
+let width_of_prim = function
+  | Char -> 1
+  | UInt8 -> 1
+  | UInt16 -> 2
+  | UInt32 -> 4
+  | UInt64 -> 8
+
+let width_of_ty = function
+  | Prim p -> width_of_prim p
+  | Buffer (p, len) -> width_of_prim p * len
+
+let width_of_field f =
+  width_of_ty f.ty
+
+let field_to_string f =
+  let rec string = function
+    |Prim Char -> "char_t"
+    |Prim UInt8 -> "uint8_t"
+    |Prim UInt16 -> "uint16_t"
+    |Prim UInt32 -> "uint32_t"
+    |Prim UInt64 -> "uint64_t"
+    |Buffer (prim, len) -> sprintf "%s[%d]" (string (Prim prim)) len
+  in
+  sprintf "%s %s" (string f.ty) f.name
+
+let loc_err loc fmt = Location.raise_errorf ~loc ("ppx_cstruct: " ^^ fmt)
+
+let parse_field loc name field_type sz =
+  match ty_of_string field_type with
+  |None -> loc_err loc "Unknown type %s" field_type
+  |Some ty -> begin
+    let ty = match ty,sz with
+      |_,None -> Prim ty
+      |prim,Some sz -> Buffer (prim, sz)
+    in
+    { name; ty; definition_loc = loc }
+  end
+
+let check_for_duplicates fields =
+  let module StringSet = Set.Make(String) in
+  let _ : StringSet.t =
+    List.fold_left (fun seen f ->
+        match f with
+        | Ignored_field -> seen
+        | Named_field {name; definition_loc; _} ->
+          if StringSet.mem name seen then
+            loc_err definition_loc "field %s is present several times in this type" name
+          else
+            StringSet.add name seen
+      )
+      StringSet.empty
+      fields
+  in
+  ()
+
+let create_struct loc endian name fields =
+  let endian = match endian with
+    |"little_endian" -> Little_endian
+    |"big_endian" -> Big_endian
+    |"host_endian" -> Host_endian
+    |"bi_endian" -> Bi_endian
+    |_ -> loc_err loc "unknown endian %s, should be little_endian, big_endian, host_endian or bi_endian" endian
+  in
+  let len, fields =
+    List.fold_left (fun (off,acc) ({name; ty; definition_loc}:raw_field) ->
+        let field =
+          if field_is_ignored name then
+            Ignored_field
+          else
+            Named_field { name; ty; off; definition_loc }
+        in
+        let off = width_of_ty ty + off in
+        let acc = acc @ [field] in
+        (off, acc)
+      ) (0,[]) fields
+  in
+  check_for_duplicates fields;
+  { fields; name = name.txt; len; endian }
+
+let ($.) l x = Longident.Ldot (l, x)
+let cstruct_id = Longident.Lident "Cstruct"
+let mode_mod s = function
+  |Big_endian -> cstruct_id$."BE"$.s
+  |Little_endian -> cstruct_id$."LE"$.s
+  |Host_endian -> cstruct_id$."HE"$.s
+  |Bi_endian -> cstruct_id$."BL"$.s
+
+let mode_mod loc x s =
+  Exp.ident ~loc {loc ; txt = mode_mod s x}
+
+type op =
+  | Op_get of named_field
+  | Op_set of named_field
+  | Op_copy of named_field
+  | Op_blit of named_field
+  | Op_sizeof
+  | Op_hexdump
+  | Op_hexdump_to_buffer
+
+let op_name s op =
+  let parts =
+    match op with
+    | Op_get f -> ["get"; s.name; f.name]
+    | Op_set f -> ["set"; s.name; f.name]
+    | Op_copy f -> ["copy"; s.name; f.name]
+    | Op_blit f -> ["blit"; s.name; f.name]
+    | Op_sizeof -> ["sizeof"; s.name]
+    | Op_hexdump -> ["hexdump"; s.name]
+    | Op_hexdump_to_buffer -> ["hexdump"; s.name; "to_buffer"]
+  in
+  String.concat "_" parts
+
+let op_pvar ~loc s op = Ast.pvar ~loc (op_name s op)
+let op_evar ~loc s op = Ast.evar ~loc (op_name s op)
+
+let get_expr loc s f =
+  let m = mode_mod loc s.endian in
+  let num x = Ast.eint ~loc x in
+  match f.ty with
+  |Buffer (_, _) ->
+    let len = width_of_field f in
+    [%expr
+      fun src -> Cstruct.sub src [%e num f.off] [%e num len]
+    ]
+  |Prim prim ->
+    [%expr
+      fun v ->
+        [%e match prim with
+            |Char -> [%expr Cstruct.get_char v [%e num f.off]]
+            |UInt8 -> [%expr Cstruct.get_uint8 v [%e num f.off]]
+            |UInt16 -> [%expr [%e m "get_uint16"] v [%e num f.off]]
+            |UInt32 -> [%expr [%e m "get_uint32"] v [%e num f.off]]
+            |UInt64 -> [%expr [%e m "get_uint64"] v [%e num f.off]]]]
+
+let type_of_int_field ~loc = function
+  |Char -> [%type: char]
+  |UInt8 -> [%type: Cstruct.uint8]
+  |UInt16 -> [%type: Cstruct.uint16]
+  |UInt32 -> [%type: Cstruct.uint32]
+  |UInt64 -> [%type: Cstruct.uint64]
+
+let set_expr loc s f =
+  let m = mode_mod loc s.endian in
+  let num x = Ast.eint ~loc x in
+  match f.ty with
+  |Buffer (_,_) ->
+    let len = width_of_field f in
+    [%expr
+      fun src srcoff dst ->
+        Cstruct.blit_from_string src srcoff dst [%e num f.off] [%e num len]]
+  |Prim prim ->
+    [%expr fun v x ->
+        [%e match prim with
+            |Char -> [%expr Cstruct.set_char v [%e num f.off] x]
+            |UInt8 -> [%expr Cstruct.set_uint8 v [%e num f.off] x]
+            |UInt16 -> [%expr [%e m "set_uint16"] v [%e num f.off] x]
+            |UInt32 -> [%expr [%e m "set_uint32"] v [%e num f.off] x]
+            |UInt64 -> [%expr [%e m "set_uint64"] v [%e num f.off] x]]]
+
+let type_of_set ~loc f =
+  match f.ty with
+  |Buffer (_,_) ->
+    [%type: string -> int -> Cstruct.t -> unit]
+  |Prim prim ->
+    let retf = type_of_int_field ~loc prim in
+    [%type: Cstruct.t -> [%t retf] -> unit]
+
+let hexdump_expr ~loc s =
+  [%expr fun v ->
+    let buf = Buffer.create 128 in
+    Buffer.add_string buf [%e Ast.estring ~loc (s.name ^ " = {\n")];
+    [%e op_evar ~loc s Op_hexdump_to_buffer] buf v;
+    print_endline (Buffer.contents buf);
+    print_endline "}"
+  ]
+
+let hexdump_to_buffer_expr ~loc s =
+  let prim_format_string = function
+    | Char -> [%expr "%c\n"]
+    | UInt8 | UInt16 -> [%expr "0x%x\n"]
+    | UInt32 -> [%expr "0x%lx\n"]
+    | UInt64 -> [%expr "0x%Lx\n"]
+  in
+  let hexdump_field = function
+    | Ignored_field ->
+      [%expr ()]
+    | Named_field f ->
+      let get_f = op_evar ~loc s (Op_get f) in
+      let expr =
+        match f.ty with
+        |Prim p ->
+          [%expr Printf.bprintf buf [%e prim_format_string p] ([%e get_f] v)]
+        |Buffer (_,_) ->
+          [%expr Printf.bprintf buf "<buffer %s>" [%e Ast.estring ~loc (field_to_string f)];
+            Cstruct.hexdump_to_buffer buf ([%e get_f] v)]
+    in
+    [%expr
+      Printf.bprintf buf "  %s = " [%e Ast.estring ~loc f.name];
+      [%e expr]]
+  in
+  [%expr fun buf v -> [%e Ast.esequence ~loc (List.map hexdump_field s.fields)]]
+
+let op_expr loc s = function
+  | Op_sizeof -> Ast.eint ~loc s.len
+  | Op_hexdump -> hexdump_expr ~loc s
+  | Op_hexdump_to_buffer -> hexdump_to_buffer_expr ~loc s
+  | Op_get f -> get_expr loc s f
+  | Op_set f -> set_expr loc s f
+  | Op_copy f ->
+    let len = width_of_field f in
+    [%expr fun src -> Cstruct.to_string src ~off:[%e Ast.eint ~loc f.off] ~len:[%e Ast.eint ~loc len] ]
+  | Op_blit f ->
+    let len = width_of_field f in
+    [%expr fun src srcoff dst ->
+      Cstruct.blit src srcoff dst [%e Ast.eint ~loc f.off] [%e Ast.eint ~loc len]]
+
+let field_ops_for =
+  function
+  | Ignored_field ->
+    []
+  | Named_field f ->
+    let if_buffer x =
+      match f.ty with
+      |Buffer (_,_) -> [x]
+      |Prim _ -> []
+    in
+    List.concat
+      [ [Op_get f]
+      ; if_buffer (Op_copy f)
+      ; [Op_set f]
+      ; if_buffer (Op_blit f)
+      ]
+
+let ops_for s =
+  ( [Op_sizeof]
+  @ List.concat (List.map field_ops_for s.fields)
+  @ [Op_hexdump_to_buffer;
+     Op_hexdump;
+    ])
+
+(** Generate functions of the form {get/set}_<struct>_<field> *)
+let output_struct_one_endian loc s =
+  List.map
+    (fun op ->
+       [%stri let[@ocaml.warning "-32"] [%p op_pvar ~loc s op] =
+                [%e op_expr loc s op]])
+    (ops_for s)
+
+let output_struct _loc s =
+  match s.endian with
+  | Bi_endian ->
+    (* In case of Bi-endian, create two modules - one for BE and one for LE *)
+    let expr_be = Mod.structure (output_struct_one_endian _loc {s with endian = Big_endian})
+    and expr_le = Mod.structure (output_struct_one_endian _loc {s with endian = Little_endian})
+
+    in [{pstr_desc = Pstr_module
+        {pmb_name = {txt = Some "BE"; loc = _loc}; pmb_expr = expr_be;
+        pmb_attributes = []; pmb_loc = _loc;}; pstr_loc = _loc;};
+        {pstr_desc = Pstr_module
+        {pmb_name = {txt = Some "LE"; loc = _loc}; pmb_expr = expr_le;
+        pmb_attributes = []; pmb_loc = _loc;}; pstr_loc = _loc;}
+        ]
+  | _ -> output_struct_one_endian _loc s
+
+let type_of_get ~loc f =
+  match f.ty with
+  |Buffer (_,_) ->
+    [%type: Cstruct.t -> Cstruct.t]
+  |Prim prim ->
+    let retf = type_of_int_field ~loc prim in
+    [%type: Cstruct.t -> [%t retf]]
+
+let op_typ ~loc = function
+  | Op_sizeof -> [%type: int]
+  | Op_hexdump_to_buffer -> [%type: Buffer.t -> Cstruct.t -> unit]
+  | Op_hexdump -> [%type: Cstruct.t -> unit]
+  | Op_get f -> type_of_get ~loc f
+  | Op_set f -> type_of_set ~loc f
+  | Op_copy _ -> [%type: Cstruct.t -> string]
+  | Op_blit _ -> [%type: Cstruct.t -> int -> Cstruct.t -> unit]
+
+(** Generate signatures of the form {get/set}_<struct>_<field> *)
+let output_struct_sig loc s =
+  List.map
+    (fun op ->
+       Sig.value
+         (Val.mk
+            (Loc.make (op_name s op) ~loc)
+            (op_typ ~loc op)))
+    (ops_for s)
+
+type enum_op =
+  | Enum_to_sexp
+  | Enum_of_sexp
+  | Enum_get
+  | Enum_set
+  | Enum_print
+  | Enum_parse
+  | Enum_compare
+
+type cenum =
+  { name : string Loc.t;
+    fields : (string Loc.t * int64) list;
+    prim : prim;
+    sexp : bool;
+  }
+
+let enum_op_name cenum =
+  let s = cenum.name.txt in
+  function
+  | Enum_to_sexp -> sprintf "sexp_of_%s" s
+  | Enum_of_sexp -> sprintf "%s_of_sexp" s
+  | Enum_get -> sprintf "int_to_%s" s
+  | Enum_set -> sprintf "%s_to_int" s
+  | Enum_print -> sprintf "%s_to_string" s
+  | Enum_parse -> sprintf "string_to_%s" s
+  | Enum_compare -> sprintf "compare_%s" s
+
+let enum_pattern ~loc {prim; _} =
+  let pat_integer f suffix i =
+    Pat.constant (Pconst_integer(f i, suffix))
+  in
+  match prim with
+  | Char ->
+    (fun i -> Ast.pchar ~loc (Char.chr (Int64.to_int i)))
+  | (UInt8 | UInt16) -> pat_integer Int64.to_string None
+  | UInt32 -> pat_integer (fun i -> Int32.to_string (Int64.to_int32 i)) (Some 'l')
+  | UInt64 -> pat_integer Int64.to_string (Some 'L')
+
+let enum_integer ~loc {prim; _} =
+  let expr_integer f suffix i =
+    Exp.constant (Pconst_integer(f i, suffix))
+  in
+  match prim with
+    | Char -> (fun i -> Ast.echar ~loc (Char.chr (Int64.to_int i)))
+    | (UInt8 | UInt16) -> expr_integer Int64.to_string None
+    | UInt32 -> expr_integer (fun i -> Int32.to_string (Int64.to_int32 i)) (Some 'l')
+    | UInt64 -> expr_integer Int64.to_string (Some 'L')
+
+let declare_enum_expr ~loc ({fields; _} as cenum) = function
+  | Enum_to_sexp ->
+    [%expr fun x -> Sexplib.Sexp.Atom ([%e Ast.evar ~loc (enum_op_name cenum Enum_print)] x) ]
+  | Enum_of_sexp ->
+    [%expr
+      fun x ->
+        match x with
+        | Sexplib.Sexp.List _ ->
+          raise (Sexplib.Pre_sexp.Of_sexp_error (Failure "expected Atom, got List", x))
+        | Sexplib.Sexp.Atom v ->
+          match [%e Ast.evar ~loc (enum_op_name cenum Enum_parse)] v with
+          | None ->
+            raise (Sexplib.Pre_sexp.Of_sexp_error (Failure "unable to parse enum string", x))
+          | Some r -> r
+    ]
+  | Enum_get ->
+    let getters = (List.map (fun ({txt = f; _},i) ->
+        Exp.case (enum_pattern ~loc cenum i) [%expr Some [%e Ast.econstr ~loc f []]]
+      ) fields) @ [Exp.case [%pat? _] [%expr None]]
+    in
+    Exp.function_ getters
+  | Enum_set ->
+    let setters = List.map (fun ({txt = f; _},i) ->
+        Exp.case (Ast.pconstr ~loc f []) (enum_integer ~loc cenum i)
+      ) fields in
+    Exp.function_ setters
+  | Enum_print ->
+    let printers = List.map (fun ({txt = f; _},_) ->
+        Exp.case (Ast.pconstr ~loc f []) (Ast.estring ~loc f)
+      ) fields in
+    Exp.function_ printers
+  | Enum_parse ->
+    let parsers = List.map (fun ({txt = f; _},_) ->
+        Exp.case (Ast.pstring ~loc f) [%expr Some [%e Ast.econstr ~loc f []]]
+      ) fields in
+    Exp.function_ (parsers @ [Exp.case [%pat? _] [%expr None]])
+  | Enum_compare -> [%expr fun x y ->
+    let to_int = [%e Ast.evar ~loc (enum_op_name cenum Enum_set)] in
+    Stdlib.compare (to_int x) (to_int y)
+  ]
+
+let enum_ops_for {sexp; _} =
+  Enum_get ::
+  Enum_set ::
+  Enum_compare ::
+  Enum_print ::
+  Enum_parse ::
+  if sexp then
+    [ Enum_to_sexp
+    ; Enum_of_sexp
+    ]
+  else
+    []
+
+let enum_type_decl {name; fields; _} =
+  let decls = List.map (fun (f,_) -> Type.constructor f) fields in
+  Type.mk ~kind:(Ptype_variant decls) name
+
+let output_enum ~loc cenum =
+  Str.type_ Recursive [enum_type_decl cenum] ::
+  List.map
+    (fun op ->
+       [%stri
+         let[@ocaml.warning "-32"] [%p Ast.pvar ~loc (enum_op_name cenum op)] =
+           [%e declare_enum_expr ~loc cenum op]
+       ])
+    (enum_ops_for cenum)
+
+let enum_op_type ~loc {name; prim; _} =
+  let cty = Ast.tconstr ~loc name.txt [] in
+  let oty = match prim with
+    | Char -> [%type: char]
+    | (UInt8|UInt16) -> [%type: int]
+    | UInt32 -> [%type: int32]
+    | UInt64 -> [%type: int64]
+  in
+  function
+  | Enum_get -> [%type: [%t oty] -> [%t cty] option]
+  | Enum_set -> [%type: [%t cty] -> [%t oty]]
+  | Enum_print -> [%type: [%t cty] -> string]
+  | Enum_parse -> [%type: string -> [%t cty] option]
+  | Enum_to_sexp -> [%type: [%t cty] -> Sexplib.Sexp.t]
+  | Enum_of_sexp -> [%type: Sexplib.Sexp.t -> [%t cty]]
+  | Enum_compare -> [%type: [%t cty] -> [%t cty] -> int]
+
+let output_enum_sig loc (cenum:cenum) =
+  Sig.type_ Recursive [enum_type_decl cenum] ::
+  List.map
+    (fun op ->
+       let name = enum_op_name cenum op in
+       let typ = enum_op_type ~loc cenum op in
+       Sig.value (Val.mk (Loc.make name ~loc) typ))
+    (enum_ops_for cenum)
+
+let constr_enum = function
+  | {pcd_name = f; pcd_args = Pcstr_tuple []; pcd_attributes = attrs; _} ->
+    let id = match attrs with
+      | [{attr_name = {txt = "id";_}; 
+          attr_payload = 
+            PStr [{ pstr_desc = 
+              Pstr_eval ({pexp_desc = Pexp_constant cst; pexp_loc = loc; _}, _); _}]
+          ;_ }] ->
+        let cst = match cst with
+          | Pconst_integer(i, _) -> Int64.of_string i
+          | _ ->
+            loc_err loc "invalid id"
+        in
+        Some cst
+      | _ ->
+        None
+    in
+    (f, id)
+  | {pcd_loc = loc; _} ->
+    loc_err loc "invalid cenum variant"
+
+let get_len = function
+  | [{attr_name = {txt = "len"; loc};
+      attr_payload = PStr
+        [{pstr_desc =
+          Pstr_eval ({pexp_desc = Pexp_constant (Pconst_integer (sz, None)); _}, _);
+          _}]
+      ; _}]
+    ->
+    let n = int_of_string sz in
+    if n > 0 then
+      Some n
+    else
+      loc_err loc "[@len] argument should be > 0"
+  | [{attr_name = {txt = "len"; loc}; _} ] ->
+    loc_err loc "[@len] argument should be an integer"
+  | _ ->
+    None
+
+let constr_field {pld_name = fname; pld_type = fty; pld_loc = loc; pld_attributes = att; _} =
+  let sz = match get_len fty.ptyp_attributes, get_len att with
+  | Some sz, None
+  | None, Some sz -> Some sz
+  | Some _, Some _ -> loc_err loc "multiple field length attribute"
+  | None, None -> None
+  in
+  let fty = match fty.ptyp_desc with
+    | Ptyp_constr ({txt = Lident fty; _}, []) -> fty
+    | _ ->
+      loc_err fty.ptyp_loc "type identifier expected"
+  in
+  parse_field loc fname.txt fty sz
+
+let cstruct decl =
+  let {ptype_name = name; ptype_kind = kind;
+       ptype_attributes = attrs; ptype_loc = loc; _} = decl in
+  let fields = match kind with
+    | Ptype_record fields -> List.map constr_field fields
+    | _ -> loc_err loc "record type declaration expected"
+  in
+  let endian = match attrs with
+    | [{attr_name = {txt = endian; _}; attr_payload = PStr []; _}] -> endian
+    | [_] -> loc_err loc "no attribute payload expected"
+    | _ -> loc_err loc "too many attributes"
+  in
+  create_struct loc endian name fields
+
+let cenum decl =
+  let {ptype_name = name; ptype_kind = kind;
+       ptype_attributes = attrs; ptype_loc = loc; _} = decl in
+  let fields = match kind with
+    | Ptype_variant fields -> fields
+    | _ ->
+      loc_err loc "expected variant type"
+  in
+  let width, sexp =
+    match attrs with
+    | ({attr_name = {txt = width; _};attr_payload= PStr [];_}) 
+      :: ({attr_name = {txt = "sexp"; _};attr_payload = PStr []; _}) :: [] ->
+      width, true
+    | ({attr_name = {txt = width; _};attr_payload= PStr [];_}) :: [] ->
+      width, false
+    | _ ->
+      loc_err loc "invalid cenum attributes"
+  in
+  let n = ref Int64.minus_one in
+  let incr_n () = n := Int64.succ !n in
+  let fields = List.map constr_enum fields in
+  let fields =
+    List.map (function
+        | (f, None)   -> incr_n (); (f, !n)
+        | (f, Some i) -> n := i; (f, i)
+      ) fields in
+  let prim = match ty_of_string width with
+    | None -> loc_err loc "enum: unknown width specifier %s" width
+    | Some p -> p
+  in
+  { name;
+    fields;
+    prim;
+    sexp;
+  }
+
+let signature_item' mapper = function
+  | {psig_desc =
+       Psig_extension (({txt = "cstruct"; _}, PStr [{pstr_desc = Pstr_type(_, [decl]); _}]), _);
+     psig_loc = loc} ->
+    output_struct_sig loc (cstruct decl)
+  | {psig_desc =
+       Psig_extension (({txt = "cenum"; _}, PStr [{pstr_desc = Pstr_type(_, [decl]); _}]), _);
+     psig_loc = loc} ->
+    output_enum_sig loc (cenum decl)
+  | other -> [mapper other]
+
+let structure_item' mapper = function
+  | {pstr_desc =
+       Pstr_extension (({txt = "cstruct"; _}, PStr [{pstr_desc = Pstr_type(_, [decl]); _}]), _);
+     pstr_loc = loc} ->
+    output_struct loc (cstruct decl)
+  | {pstr_desc =
+       Pstr_extension (({txt = "cenum"; _}, PStr [{pstr_desc = Pstr_type(_, [decl]); _}]), _);
+     pstr_loc = loc ;
+     _ } ->
+    output_enum ~loc (cenum decl)
+  | other -> [mapper other]
+
+class mapper = object
+  inherit Ast_traverse.map as super
+
+  method! signature s =
+    List.concat (List.map (signature_item' super#signature_item) s)
+
+  method! structure s =
+    List.concat (List.map (structure_item' super#structure_item) s)
+end
+
+let () =
+  let mapper = new mapper in
+  Driver.register_transformation "ppx_cstruct" ~impl:mapper#structure ~intf:mapper#signature