Repository: mransan/ocaml-protoc
Pull request #213 (2023-09-09 03:02:19)

large refactor, plugins, single-file codegen

The goal of this PR is to streamline the code generation for `ocaml-protoc` and
cleanup the code. This would help me add new code generation _plugins_[^1] and also
reduce the sprawl of files produced by `ocaml-protoc`.

This is very much breaking:
- options are now long (`--pb` and `--pp` instead of `-pb` and `-pp`)
- all code for a .proto now lives in a single .ml/.mli file

In addition, the compiler lib can be directly used to generate code from a
.proto, giving it a list of plugins, without requiring `Cmdline.t`.

[^1]: my plans are: expose mutable records (or have an option to directly make all records mutable in the first place); and codegen for services


c-cube: @Lupus this is not finished but it's what I had in mind. With this it should be easier to add more `Plugin.t` to generate services, etc.

Patch:
diff --git a/benchs/benchs.ml b/benchs/benchs.ml
index 5c09778..2177d71 100644
--- a/benchs/benchs.ml
+++ b/benchs/benchs.ml
@@ -263,18 +263,18 @@ let () =
        ]
 
 module Nested = struct
-  type person = Foo_types.person = {
+  type person = Foo.person = {
     name: string;
     age: int64;
   }
 
-  type store = Foo_types.store = {
+  type store = Foo.store = {
     address: string;
     employees: person list;
     clients: person list;
   }
 
-  type company = Foo_types.company = {
+  type company = Foo.company = {
     name: string;
     stores: store list;
   }
@@ -610,7 +610,7 @@ module Nested = struct
     let dec_s s =
       Pbrt.Decoder.(
         let dec = of_string s in
-        Foo_pb.decode_company dec)
+        Foo.decode_pb_company dec)
     in
     let c_basic = dec_s s_basic in
     let c_buffers_nested = dec_s s_buffers_nested in
diff --git a/benchs/dune b/benchs/dune
index c125c8d..8a8d064 100644
--- a/benchs/dune
+++ b/benchs/dune
@@ -4,7 +4,7 @@
  (libraries ocaml-protoc benchmark))
 
 (rule
- (targets foo_types.ml foo_types.mli foo_pb.ml foo_pb.mli foo_pp.ml)
+ (targets foo.ml foo.mli)
  (deps foo.proto)
  (action
-  (run ocaml-protoc %{deps} -binary -pp -ml_out .)))
+  (run ocaml-protoc %{deps} --binary --pp --ml_out .)))
diff --git a/dune-project b/dune-project
index f364fd1..f70d303 100644
--- a/dune-project
+++ b/dune-project
@@ -1,3 +1,35 @@
 (lang dune 2.0)
 (name ocaml-protoc)
-(allow_approximate_merlin)
+(generate_opam_files true)
+(version 2.4)
+
+(maintainers "Maxime Ransan <maxime.ransan@gmail.com>")
+(authors "Maxime Ransan <maxime.ransan@gmail.com>")
+(source (github mransan/ocaml-protoc))
+(license MIT)
+
+(package
+  (name ocaml-protoc)
+  (synopsis "Pure OCaml compiler for .proto files")
+  (depends
+    (ocaml (>= 4.03)))
+  (tags (protoc protobuf codegen)))
+
+(package
+  (name pbrt)
+  (synopsis "Runtime library for Protobuf tooling")
+  (depends 
+    stdlib-shims
+    (ocaml (>= 4.03)))
+  (tags (protobuf encode decode)))
+
+(package
+  (name pbrt_yojson)
+  (synopsis "Runtime library for ocaml-protoc to support JSON encoding/decoding")
+  (depends 
+    (ocaml (>= 4.03))
+    (yojson (>= 1.6))
+    base64)
+  (tags (protobuf encode decode)))
+
+
diff --git a/ocaml-protoc.opam b/ocaml-protoc.opam
index bedad9a..dfe571c 100644
--- a/ocaml-protoc.opam
+++ b/ocaml-protoc.opam
@@ -1,29 +1,29 @@
+# This file is generated by dune, edit dune-project instead
 opam-version: "2.0"
-name: "ocaml-protoc"
 version: "2.4"
-synopsis: "Protobuf compiler for OCaml"
-description: "Protobuf compiler for OCaml"
-maintainer: "Maxime Ransan <maxime.ransan@gmail.com>"
-post-messages: [
-"The runtime library is now called \"pbrt\"."
-]
-authors:[
-  "Maxime Ransan <maxime.ransan@gmail.com>"
-]
-homepage: "https://github.com/mransan/ocaml-protoc"
-bug-reports:"https://github.com/mransan/ocaml-protoc/issues"
-dev-repo:"git+https://github.com/mransan/ocaml-protoc.git"
+synopsis: "Pure OCaml compiler for .proto files"
+maintainer: ["Maxime Ransan <maxime.ransan@gmail.com>"]
+authors: ["Maxime Ransan <maxime.ransan@gmail.com>"]
 license: "MIT"
-build: [
-  ["dune" "subst"] {dev}
-  ["dune" "build" "@install" "-p" name "-j" jobs]
-  ["dune" "runtest" "-p" name "-j" jobs] {with-test}
-  ["dune" "build" "@doc" "-p" name "-j" jobs] {with-doc}
-]
+tags: ["protoc" "protobuf" "codegen"]
+homepage: "https://github.com/mransan/ocaml-protoc"
+bug-reports: "https://github.com/mransan/ocaml-protoc/issues"
 depends: [
-  "ocaml" {>="4.08.0"}
-  "dune"  {>="2.0"}
-  "stdlib-shims"
-  "pbrt" {= version}
-  "odoc" {with-doc}
+  "dune" {>= "2.0"}
+  "ocaml" {>= "4.03"}
+]
+build: [
+  ["dune" "subst"] {pinned}
+  [
+    "dune"
+    "build"
+    "-p"
+    name
+    "-j"
+    jobs
+    "@install"
+    "@runtest" {with-test}
+    "@doc" {with-doc}
+  ]
 ]
+dev-repo: "git+https://github.com/mransan/ocaml-protoc.git"
diff --git a/pbrt.opam b/pbrt.opam
index c4c138a..f772538 100644
--- a/pbrt.opam
+++ b/pbrt.opam
@@ -1,31 +1,30 @@
+# This file is generated by dune, edit dune-project instead
 opam-version: "2.0"
-name: "pbrt"
 version: "2.4"
 synopsis: "Runtime library for Protobuf tooling"
-description: "Runtime library for Protobuf tooling"
-maintainer: "Maxime Ransan <maxime.ransan@gmail.com>"
-post-messages: [
-"Pbrt: runtime library for ocaml-protoc.
-
-A shim library named \"ocaml-protoc\" still exists, to ease the
-migration."
-]
-authors:[
-  "Maxime Ransan <maxime.ransan@gmail.com>"
-]
-homepage: "https://github.com/mransan/ocaml-protoc"
-bug-reports:"https://github.com/mransan/ocaml-protoc/issues"
-dev-repo:"git+https://github.com/mransan/ocaml-protoc.git"
+maintainer: ["Maxime Ransan <maxime.ransan@gmail.com>"]
+authors: ["Maxime Ransan <maxime.ransan@gmail.com>"]
 license: "MIT"
-build: [
-  ["dune" "subst"] {dev}
-  ["dune" "build" "@install" "-p" name "-j" jobs]
-  ["dune" "runtest" "-p" "pbrt,ocaml-protoc" "-j" jobs] {with-test}
-  ["dune" "build" "@doc" "-p" name "-j" jobs] {with-doc}
-]
+tags: ["protobuf" "encode" "decode"]
+homepage: "https://github.com/mransan/ocaml-protoc"
+bug-reports: "https://github.com/mransan/ocaml-protoc/issues"
 depends: [
-  "ocaml" {>="4.03.0"}
-  "dune"  {>="2.0"}
+  "dune" {>= "2.0"}
   "stdlib-shims"
-  "odoc" {with-doc}
+  "ocaml" {>= "4.03"}
+]
+build: [
+  ["dune" "subst"] {pinned}
+  [
+    "dune"
+    "build"
+    "-p"
+    name
+    "-j"
+    jobs
+    "@install"
+    "@runtest" {with-test}
+    "@doc" {with-doc}
+  ]
 ]
+dev-repo: "git+https://github.com/mransan/ocaml-protoc.git"
diff --git a/pbrt_yojson.opam b/pbrt_yojson.opam
index e12b9dc..103db79 100644
--- a/pbrt_yojson.opam
+++ b/pbrt_yojson.opam
@@ -1,16 +1,22 @@
+# This file is generated by dune, edit dune-project instead
 opam-version: "2.0"
-version: "0.2.0"
-maintainer: "Maxime Ransan <maxime.ransan@gmail.com>"
-authors:[
-  "Maxime Ransan <maxime.ransan@gmail.com>"
-]
-synopsis: "Runtime library for ocaml-protoc to support JSON encoding/decoding"
-homepage: "https://github.com/mransan/ocaml-protoc-yojson"
-bug-reports:"https://github.com/mransan/ocaml-protoc-yojson/issues"
-dev-repo: "git+https://github.com/mransan/ocaml-protoc-yojson.git"
+version: "2.4"
+synopsis:
+  "Runtime library for ocaml-protoc to support JSON encoding/decoding"
+maintainer: ["Maxime Ransan <maxime.ransan@gmail.com>"]
+authors: ["Maxime Ransan <maxime.ransan@gmail.com>"]
 license: "MIT"
+tags: ["protobuf" "encode" "decode"]
+homepage: "https://github.com/mransan/ocaml-protoc"
+bug-reports: "https://github.com/mransan/ocaml-protoc/issues"
+depends: [
+  "dune" {>= "2.0"}
+  "ocaml" {>= "4.03"}
+  "yojson" {>= "1.6"}
+  "base64"
+]
 build: [
-  ["dune" "subst"] {dev}
+  ["dune" "subst"] {pinned}
   [
     "dune"
     "build"
@@ -23,9 +29,4 @@ build: [
     "@doc" {with-doc}
   ]
 ]
-depends: [
-  "dune" { >= "2.0" }
-  "yojson" {>= "1.6.0"}
-  "base64"
-  "ocaml" { >= "4.02.1" }
-]
+dev-repo: "git+https://github.com/mransan/ocaml-protoc.git"
diff --git a/src/compilerlib/dune b/src/compilerlib/dune
index 30532ec..3064fcb 100644
--- a/src/compilerlib/dune
+++ b/src/compilerlib/dune
@@ -7,10 +7,10 @@
  (public_name ocaml-protoc.compiler-lib)
  (synopsis "Compiler library for ocaml-protoc, to turn .proto files into OCaml code")
  (wrapped true)
- (modules pb_codegen_backend pb_codegen_decode_binary pb_codegen_decode_bs
+ (modules pb_codegen_all pb_codegen_backend pb_codegen_decode_binary pb_codegen_decode_bs
    pb_codegen_decode_yojson pb_codegen_default pb_codegen_encode_binary
    pb_codegen_encode_bs pb_codegen_encode_yojson pb_codegen_formatting
-   pb_codegen_ocaml_type pb_codegen_pp pb_codegen_sig pb_codegen_types
+   pb_codegen_ocaml_type pb_codegen_pp pb_codegen_plugin pb_codegen_types
    pb_codegen_util pb_exception pb_field_type pb_location pb_logger pb_option
    pb_parsing pb_parsing_lexer pb_parsing_parser pb_parsing_parse_tree
    pb_parsing_util pb_typing_graph pb_typing pb_typing_recursion
diff --git a/src/compilerlib/pb_codegen_all.ml b/src/compilerlib/pb_codegen_all.ml
new file mode 100644
index 0000000..2f2f5f0
--- /dev/null
+++ b/src/compilerlib/pb_codegen_all.ml
@@ -0,0 +1,129 @@
+(*
+  The MIT License (MIT)
+  
+  Copyright (c) 2016 Maxime Ransan <maxime.ransan@gmail.com>
+  
+  Permission is hereby granted, free of charge, to any person obtaining a copy
+  of this software and associated documentation files (the "Software"), to deal
+  in the Software without restriction, including without limitation the rights
+  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+  copies of the Software, and to permit persons to whom the Software is
+  furnished to do so, subject to the following conditions:
+
+  The above copyright notice and this permission notice shall be included in all
+  copies or substantial portions of the Software.
+
+  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+  SOFTWARE.
+
+*)
+
+module E = Pb_exception
+module Ot = Pb_codegen_ocaml_type
+module F = Pb_codegen_formatting
+module Plugin = Pb_codegen_plugin
+
+type codegen_f = Plugin.codegen_f
+
+type ocaml_mod = {
+  ml: F.scope;
+  mli: F.scope;
+}
+
+let new_ocaml_mod ~proto_file_options ~proto_file_name () : ocaml_mod =
+  let self = { ml = F.empty_scope (); mli = F.empty_scope () } in
+
+  let print_ppx sc =
+    match Pb_option.get proto_file_options "ocaml_file_ppx" with
+    | None -> ()
+    | Some Pb_option.(Scalar_value (Constant_string s)) ->
+      F.linep sc "[@@@%s]" s
+    | _ -> E.invalid_ppx_extension_option proto_file_name
+  in
+
+  (* write preludes *)
+  F.line self.ml "[@@@ocaml.warning \"-27-30-39\"]";
+  F.empty_line self.ml;
+  print_ppx self.ml;
+  F.empty_line self.mli;
+
+  F.linep self.mli "(** Code for %s *)" (Filename.basename proto_file_name);
+  F.empty_line self.mli;
+  F.linep self.mli "(* generated from %S, do not edit *)" proto_file_name;
+  F.empty_line self.mli;
+  print_ppx self.mli;
+  F.empty_line self.mli;
+
+  self
+
+let generate_for_all_types (ocaml_types : Ot.type_ list list) (sc : F.scope)
+    (f : codegen_f) ocamldoc_title : unit =
+  (match ocamldoc_title with
+  | None -> ()
+  | Some ocamldoc_title ->
+    F.empty_line sc;
+    F.linep sc "(** {2 %s} *)" ocamldoc_title;
+    F.empty_line sc);
+
+  List.iter
+    (fun types ->
+      let (_ : bool) =
+        List.fold_left
+          (fun first type_ ->
+            let has_encoded =
+              if first then
+                f type_ sc
+              else
+                f ~and_:() type_ sc
+            in
+            F.empty_line sc;
+            first && not has_encoded)
+          true types
+      in
+      ())
+    ocaml_types
+
+let generate_type_and_default (self : ocaml_mod) ocaml_types : unit =
+  generate_for_all_types ocaml_types self.ml Pb_codegen_types.gen_struct None;
+  generate_for_all_types ocaml_types self.ml Pb_codegen_default.gen_struct None;
+
+  generate_for_all_types ocaml_types self.mli Pb_codegen_types.gen_sig
+    (Some Pb_codegen_types.ocamldoc_title);
+  generate_for_all_types ocaml_types self.mli Pb_codegen_default.gen_sig
+    (Some Pb_codegen_default.ocamldoc_title);
+  ()
+
+let generate_mutable_records (self : ocaml_mod) ocaml_types : unit =
+  let ocaml_types = List.flatten ocaml_types in
+  List.iter
+    (fun { Ot.spec; _ } ->
+      match spec with
+      | Ot.Record r ->
+        Pb_codegen_types.gen_record_mutable r self.ml;
+        F.empty_line self.ml;
+        Pb_codegen_default.gen_record_mutable r self.ml;
+        F.empty_line self.ml
+      | _ -> ())
+    ocaml_types
+
+let generate_plugin (self : ocaml_mod) ocaml_types (p : Plugin.t) : unit =
+  let (module P) = p in
+  F.line self.ml "[@@@ocaml.warning \"-27-30-39\"]";
+  generate_for_all_types ocaml_types self.ml P.gen_struct
+    (Some P.ocamldoc_title);
+  generate_for_all_types ocaml_types self.mli P.gen_sig (Some P.ocamldoc_title);
+  ()
+
+let codegen ocaml_types ~proto_file_options ~proto_file_name
+    (plugins : Plugin.t list) : ocaml_mod =
+  let self = new_ocaml_mod ~proto_file_options ~proto_file_name () in
+  generate_type_and_default self ocaml_types;
+  if List.exists Pb_codegen_plugin.requires_mutable_records plugins then
+    generate_mutable_records self ocaml_types;
+  List.iter (generate_plugin self ocaml_types) plugins;
+  self
diff --git a/src/compilerlib/pb_codegen_all.mli b/src/compilerlib/pb_codegen_all.mli
new file mode 100644
index 0000000..6394720
--- /dev/null
+++ b/src/compilerlib/pb_codegen_all.mli
@@ -0,0 +1,19 @@
+(** Generate all the code for a given OCaml module. *)
+
+module Ot = Pb_codegen_ocaml_type
+module Plugin = Pb_codegen_plugin
+module F = Pb_codegen_formatting
+
+type ocaml_mod = {
+  ml: F.scope;
+  mli: F.scope;
+}
+
+val codegen :
+  Ot.type_ list list ->
+  proto_file_options:Pb_option.set ->
+  proto_file_name:string ->
+  Plugin.t list ->
+  ocaml_mod
+(** [codegen types] returns a full code listing for the [.ml]
+    and the [.mli] files. *)
diff --git a/src/compilerlib/pb_codegen_backend.ml b/src/compilerlib/pb_codegen_backend.ml
index 6211cd6..10243c8 100644
--- a/src/compilerlib/pb_codegen_backend.ml
+++ b/src/compilerlib/pb_codegen_backend.ml
@@ -28,14 +28,14 @@ module Ot = Pb_codegen_ocaml_type
 module Tt = Pb_typing_type_tree
 module Typing_util = Pb_typing_util
 
-(* [rev_split_by_naming_convention s] will split [s] according to the protobuf
- *  coding style convention. The rule split are
- *  {ul
- *  {- character ['_'] is a separator}
- *  {- the first uppercase letter after a lower case is a separator
- *     (ie FooBar will be split into [ ["Bar";"Foo"] ]}
- *  }
- *)
+(** [rev_split_by_naming_convention s] will split [s] according to the protobuf
+    coding style convention. The rule split are
+    {ul
+    {- character ['_'] is a separator}
+    {- the first uppercase letter after a lower case is a separator
+       (ie FooBar will be split into [ ["Bar";"Foo"] ]}
+    }
+*)
 let rev_split_by_naming_convention s =
   let is_uppercase c = 64 < Char.code c && Char.code c < 91 in
 
diff --git a/src/compilerlib/pb_codegen_backend.mli b/src/compilerlib/pb_codegen_backend.mli
index 0f6ffdd..fc0e1bf 100644
--- a/src/compilerlib/pb_codegen_backend.mli
+++ b/src/compilerlib/pb_codegen_backend.mli
@@ -23,9 +23,9 @@
 
 *)
 
-(** Compile protobuf parse tree to the ocaml type *)
+(** Compile protobuf typed tree to the OCaml AST.
 
-(** This module focuses on the compilation steps which transforms a
+    This module focuses on the compilation steps which transforms a
     fully resolved Protobuf message into an OCaml representation.
 
     After compilation this module also expose code generation
diff --git a/src/compilerlib/pb_codegen_decode_binary.ml b/src/compilerlib/pb_codegen_decode_binary.ml
index 288f863..b13bdef 100644
--- a/src/compilerlib/pb_codegen_decode_binary.ml
+++ b/src/compilerlib/pb_codegen_decode_binary.ml
@@ -2,7 +2,6 @@ module Ot = Pb_codegen_ocaml_type
 module F = Pb_codegen_formatting
 
 let sp = Pb_codegen_util.sp
-let file_suffix = "pb"
 
 let runtime_function_for_basic_type bt pk =
   match pk, bt with
@@ -44,10 +43,8 @@ let decode_field_expression field_type pk =
   match field_type with
   | Ot.Ft_user_defined_type t ->
     let f_name =
-      let function_prefix = "decode" in
-      let module_suffix = file_suffix in
-      Pb_codegen_util.function_name_of_user_defined ~function_prefix
-        ~module_suffix t
+      let function_prefix = "decode_pb" in
+      Pb_codegen_util.function_name_of_user_defined ~function_prefix t
     in
     (match t.Ot.udt_type with
     | `Message -> f_name ^ " (Pbrt.Decoder.nested d)"
@@ -66,7 +63,7 @@ let gen_field_common sc encoding_number payload_kind message_name
     ?(is_packed = false) f =
   F.linep sc "| Some (%i, Pbrt.%s) -> begin" encoding_number
     (pbrt_payload_kind payload_kind is_packed);
-  F.scope sc (fun sc -> f sc);
+  F.sub_scope sc f;
   F.line sc "end";
   F.linep sc "| Some (%i, pk) -> " encoding_number;
   F.linep sc "  Pbrt.Decoder.unexpected_payload \"%s\" pk"
@@ -113,7 +110,7 @@ let gen_rft_repeated sc r_name rf_label repeated_field =
   | Ot.Rt_repeated_field, true ->
     gen_field_common sc encoding_number pk r_name ~is_packed (fun sc ->
         F.line sc "Pbrt.Decoder.packed_fold (fun () d -> ";
-        F.scope sc (fun sc ->
+        F.sub_scope sc (fun sc ->
             F.linep sc "Pbrt.Repeated_field.add (%s) v.%s;"
               (decode_field_expression field_type pk)
               rf_label);
@@ -138,7 +135,7 @@ let gen_rft_associative sc r_name rf_label associative_field =
    * the performance. *)
   gen_field_common sc encoding_number Ot.Pk_bytes r_name (fun sc ->
       F.line sc "let decode_value = (fun d ->";
-      F.scope sc (fun sc ->
+      F.sub_scope sc (fun sc ->
           F.line sc @@ decode_field_expression value_type value_pk);
       F.line sc ") in";
       let decode_expression =
@@ -149,13 +146,14 @@ let gen_rft_associative sc r_name rf_label associative_field =
       match at with
       | Ot.At_list ->
         F.linep sc "v.%s <- (" rf_label;
-        F.scope sc (fun sc -> F.linep sc "%s::v.%s;" decode_expression rf_label);
+        F.sub_scope sc (fun sc ->
+            F.linep sc "%s::v.%s;" decode_expression rf_label);
         F.line sc ");"
       | Ot.At_hashtable ->
         F.linep sc "let a, b = %s in" decode_expression;
         F.linep sc "Hashtbl.add v.%s a b;" rf_label)
 
-let gen_rft_variant sc module_prefix r_name rf_label { Ot.v_constructors; _ } =
+let gen_rft_variant sc r_name rf_label { Ot.v_constructors; _ } =
   List.iter
     (fun variant_constructor ->
       let {
@@ -171,19 +169,17 @@ let gen_rft_variant sc module_prefix r_name rf_label { Ot.v_constructors; _ } =
           match vc_field_type with
           | Ot.Vct_nullary ->
             F.line sc "Pbrt.Decoder.empty_nested d;";
-            F.linep sc "v.%s <- %s_types.%s;" rf_label module_prefix
-              vc_constructor
+            F.linep sc "v.%s <- %s;" rf_label vc_constructor
           | Ot.Vct_non_nullary_constructor field_type ->
-            F.linep sc "v.%s <- %s_types.%s (%s);" rf_label module_prefix
-              vc_constructor
+            F.linep sc "v.%s <- %s (%s);" rf_label vc_constructor
               (decode_field_expression field_type pk)))
     v_constructors
 
-let gen_record ?and_ module_prefix { Ot.r_name; r_fields } sc =
+let gen_record ?and_ { Ot.r_name; r_fields } sc =
   (* list fields have a special treatement when decoding since each new element
-   * of a repeated field is appended to the front of the list. In order
-   * to retreive the right order efficiently we reverse all the repeated field
-   * lists values when the message is done being decoded. *)
+     of a repeated field is appended to the front of the list. In order
+     to retreive the right order efficiently we reverse all the repeated field
+     lists values when the message is done being decoded. *)
   let all_lists =
     List.fold_left
       (fun acc { Ot.rf_label; rf_field_type; _ } ->
@@ -205,8 +201,8 @@ let gen_record ?and_ module_prefix { Ot.r_name; r_fields } sc =
 
   let mutable_record_name = Pb_codegen_util.mutable_record_name r_name in
 
-  F.linep sc "%s decode_%s d =" (Pb_codegen_util.let_decl_of_and and_) r_name;
-  F.scope sc (fun sc ->
+  F.linep sc "%s decode_pb_%s d =" (Pb_codegen_util.let_decl_of_and and_) r_name;
+  F.sub_scope sc (fun sc ->
       F.linep sc "let v = default_%s () in" mutable_record_name;
       F.line sc "let continue__= ref true in";
 
@@ -221,12 +217,12 @@ let gen_record ?and_ module_prefix { Ot.r_name; r_fields } sc =
        * function loop iterate over all fields returned by the Protobuf
        * runtime. *)
       F.line sc "while !continue__ do";
-      F.scope sc (fun sc ->
+      F.sub_scope sc (fun sc ->
           F.line sc "match Pbrt.Decoder.key d with";
 
           (* termination condition *)
           F.line sc "| None -> (";
-          F.scope sc (fun sc ->
+          F.sub_scope sc (fun sc ->
               List.iter
                 (fun field_name ->
                   F.linep sc "v.%s <- List.rev v.%s;" field_name field_name)
@@ -243,8 +239,7 @@ let gen_record ?and_ module_prefix { Ot.r_name; r_fields } sc =
               | Ot.Rft_optional x -> gen_rft_optional sc r_name rf_label x
               | Ot.Rft_repeated x -> gen_rft_repeated sc r_name rf_label x
               | Ot.Rft_associative x -> gen_rft_associative sc r_name rf_label x
-              | Ot.Rft_variant x ->
-                gen_rft_variant sc module_prefix r_name rf_label x)
+              | Ot.Rft_variant x -> gen_rft_variant sc r_name rf_label x)
             r_fields;
           F.line sc
             ("| Some (_, payload_kind) -> " ^ "Pbrt.Decoder.skip d payload_kind"));
@@ -261,23 +256,25 @@ let gen_record ?and_ module_prefix { Ot.r_name; r_fields } sc =
         all_required_rf_labels;
 
       F.line sc "({";
-      F.scope sc (fun sc ->
+      F.sub_scope sc (fun sc ->
           List.iter
             (fun { Ot.rf_label; _ } ->
-              F.linep sc "%s_types.%s = v.%s;" module_prefix rf_label rf_label)
+              F.linep sc "%s = v.%s;" rf_label rf_label)
             r_fields);
-      F.linep sc "} : %s_types.%s)" module_prefix r_name)
+      F.linep sc "} : %s)" r_name)
 
 let gen_unit ?and_ { Ot.er_name } sc =
-  F.linep sc "%s decode_%s d =" (Pb_codegen_util.let_decl_of_and and_) er_name;
-  F.scope sc (fun sc ->
+  F.linep sc "%s decode_pb_%s d ="
+    (Pb_codegen_util.let_decl_of_and and_)
+    er_name;
+  F.sub_scope sc (fun sc ->
       F.line sc "match Pbrt.Decoder.key d with";
       F.line sc "| None -> ();";
       F.line sc "| Some (_, pk) -> ";
       F.linep sc "  Pbrt.Decoder.unexpected_payload \"%s\" pk"
         (sp "Unexpected fields in empty message(%s)" er_name))
 
-let gen_variant ?and_ module_prefix { Ot.v_name; v_constructors } sc =
+let gen_variant ?and_ { Ot.v_name; v_constructors } sc =
   let process_ctor sc variant_constructor =
     let {
       Ot.vc_constructor;
@@ -291,26 +288,26 @@ let gen_variant ?and_ module_prefix { Ot.v_name; v_constructors } sc =
     match vc_field_type with
     | Ot.Vct_nullary ->
       F.linep sc "| Some (%i, _) -> begin " vc_encoding_number;
-      F.scope sc (fun sc ->
+      F.sub_scope sc (fun sc ->
           F.line sc "Pbrt.Decoder.empty_nested d ;";
-          F.linep sc "(%s_types.%s : %s_types.%s)" module_prefix vc_constructor
-            module_prefix v_name);
+          F.linep sc "(%s : %s)" vc_constructor v_name);
       F.line sc "end"
     | Ot.Vct_non_nullary_constructor field_type ->
-      F.linep sc "| Some (%i, _) -> (%s_types.%s (%s) : %s_types.%s) "
-        vc_encoding_number module_prefix vc_constructor
+      F.linep sc "| Some (%i, _) -> (%s (%s) : %s) " vc_encoding_number
+        vc_constructor
         (decode_field_expression field_type pk)
-        module_prefix v_name
+        v_name
   in
 
-  F.linep sc "%s decode_%s d = " (Pb_codegen_util.let_decl_of_and and_) v_name;
-  F.scope sc (fun sc ->
+  F.linep sc "%s decode_pb_%s d = "
+    (Pb_codegen_util.let_decl_of_and and_)
+    v_name;
+  F.sub_scope sc (fun sc ->
       F.linep sc "let rec loop () = ";
-      F.scope sc (fun sc ->
-          F.linep sc "let ret:%s_types.%s = match Pbrt.Decoder.key d with"
-            module_prefix v_name;
+      F.sub_scope sc (fun sc ->
+          F.linep sc "let ret:%s = match Pbrt.Decoder.key d with" v_name;
 
-          F.scope sc (fun sc ->
+          F.sub_scope sc (fun sc ->
               F.linep sc "| None -> Pbrt.Decoder.malformed_variant \"%s\""
                 v_name;
               List.iter (fun ctor -> process_ctor sc ctor) v_constructors;
@@ -324,30 +321,31 @@ let gen_variant ?and_ module_prefix { Ot.v_name; v_constructors } sc =
       F.line sc "in";
       F.line sc "loop ()")
 
-let gen_const_variant ?and_ module_prefix { Ot.cv_name; cv_constructors } sc =
-  F.linep sc "%s decode_%s d = " (Pb_codegen_util.let_decl_of_and and_) cv_name;
-  F.scope sc (fun sc ->
+let gen_const_variant ?and_ { Ot.cv_name; cv_constructors } sc =
+  F.linep sc "%s decode_pb_%s d = "
+    (Pb_codegen_util.let_decl_of_and and_)
+    cv_name;
+  F.sub_scope sc (fun sc ->
       F.line sc "match Pbrt.Decoder.int_as_varint d with";
       List.iter
         (fun { Ot.cvc_name; cvc_binary_value; _ } ->
-          F.linep sc "| %i -> (%s_types.%s:%s_types.%s)" cvc_binary_value
-            module_prefix cvc_name module_prefix cv_name)
+          F.linep sc "| %i -> (%s:%s)" cvc_binary_value cvc_name cv_name)
         cv_constructors;
       F.linep sc "| _ -> Pbrt.Decoder.malformed_variant \"%s\"" cv_name)
 
 let gen_struct ?and_ t sc =
-  let { Ot.module_prefix; spec; _ } = t in
+  let { Ot.spec; _ } = t in
 
   let has_encoded =
     match spec with
     | Ot.Record r ->
-      gen_record ?and_ module_prefix r sc;
+      gen_record ?and_ r sc;
       true
     | Ot.Variant v ->
-      gen_variant ?and_ module_prefix v sc;
+      gen_variant ?and_ v sc;
       true
     | Ot.Const_variant v ->
-      gen_const_variant ?and_ module_prefix v sc;
+      gen_const_variant ?and_ v sc;
       true
     | Ot.Unit u ->
       gen_unit ?and_ u sc;
@@ -359,13 +357,13 @@ let gen_struct ?and_ t sc =
 let gen_sig ?and_ t sc =
   let _ = and_ in
 
-  let { Ot.module_prefix; spec; _ } = t in
+  let { Ot.spec; _ } = t in
 
   let f type_name =
-    F.linep sc "val decode_%s : Pbrt.Decoder.t -> %s_types.%s" type_name
-      module_prefix type_name;
+    F.linep sc "val decode_pb_%s : Pbrt.Decoder.t -> %s" type_name type_name;
     F.linep sc
-      ("(** [decode_%s decoder] decodes a " ^^ "[%s] value from [decoder] *)")
+      ("(** [decode_pb_%s decoder] decodes a "
+     ^^ "[%s] binary value from [decoder] *)")
       type_name type_name
   in
 
@@ -388,3 +386,13 @@ let gen_sig ?and_ t sc =
   has_encoded
 
 let ocamldoc_title = "Protobuf Decoding"
+let requires_mutable_records = true
+
+let plugin : Pb_codegen_plugin.t =
+  let module P = struct
+    let gen_sig = gen_sig
+    let gen_struct = gen_struct
+    let ocamldoc_title = ocamldoc_title
+    let requires_mutable_records = requires_mutable_records
+  end in
+  (module P)
diff --git a/src/compilerlib/pb_codegen_decode_binary.mli b/src/compilerlib/pb_codegen_decode_binary.mli
index ee072d4..3fe5483 100644
--- a/src/compilerlib/pb_codegen_decode_binary.mli
+++ b/src/compilerlib/pb_codegen_decode_binary.mli
@@ -1,3 +1,5 @@
 (** Code generator for the [decode] function *)
 
-include Pb_codegen_sig.S
+include Pb_codegen_plugin.S
+
+val plugin : Pb_codegen_plugin.t
diff --git a/src/compilerlib/pb_codegen_decode_bs.ml b/src/compilerlib/pb_codegen_decode_bs.ml
index 066fe03..645f782 100644
--- a/src/compilerlib/pb_codegen_decode_bs.ml
+++ b/src/compilerlib/pb_codegen_decode_bs.ml
@@ -23,9 +23,7 @@ let value_expression ~r_name ~rf_label field_type =
     let { Ot.udt_type; _ } = udt in
     let f_name =
       let function_prefix = "decode" in
-      let module_suffix = "bs" in
-      Pb_codegen_util.function_name_of_user_defined ~function_prefix
-        ~module_suffix udt
+      Pb_codegen_util.function_name_of_user_defined ~function_prefix udt
     in
     (match udt_type with
     | `Message ->
@@ -62,9 +60,9 @@ let gen_rft_repeated sc ~r_name ~rf_label repeated_field =
 
   F.linep sc "| \"%s\" -> begin" json_label;
 
-  F.scope sc (fun sc ->
+  F.sub_scope sc (fun sc ->
       F.line sc "let a = ";
-      F.scope sc (fun sc ->
+      F.sub_scope sc (fun sc ->
           F.linep sc "let a = Js.Dict.unsafeGet json \"%s\" in " json_label;
           F.linep sc "Pbrt_bs.array_ a \"%s\" \"%s\"" r_name rf_label);
       F.line sc "in";
@@ -105,18 +103,18 @@ let gen_rft_variant sc ~r_name ~rf_label { Ot.v_constructors; _ } =
     v_constructors
 
 (* Generate decode function for a record *)
-let gen_record ?and_ module_prefix { Ot.r_name; r_fields } sc =
+let gen_record ?and_ { Ot.r_name; r_fields } sc =
   let mutable_record_name = Pb_codegen_util.mutable_record_name r_name in
 
   F.linep sc "%s decode_%s json =" (Pb_codegen_util.let_decl_of_and and_) r_name;
 
-  F.scope sc (fun sc ->
+  F.sub_scope sc (fun sc ->
       F.linep sc "let v = default_%s () in" mutable_record_name;
       F.line sc "let keys = Js.Dict.keys json in";
       F.line sc "let last_key_index = Array.length keys - 1 in";
 
       F.line sc "for i = 0 to last_key_index do";
-      F.scope sc (fun sc ->
+      F.sub_scope sc (fun sc ->
           F.line sc "match Array.unsafe_get keys i with";
 
           (* Generate pattern match for all the possible message field *)
@@ -146,23 +144,22 @@ let gen_record ?and_ module_prefix { Ot.r_name; r_fields } sc =
 
       (* Transform the mutable record in an immutable one *)
       F.line sc "({";
-      F.scope sc (fun sc ->
+      F.sub_scope sc (fun sc ->
           List.iter
             (fun { Ot.rf_label; _ } ->
-              F.linep sc "%s_types.%s = v.%s;" module_prefix rf_label rf_label)
+              F.linep sc "%s = v.%s;" rf_label rf_label)
             r_fields);
-      F.linep sc "} : %s_types.%s)" module_prefix r_name)
+      F.linep sc "} : %s)" r_name)
 
 (* Generate decode function for a variant type *)
-let gen_variant ?and_ module_prefix { Ot.v_name; v_constructors } sc =
+let gen_variant ?and_ { Ot.v_name; v_constructors } sc =
   (* helper function for each constructor case *)
   let process_v_constructor sc { Ot.vc_constructor; vc_field_type; _ } =
     let json_label = Pb_codegen_util.camel_case_of_constructor vc_constructor in
 
     match vc_field_type with
     | Ot.Vct_nullary ->
-      F.linep sc "| \"%s\" -> (%s_types.%s : %s_types.%s)" json_label
-        module_prefix vc_constructor module_prefix v_name
+      F.linep sc "| \"%s\" -> (%s : %s)" json_label vc_constructor v_name
     | Ot.Vct_non_nullary_constructor field_type ->
       let value_expression =
         let r_name = v_name and rf_label = vc_constructor in
@@ -171,13 +168,12 @@ let gen_variant ?and_ module_prefix { Ot.v_name; v_constructors } sc =
 
       F.linep sc "| \"%s\" -> " json_label;
       F.linep sc "  let json = Js.Dict.unsafeGet json \"%s\" in" json_label;
-      F.linep sc "  (%s_types.%s (%s) : %s_types.%s)" module_prefix
-        vc_constructor value_expression module_prefix v_name
+      F.linep sc "  (%s (%s) : %s)" vc_constructor value_expression v_name
   in
 
   F.linep sc "%s decode_%s json =" (Pb_codegen_util.let_decl_of_and and_) v_name;
 
-  F.scope sc (fun sc ->
+  F.sub_scope sc (fun sc ->
       F.line sc "let keys = Js.Dict.keys json in";
 
       (* even though a variant should be an object with a single field,
@@ -185,11 +181,11 @@ let gen_variant ?and_ module_prefix { Ot.v_name; v_constructors } sc =
        * we still need a loop to iterate over the key/value, even if in 99.99%
        * of the cases it will be a single iteration *)
       F.line sc "let rec loop = function ";
-      F.scope sc (fun sc ->
+      F.sub_scope sc (fun sc ->
           F.linep sc "| -1 -> Pbrt_bs.E.malformed_variant \"%s\"" v_name;
           F.line sc "| i -> ";
 
-          F.scope sc (fun sc ->
+          F.sub_scope sc (fun sc ->
               F.line sc "begin match Array.unsafe_get keys i with";
               List.iter (process_v_constructor sc) v_constructors;
               F.empty_line sc;
@@ -199,21 +195,20 @@ let gen_variant ?and_ module_prefix { Ot.v_name; v_constructors } sc =
       F.line sc "in";
       F.line sc "loop (Array.length keys - 1)")
 
-let gen_const_variant ?and_ module_prefix { Ot.cv_name; cv_constructors } sc =
+let gen_const_variant ?and_ { Ot.cv_name; cv_constructors } sc =
   F.linep sc "%s decode_%s (json:Js.Json.t) ="
     (Pb_codegen_util.let_decl_of_and and_)
     cv_name;
 
-  F.scope sc (fun sc ->
+  F.sub_scope sc (fun sc ->
       F.linep sc "match Pbrt_bs.string json \"%s\" \"value\" with" cv_name;
 
       List.iter
         (fun { Ot.cvc_name; cvc_string_value; _ } ->
-          F.linep sc "| \"%s\" -> (%s_types.%s : %s_types.%s)" cvc_string_value
-            module_prefix cvc_name module_prefix cv_name)
+          F.linep sc "| \"%s\" -> (%s : %s)" cvc_string_value cvc_name cv_name)
         cv_constructors;
 
-      F.linep sc "| \"\" -> %s_types.%s" module_prefix
+      F.linep sc "| \"\" -> %s"
         (let { Ot.cvc_name; _ } = List.hd cv_constructors in
          cvc_name);
 
@@ -225,19 +220,19 @@ let gen_unit ?and_ { Ot.er_name } sc =
   @@ sp "%s decode_%s d =" (Pb_codegen_util.let_decl_of_and and_) er_name;
   F.line sc (sp "failwith \"support for empty messages not implemented\"")
 
-let gen_struct ?and_ t sc =
-  let { Ot.module_prefix; spec; _ } = t in
+let gen_struct ?and_ t sc : bool =
+  let { Ot.spec; _ } = t in
 
   let has_encoded =
     match spec with
     | Ot.Record r ->
-      gen_record ?and_ module_prefix r sc;
+      gen_record ?and_ r sc;
       true
     | Ot.Variant v ->
-      gen_variant ?and_ module_prefix v sc;
+      gen_variant ?and_ v sc;
       true
     | Ot.Const_variant v ->
-      gen_const_variant ?and_ module_prefix v sc;
+      gen_const_variant ?and_ v sc;
       true
     | Ot.Unit u ->
       gen_unit ?and_ u sc;
@@ -248,11 +243,10 @@ let gen_struct ?and_ t sc =
 let gen_sig ?and_ t sc =
   let _ = and_ in
 
-  let { Ot.module_prefix; spec; _ } = t in
+  let { Ot.spec; _ } = t in
 
   let f type_name =
-    F.linep sc "val decode_%s : Js.Json.t Js.Dict.t -> %s_types.%s" type_name
-      module_prefix type_name;
+    F.linep sc "val decode_%s : Js.Json.t Js.Dict.t -> %s" type_name type_name;
     F.linep sc
       ("(** [decode_%s decoder] decodes a " ^^ "[%s] value from [decoder] *)")
       type_name type_name
@@ -266,8 +260,7 @@ let gen_sig ?and_ t sc =
     f v_name;
     true
   | Ot.Const_variant { Ot.cv_name; _ } ->
-    F.linep sc "val decode_%s : Js.Json.t -> %s_types.%s" cv_name module_prefix
-      cv_name;
+    F.linep sc "val decode_%s : Js.Json.t -> %s" cv_name cv_name;
     F.linep sc "(** [decode_%s value] decodes a [%s] from a Json value*)"
       cv_name cv_name;
     true
@@ -276,4 +269,13 @@ let gen_sig ?and_ t sc =
     true
 
 let ocamldoc_title = "BS Decoding"
-let file_suffix = "bs"
+let requires_mutable_records = true
+
+let plugin : Pb_codegen_plugin.t =
+  let module P = struct
+    let gen_sig = gen_sig
+    let gen_struct = gen_struct
+    let ocamldoc_title = ocamldoc_title
+    let requires_mutable_records = requires_mutable_records
+  end in
+  (module P)
diff --git a/src/compilerlib/pb_codegen_decode_bs.mli b/src/compilerlib/pb_codegen_decode_bs.mli
new file mode 100644
index 0000000..7c0a8a1
--- /dev/null
+++ b/src/compilerlib/pb_codegen_decode_bs.mli
@@ -0,0 +1,4 @@
+
+include Pb_codegen_plugin.S
+
+val plugin : Pb_codegen_plugin.t
diff --git a/src/compilerlib/pb_codegen_decode_yojson.ml b/src/compilerlib/pb_codegen_decode_yojson.ml
index fbb43cb..c755ea6 100644
--- a/src/compilerlib/pb_codegen_decode_yojson.ml
+++ b/src/compilerlib/pb_codegen_decode_yojson.ml
@@ -2,7 +2,6 @@ module Ot = Pb_codegen_ocaml_type
 module F = Pb_codegen_formatting
 
 let sp = Pb_codegen_util.sp
-let file_suffix = "yojson"
 
 (* Function which returns all the possible pattern match for reading a JSON
    value into an OCaml value. The protobuf JSON encoding rules
@@ -31,10 +30,8 @@ let field_pattern_match ~r_name ~rf_label field_type =
      runtime functions) *)
   | Ot.Ft_user_defined_type udt ->
     let f_name =
-      let function_prefix = "decode" in
-      let module_suffix = file_suffix in
-      Pb_codegen_util.function_name_of_user_defined ~function_prefix
-        ~module_suffix udt
+      let function_prefix = "decode_json" in
+      Pb_codegen_util.function_name_of_user_defined ~function_prefix udt
     in
     let value_expression = "(" ^ f_name ^ " json_value)" in
     "json_value", value_expression
@@ -58,7 +55,7 @@ let gen_rft_repeated_field sc ~r_name ~rf_label repeated_field =
 
   F.linep sc "| (\"%s\", `List l) -> begin" json_label;
 
-  F.scope sc (fun sc ->
+  F.sub_scope sc (fun sc ->
       F.linep sc "v.%s <- List.map (function" rf_label;
       let match_variable_name, exp =
         field_pattern_match ~r_name ~rf_label field_type
@@ -101,13 +98,13 @@ let gen_rft_variant_field sc ~r_name ~rf_label { Ot.v_constructors; _ } =
     v_constructors
 
 (* Generate decode function for a record *)
-let gen_record ?and_ module_prefix { Ot.r_name; r_fields } sc =
+let gen_record ?and_ { Ot.r_name; r_fields } sc =
   let mutable_record_name = Pb_codegen_util.mutable_record_name r_name in
 
   F.line sc
-  @@ sp "%s decode_%s d =" (Pb_codegen_util.let_decl_of_and and_) r_name;
+  @@ sp "%s decode_json_%s d =" (Pb_codegen_util.let_decl_of_and and_) r_name;
 
-  F.scope sc (fun sc ->
+  F.sub_scope sc (fun sc ->
       F.linep sc "let v = default_%s () in" mutable_record_name;
       F.line sc @@ "let assoc = match d with";
       F.line sc @@ "  | `Assoc assoc -> assoc";
@@ -116,7 +113,7 @@ let gen_record ?and_ module_prefix { Ot.r_name; r_fields } sc =
       F.line sc @@ "in";
 
       F.line sc "List.iter (function ";
-      F.scope sc (fun sc ->
+      F.sub_scope sc (fun sc ->
           (* Generate pattern match for all the possible message field *)
           List.iter
             (fun { Ot.rf_label; rf_field_type; _ } ->
@@ -144,29 +141,28 @@ let gen_record ?and_ module_prefix { Ot.r_name; r_fields } sc =
 
       (* Transform the mutable record in an immutable one *)
       F.line sc "({";
-      F.scope sc (fun sc ->
+      F.sub_scope sc (fun sc ->
           List.iter
             (fun { Ot.rf_label; _ } ->
-              F.linep sc "%s_types.%s = v.%s;" module_prefix rf_label rf_label)
+              F.linep sc "%s = v.%s;" rf_label rf_label)
             r_fields);
-      F.linep sc "} : %s_types.%s)" module_prefix r_name)
+      F.linep sc "} : %s)" r_name)
 
 (* Generate decode function for an empty record *)
 let gen_unit ?and_ { Ot.er_name } sc =
   F.line sc
-  @@ sp "%s decode_%s d =" (Pb_codegen_util.let_decl_of_and and_) er_name;
+  @@ sp "%s decode_json_%s d =" (Pb_codegen_util.let_decl_of_and and_) er_name;
   F.line sc (sp "Pbrt_yojson.unit d \"%s\" \"%s\"" er_name "empty record")
 
 (* Generate decode function for a variant type *)
-let gen_variant ?and_ module_prefix { Ot.v_name; v_constructors } sc =
+let gen_variant ?and_ { Ot.v_name; v_constructors } sc =
   (* helper function for each constructor case *)
   let process_v_constructor sc { Ot.vc_constructor; vc_field_type; _ } =
     let json_label = Pb_codegen_util.camel_case_of_constructor vc_constructor in
 
     match vc_field_type with
     | Ot.Vct_nullary ->
-      F.linep sc "| (\"%s\", _)::_-> (%s_types.%s : %s_types.%s)" json_label
-        module_prefix vc_constructor module_prefix v_name
+      F.linep sc "| (\"%s\", _)::_-> (%s : %s)" json_label vc_constructor v_name
     | Ot.Vct_non_nullary_constructor field_type ->
       let match_, exp =
         let r_name = v_name and rf_label = vc_constructor in
@@ -174,13 +170,14 @@ let gen_variant ?and_ module_prefix { Ot.v_name; v_constructors } sc =
       in
 
       F.linep sc "| (\"%s\", %s)::_ -> " json_label match_;
-      F.linep sc "  (%s_types.%s (%s) : %s_types.%s)" module_prefix
-        vc_constructor exp module_prefix v_name
+      F.linep sc "  (%s (%s) : %s)" vc_constructor exp v_name
   in
 
-  F.linep sc "%s decode_%s json =" (Pb_codegen_util.let_decl_of_and and_) v_name;
+  F.linep sc "%s decode_json_%s json ="
+    (Pb_codegen_util.let_decl_of_and and_)
+    v_name;
 
-  F.scope sc (fun sc ->
+  F.sub_scope sc (fun sc ->
       (* even though a variant should be an object with a single field,
        * it is possible other fields are present in the JSON object. Therefore
        * we still need a loop to iterate over the key/value, even if in 99.99%
@@ -192,7 +189,7 @@ let gen_variant ?and_ module_prefix { Ot.v_name; v_constructors } sc =
       F.line sc "in";
 
       F.line sc "let rec loop = function";
-      F.scope sc (fun sc ->
+      F.sub_scope sc (fun sc ->
           (* termination condition *)
           F.linep sc "| [] -> Pbrt_yojson.E.malformed_variant \"%s\"" v_name;
 
@@ -203,32 +200,32 @@ let gen_variant ?and_ module_prefix { Ot.v_name; v_constructors } sc =
       F.line sc "in";
       F.line sc "loop assoc")
 
-let gen_const_variant ?and_ module_prefix { Ot.cv_name; cv_constructors } sc =
-  F.linep sc "%s decode_%s json ="
+let gen_const_variant ?and_ { Ot.cv_name; cv_constructors } sc =
+  F.linep sc "%s decode_json_%s json ="
     (Pb_codegen_util.let_decl_of_and and_)
     cv_name;
 
-  F.scope sc (fun sc ->
+  F.sub_scope sc (fun sc ->
       F.line sc "match json with";
       List.iter
         (fun { Ot.cvc_name; cvc_string_value; _ } ->
-          F.linep sc "| `String \"%s\" -> (%s_types.%s : %s_types.%s)"
-            cvc_string_value module_prefix cvc_name module_prefix cv_name)
+          F.linep sc "| `String \"%s\" -> (%s : %s)" cvc_string_value cvc_name
+            cv_name)
         cv_constructors;
       F.linep sc "| _ -> Pbrt_yojson.E.malformed_variant \"%s\"" cv_name)
 
 let gen_struct ?and_ t sc =
-  let { Ot.module_prefix; spec; _ } = t in
+  let { Ot.spec; _ } = t in
   let has_encoded =
     match spec with
     | Ot.Record r ->
-      gen_record ?and_ module_prefix r sc;
+      gen_record ?and_ r sc;
       true
     | Ot.Variant v ->
-      gen_variant ?and_ module_prefix v sc;
+      gen_variant ?and_ v sc;
       true
     | Ot.Const_variant v ->
-      gen_const_variant ?and_ module_prefix v sc;
+      gen_const_variant ?and_ v sc;
       true
     | Ot.Unit u ->
       gen_unit ?and_ u sc;
@@ -238,13 +235,13 @@ let gen_struct ?and_ t sc =
 
 let gen_sig ?and_ t sc =
   let _ = and_ in
-  let { Ot.module_prefix; spec; _ } = t in
+  let { Ot.spec; _ } = t in
 
   let f type_name =
-    F.linep sc "val decode_%s : Yojson.Basic.json -> %s_types.%s" type_name
-      module_prefix type_name;
+    F.linep sc "val decode_json_%s : Yojson.Basic.t -> %s" type_name type_name;
     F.linep sc
-      ("(** [decode_%s decoder] decodes a " ^^ "[%s] value from [decoder] *)")
+      ("(** [decode_json_%s decoder] decodes a "
+     ^^ "[%s] value from [decoder] *)")
       type_name type_name
   in
 
@@ -263,3 +260,13 @@ let gen_sig ?and_ t sc =
     true
 
 let ocamldoc_title = "JSON Decoding"
+let requires_mutable_records = true
+
+let plugin : Pb_codegen_plugin.t =
+  let module P = struct
+    let gen_sig = gen_sig
+    let gen_struct = gen_struct
+    let ocamldoc_title = ocamldoc_title
+    let requires_mutable_records = requires_mutable_records
+  end in
+  (module P)
diff --git a/src/compilerlib/pb_codegen_decode_yojson.mli b/src/compilerlib/pb_codegen_decode_yojson.mli
index 8bf43d6..fcff7c2 100644
--- a/src/compilerlib/pb_codegen_decode_yojson.mli
+++ b/src/compilerlib/pb_codegen_decode_yojson.mli
@@ -1,3 +1,5 @@
 (** Code generator for the [decode] JSON function *)
 
-include Pb_codegen_sig.S
+include Pb_codegen_plugin.S
+
+val plugin : Pb_codegen_plugin.t
diff --git a/src/compilerlib/pb_codegen_default.ml b/src/compilerlib/pb_codegen_default.ml
index 7b4ef51..19d0b4a 100644
--- a/src/compilerlib/pb_codegen_default.ml
+++ b/src/compilerlib/pb_codegen_default.ml
@@ -3,8 +3,6 @@ module F = Pb_codegen_formatting
 module E = Pb_exception
 open Pb_codegen_util
 
-let file_suffix = Pb_codegen_types.file_suffix
-
 let default_value_of_basic_type ?field_name basic_type field_default =
   match basic_type, field_default with
   | Ot.Bt_string, None -> "\"\""
@@ -30,35 +28,25 @@ let default_value_of_basic_type ?field_name basic_type field_default =
 
 (* Generate the string which is the default value for a given field
    type and default information. *)
-let default_value_of_field_type ~gen_file_suffix ~module_prefix ?field_name
-    field_type field_default =
+let default_value_of_field_type ?field_name field_type field_default =
   match field_type with
   | Ot.Ft_user_defined_type udt ->
-    let f_name =
-      let function_prefix = "default" in
-      let module_suffix = file_suffix in
-      function_name_of_user_defined ~function_prefix ~module_suffix udt ^ " ()"
-    in
-
-    let same_file = gen_file_suffix = file_suffix in
-    (match same_file, udt.Ot.udt_module_prefix with
-    | false, None -> module_prefix ^ "_" ^ file_suffix ^ "." ^ f_name
-    | _ -> f_name)
+    let function_prefix = "default" in
+    function_name_of_user_defined ~function_prefix udt ^ " ()"
   | Ot.Ft_unit -> "()"
   | Ot.Ft_basic_type bt ->
     default_value_of_basic_type ?field_name bt field_default
   | Ot.Ft_wrapper_type _ -> "None"
 
-(* This function returns [(field_name, field_default_value, field_type)] for
+(** This function returns [(field_name, field_default_value, field_type)] for
    a record field. *)
-let record_field_default_info ~gen_file_suffix ~module_prefix record_field =
+let record_field_default_info record_field =
   let { Ot.rf_label; Ot.rf_field_type; _ } = record_field in
   let type_string = Pb_codegen_util.string_of_record_field_type rf_field_type in
   let field_name = rf_label in
 
   let dfvft field_type field_default =
-    default_value_of_field_type ~gen_file_suffix ~module_prefix ~field_name
-      field_type field_default
+    default_value_of_field_type ~field_name field_type field_default
   in
 
   let default_value =
@@ -79,62 +67,48 @@ let record_field_default_info ~gen_file_suffix ~module_prefix record_field =
       (match at with
       | Ot.At_list -> "[]"
       | Ot.At_hashtable -> "Hashtbl.create 128")
-      (* TODO This initial value could be configurable either via 
+      (* TODO This initial value could be configurable either via
        * the default function or via a protobuf option. *)
     | Ot.Rft_variant { Ot.v_constructors; _ } ->
       (match v_constructors with
       | [] -> assert false
       | { Ot.vc_constructor; vc_field_type; _ } :: _ ->
-        let default_value =
-          match vc_field_type with
-          | Ot.Vct_nullary -> vc_constructor
-          | Ot.Vct_non_nullary_constructor field_type ->
-            sp "%s (%s)" vc_constructor (dfvft field_type None)
-        in
-        if gen_file_suffix = Pb_codegen_types.file_suffix then
-          default_value
-        else
-          module_prefix ^ "_types." ^ default_value)
+        (match vc_field_type with
+        | Ot.Vct_nullary -> vc_constructor
+        | Ot.Vct_non_nullary_constructor field_type ->
+          sp "%s (%s)" vc_constructor (dfvft field_type None)))
   in
   field_name, default_value, type_string
 
-let gen_record_mutable ~gen_file_suffix ~module_prefix { Ot.r_name; r_fields }
-    sc =
+let gen_record_mutable { Ot.r_name; r_fields } sc : unit =
   let fields_default_info =
-    List.map
-      (fun r_field ->
-        record_field_default_info ~gen_file_suffix ~module_prefix r_field)
-      r_fields
+    List.map (fun r_field -> record_field_default_info r_field) r_fields
   in
 
   let rn = Pb_codegen_util.mutable_record_name r_name in
   F.linep sc "let default_%s () : %s = {" rn rn;
 
-  F.scope sc (fun sc ->
+  F.sub_scope sc (fun sc ->
       List.iter
         (fun (fname, fvalue, _) -> F.linep sc "%s = %s;" fname fvalue)
         fields_default_info);
   F.line sc "}"
 
-let gen_record ?and_ module_prefix { Ot.r_name; r_fields } sc =
+let gen_record ?and_ { Ot.r_name; r_fields } sc : unit =
   let fields_default_info =
-    List.map
-      (fun r_field ->
-        record_field_default_info ~gen_file_suffix:file_suffix ~module_prefix
-          r_field)
-      r_fields
+    List.map (fun r_field -> record_field_default_info r_field) r_fields
   in
 
   F.linep sc "%s default_%s " (let_decl_of_and and_) r_name;
 
-  F.scope sc (fun sc ->
+  F.sub_scope sc (fun sc ->
       List.iter
         (fun (fname, fvalue, ftype) ->
           F.linep sc "?%s:((%s:%s) = %s)" fname fname ftype fvalue)
         fields_default_info;
       F.linep sc "() : %s  = {" r_name);
 
-  F.scope sc (fun sc ->
+  F.sub_scope sc (fun sc ->
       List.iter
         (fun (fname, _, _) -> F.linep sc "%s;" fname)
         fields_default_info);
@@ -144,7 +118,7 @@ let gen_record ?and_ module_prefix { Ot.r_name; r_fields } sc =
 let gen_unit ?and_ { Ot.er_name } sc =
   F.linep sc "%s default_%s = ()" (let_decl_of_and and_) er_name
 
-let gen_variant ?and_ module_prefix { Ot.v_name; Ot.v_constructors } sc =
+let gen_variant ?and_ { Ot.v_name; Ot.v_constructors } sc =
   match v_constructors with
   | [] -> failwith "programmatic TODO error"
   | { Ot.vc_constructor; vc_field_type; _ } :: _ ->
@@ -155,9 +129,7 @@ let gen_variant ?and_ module_prefix { Ot.v_name; Ot.v_constructors } sc =
     | Ot.Vct_non_nullary_constructor field_type ->
       let default_value =
         let field_name = v_name in
-        let gen_file_suffix = file_suffix in
-        default_value_of_field_type ~gen_file_suffix ~module_prefix ~field_name
-          field_type None
+        default_value_of_field_type ~field_name field_type None
       in
       (* TODO need to fix the deault value *)
       F.linep sc "%s default_%s () : %s = %s (%s)" decl v_name v_name
@@ -172,16 +144,17 @@ let gen_const_variant ?and_ { Ot.cv_name; Ot.cv_constructors } sc =
   F.linep sc "%s default_%s () = (%s:%s)" (let_decl_of_and and_) cv_name
     first_constructor_name cv_name
 
-let gen_struct ?and_ t sc =
-  let { Ot.module_prefix; spec; _ } = t in
+let gen_struct_full ~with_mutable_records ?and_ t sc =
+  let { Ot.spec; _ } = t in
 
   let has_encoded =
     match spec with
     | Ot.Record r ->
-      gen_record ?and_ module_prefix r sc;
+      gen_record ?and_ r sc;
+      if with_mutable_records then gen_record_mutable r sc;
       true
     | Ot.Variant v ->
-      gen_variant ?and_ module_prefix v sc;
+      gen_variant ?and_ v sc;
       true
     | Ot.Const_variant v ->
       gen_const_variant v sc;
@@ -192,18 +165,17 @@ let gen_struct ?and_ t sc =
   in
   has_encoded
 
-let gen_sig_record sc module_prefix { Ot.r_name; r_fields } =
+let gen_struct ?and_ t sc =
+  gen_struct_full ?and_ ~with_mutable_records:false t sc
+
+let gen_sig_record sc { Ot.r_name; r_fields } =
   F.linep sc "val default_%s : " r_name;
 
-  let fields_default_info =
-    List.map
-      (fun r_field ->
-        record_field_default_info ~gen_file_suffix:file_suffix ~module_prefix
-          r_field)
-      r_fields
+  let fields_default_info : _ list =
+    List.map (fun r_field -> record_field_default_info r_field) r_fields
   in
 
-  F.scope sc (fun sc ->
+  F.sub_scope sc (fun sc ->
       List.iter
         (fun (field_name, _, field_type) ->
           F.linep sc "?%s:%s ->" field_name field_type)
@@ -219,20 +191,19 @@ let gen_sig_unit sc { Ot.er_name } =
   let rn = er_name in
   F.linep sc "(** [default_%s ()] is the default value for type [%s] *)" rn rn
 
-let gen_sig ?and_ t sc =
-  let _ = and_ in
+let gen_sig ?and_:_ t sc =
   let f type_name =
     F.linep sc "val default_%s : unit -> %s" type_name type_name;
     F.linep sc "(** [default_%s ()] is the default value for type [%s] *)"
       type_name type_name
   in
 
-  let { Ot.spec; module_prefix; _ } = t in
+  let { Ot.spec; _ } = t in
 
   let has_encoded =
     match spec with
     | Ot.Record r ->
-      gen_sig_record sc module_prefix r;
+      gen_sig_record sc r;
       true
     | Ot.Variant v ->
       f v.Ot.v_name;
@@ -247,4 +218,5 @@ let gen_sig ?and_ t sc =
 
   has_encoded
 
-let ocamldoc_title = "Default values"
+let ocamldoc_title = "Basic values"
+let requires_mutable_records = false
diff --git a/src/compilerlib/pb_codegen_default.mli b/src/compilerlib/pb_codegen_default.mli
index eff317b..67b71e8 100644
--- a/src/compilerlib/pb_codegen_default.mli
+++ b/src/compilerlib/pb_codegen_default.mli
@@ -1,10 +1,6 @@
-(** Code generator for the [default] function *)
+(** Code generator for the [default] functions (i.e builders) *)
 
-include Pb_codegen_sig.S
+include Pb_codegen_plugin.S
 
 val gen_record_mutable :
-  gen_file_suffix:string ->
-  module_prefix:string ->
-  Pb_codegen_ocaml_type.record ->
-  Pb_codegen_formatting.scope ->
-  unit
+  Pb_codegen_ocaml_type.record -> Pb_codegen_formatting.scope -> unit
diff --git a/src/compilerlib/pb_codegen_encode_binary.ml b/src/compilerlib/pb_codegen_encode_binary.ml
index 98e13e2..d7fd7f3 100644
--- a/src/compilerlib/pb_codegen_encode_binary.ml
+++ b/src/compilerlib/pb_codegen_encode_binary.ml
@@ -1,7 +1,6 @@
 module Ot = Pb_codegen_ocaml_type
 module F = Pb_codegen_formatting
 
-let file_suffix = Pb_codegen_decode_binary.file_suffix
 let constructor_name s = String.capitalize_ascii @@ String.lowercase_ascii s
 let sp = Pb_codegen_util.sp
 
@@ -60,10 +59,8 @@ let gen_encode_field_type ?with_key sc var_name encoding_number pk is_packed
   | Ot.Ft_user_defined_type udt ->
     encode_key sc;
     let f_name =
-      let function_prefix = "encode" in
-      let module_suffix = file_suffix in
-      Pb_codegen_util.function_name_of_user_defined ~function_prefix
-        ~module_suffix udt
+      let function_prefix = "encode_pb" in
+      Pb_codegen_util.function_name_of_user_defined ~function_prefix udt
     in
     (match udt.Ot.udt_type with
     | `Message ->
@@ -92,7 +89,7 @@ let gen_rft_required sc var_name (field_type, encoding_number, pk, _) =
 let gen_rft_optional sc var_name (field_type, encoding_number, pk, _) =
   F.linep sc "begin match %s with" var_name;
   F.linep sc "| Some x -> ";
-  F.scope sc (fun sc ->
+  F.sub_scope sc (fun sc ->
       gen_encode_field_type ~with_key:() sc "x" encoding_number pk false
         field_type);
   F.line sc "| None -> ();";
@@ -104,13 +101,13 @@ let gen_rft_repeated sc var_name repeated_field =
   match rt, is_packed with
   | Ot.Rt_list, false ->
     F.line sc "List.iter (fun x -> ";
-    F.scope sc (fun sc ->
+    F.sub_scope sc (fun sc ->
         gen_encode_field_type ~with_key:() sc "x" encoding_number pk is_packed
           field_type);
     F.linep sc ") %s;" var_name
   | Ot.Rt_repeated_field, false ->
     F.line sc "Pbrt.Repeated_field.iter (fun x -> ";
-    F.scope sc (fun sc ->
+    F.sub_scope sc (fun sc ->
         gen_encode_field_type ~with_key:() sc "x" encoding_number pk is_packed
           field_type);
     F.linep sc ") %s;" var_name
@@ -119,9 +116,9 @@ let gen_rft_repeated sc var_name repeated_field =
     (* When packed the key is encoded once.
        *)
     F.line sc "Pbrt.Encoder.nested (fun encoder ->";
-    F.scope sc (fun sc ->
+    F.sub_scope sc (fun sc ->
         F.line sc "List.iter (fun x -> ";
-        F.scope sc (fun sc ->
+        F.sub_scope sc (fun sc ->
             gen_encode_field_type sc "x" encoding_number pk is_packed field_type);
         F.linep sc ") %s;" var_name);
     F.line sc ") encoder;"
@@ -130,14 +127,14 @@ let gen_rft_repeated sc var_name repeated_field =
     (* When packed the key is encoded once.
        *)
     F.line sc "Pbrt.Encoder.nested (fun encoder ->";
-    F.scope sc (fun sc ->
+    F.sub_scope sc (fun sc ->
         F.line sc "Pbrt.Repeated_field.iter (fun x -> ";
-        F.scope sc (fun sc ->
+        F.sub_scope sc (fun sc ->
             gen_encode_field_type sc "x" encoding_number pk is_packed field_type);
         F.linep sc ") %s;" var_name);
     F.line sc ") encoder;"
 
-let gen_rft_variant sc module_prefix var_name { Ot.v_constructors; _ } =
+let gen_rft_variant sc var_name { Ot.v_constructors; _ } =
   F.linep sc "begin match %s with" var_name;
   List.iter
     (fun constructor ->
@@ -152,13 +149,13 @@ let gen_rft_variant sc module_prefix var_name { Ot.v_constructors; _ } =
 
       match vc_field_type with
       | Ot.Vct_nullary ->
-        F.linep sc "| %s_types.%s ->" module_prefix vc_constructor;
-        F.scope sc (fun sc ->
+        F.linep sc "| %s ->" vc_constructor;
+        F.sub_scope sc (fun sc ->
             gen_encode_field_key sc vc_encoding_number vc_payload_kind false;
             F.line sc "Pbrt.Encoder.empty_nested encoder")
       | Ot.Vct_non_nullary_constructor field_type ->
-        F.linep sc "| %s_types.%s x ->" module_prefix vc_constructor;
-        F.scope sc (fun sc ->
+        F.linep sc "| %s x ->" vc_constructor;
+        F.sub_scope sc (fun sc ->
             gen_encode_field_type sc ~with_key:() "x" vc_encoding_number
               vc_payload_kind false field_type))
     v_constructors;
@@ -172,14 +169,14 @@ let gen_rft_associative sc var_name associative_field =
   F.linep sc "let encode_key = %s in"
     (runtime_function_for_basic_type key_type key_pk);
   F.line sc "let encode_value = (fun x encoder ->";
-  F.scope sc (fun sc ->
+  F.sub_scope sc (fun sc ->
       gen_encode_field_type sc "x" (-1 (* TODO *)) value_pk false value_type);
   F.line sc ") in";
 
   (match at with
   | Ot.At_list -> F.line sc "List.iter (fun (k, v) ->"
   | Ot.At_hashtable -> F.line sc "Hashtbl.iter (fun k v ->");
-  F.scope sc (fun sc ->
+  F.sub_scope sc (fun sc ->
       gen_encode_field_key sc encoding_number Ot.Pk_bytes false;
       F.linep sc "let map_entry = (k, Pbrt.%s), (v, Pbrt.%s) in"
         (Pb_codegen_util.string_of_payload_kind ~capitalize:() key_pk false)
@@ -189,43 +186,43 @@ let gen_rft_associative sc var_name associative_field =
        ^ "~encode_value map_entry encoder"));
   F.linep sc ") %s;" var_name
 
-let gen_record ?and_ module_prefix { Ot.r_name; r_fields } sc =
+let gen_record ?and_ { Ot.r_name; r_fields } sc =
   let rn = r_name in
-  F.linep sc "%s encode_%s (v:%s_types.%s) encoder = "
+  F.linep sc "%s encode_pb_%s (v:%s) encoder = "
     (Pb_codegen_util.let_decl_of_and and_)
-    rn module_prefix rn;
+    rn rn;
 
-  F.scope sc (fun sc ->
+  F.sub_scope sc (fun sc ->
       List.iter
         (fun record_field ->
           let { Ot.rf_label; rf_field_type; _ } = record_field in
 
-          let var_name = sp "v.%s_types.%s" module_prefix rf_label in
+          let var_name = sp "v.%s" rf_label in
           match rf_field_type with
           | Ot.Rft_nolabel x -> gen_rft_nolabel sc var_name x
           | Ot.Rft_required x -> gen_rft_required sc var_name x
           | Ot.Rft_optional x -> gen_rft_optional sc var_name x
           | Ot.Rft_repeated x -> gen_rft_repeated sc var_name x
-          | Ot.Rft_variant x -> gen_rft_variant sc module_prefix var_name x
+          | Ot.Rft_variant x -> gen_rft_variant sc var_name x
           | Ot.Rft_associative x -> gen_rft_associative sc var_name x)
         r_fields (* List.iter *);
       F.line sc "()")
 (* encode function *)
 
-let gen_unit ?and_ module_prefix { Ot.er_name } sc =
+let gen_unit ?and_ { Ot.er_name } sc =
   let rn = er_name in
-  F.linep sc "%s encode_%s (v:%s_types.%s) encoder = "
+  F.linep sc "%s encode_pb_%s (v:%s) encoder = "
     (Pb_codegen_util.let_decl_of_and and_)
-    rn module_prefix rn;
+    rn rn;
   F.line sc "()"
 
-let gen_variant ?and_ module_prefix variant sc =
+let gen_variant ?and_ variant sc =
   let { Ot.v_name; Ot.v_constructors } = variant in
   let vn = v_name in
-  F.linep sc "%s encode_%s (v:%s_types.%s) encoder = "
+  F.linep sc "%s encode_pb_%s (v:%s) encoder = "
     (Pb_codegen_util.let_decl_of_and and_)
-    vn module_prefix vn;
-  F.scope sc (fun sc ->
+    vn vn;
+  F.sub_scope sc (fun sc ->
       F.line sc "begin match v with";
       List.iter
         (fun variant_constructor ->
@@ -239,64 +236,64 @@ let gen_variant ?and_ module_prefix variant sc =
           in
           match vc_field_type with
           | Ot.Vct_nullary ->
-            F.linep sc "| %s_types.%s ->" module_prefix vc_constructor;
-            F.scope sc (fun sc ->
+            F.linep sc "| %s ->" vc_constructor;
+            F.sub_scope sc (fun sc ->
                 gen_encode_field_key sc vc_encoding_number vc_payload_kind false;
                 F.line sc "Pbrt.Encoder.empty_nested encoder")
           | Ot.Vct_non_nullary_constructor field_type ->
-            F.linep sc "| %s_types.%s x ->" module_prefix vc_constructor;
-            F.scope sc (fun sc ->
+            F.linep sc "| %s x ->" vc_constructor;
+            F.sub_scope sc (fun sc ->
                 gen_encode_field_type sc ~with_key:() "x" vc_encoding_number
                   vc_payload_kind false field_type))
         v_constructors;
       F.line sc "end")
 
-let gen_const_variant ?and_ module_prefix cv sc =
+let gen_const_variant ?and_ cv sc =
   let { Ot.cv_name; cv_constructors } = cv in
 
-  F.linep sc "%s encode_%s (v:%s_types.%s) encoder ="
+  F.linep sc "%s encode_pb_%s (v:%s) encoder ="
     (Pb_codegen_util.let_decl_of_and and_)
-    cv_name module_prefix cv_name;
-  F.scope sc (fun sc ->
+    cv_name cv_name;
+  F.sub_scope sc (fun sc ->
       F.line sc "match v with";
       List.iter
         (fun { Ot.cvc_name; cvc_binary_value; _ } ->
           F.line sc
             (if cvc_binary_value > 0 then
-              sp "| %s_types.%s -> Pbrt.Encoder.int_as_varint %i encoder"
-                module_prefix cvc_name cvc_binary_value
+              sp "| %s -> Pbrt.Encoder.int_as_varint %i encoder" cvc_name
+                cvc_binary_value
             else
-              sp "| %s_types.%s -> Pbrt.Encoder.int_as_varint (%i) encoder"
-                module_prefix cvc_name cvc_binary_value))
+              sp "| %s -> Pbrt.Encoder.int_as_varint (%i) encoder" cvc_name
+                cvc_binary_value))
         cv_constructors)
 
 let gen_struct ?and_ t sc =
-  let { Ot.module_prefix; spec; _ } = t in
+  let { Ot.spec; _ } = t in
   let has_encoded =
     match spec with
     | Ot.Record r ->
-      gen_record ?and_ module_prefix r sc;
+      gen_record ?and_ r sc;
       true
     | Ot.Variant v ->
-      gen_variant ?and_ module_prefix v sc;
+      gen_variant ?and_ v sc;
       true
     | Ot.Const_variant v ->
-      gen_const_variant ?and_ module_prefix v sc;
+      gen_const_variant ?and_ v sc;
       true
     | Ot.Unit u ->
-      gen_unit ?and_ module_prefix u sc;
+      gen_unit ?and_ u sc;
       true
   in
   has_encoded
 
 let gen_sig ?and_ t sc =
   let _ = and_ in
-  let { Ot.module_prefix; spec; _ } = t in
+  let { Ot.spec; _ } = t in
   let f type_name =
-    F.linep sc "val encode_%s : %s_types.%s -> Pbrt.Encoder.t -> unit" type_name
-      module_prefix type_name;
+    F.linep sc "val encode_pb_%s : %s -> Pbrt.Encoder.t -> unit" type_name
+      type_name;
     F.linep sc
-      "(** [encode_%s v encoder] encodes [v] with the given [encoder] *)"
+      "(** [encode_pb_%s v encoder] encodes [v] with the given [encoder] *)"
       type_name
   in
   let has_encoded =
@@ -317,3 +314,13 @@ let gen_sig ?and_ t sc =
   has_encoded
 
 let ocamldoc_title = "Protobuf Encoding"
+let requires_mutable_records = false
+
+let plugin : Pb_codegen_plugin.t =
+  let module P = struct
+    let gen_sig = gen_sig
+    let gen_struct = gen_struct
+    let ocamldoc_title = ocamldoc_title
+    let requires_mutable_records = requires_mutable_records
+  end in
+  (module P)
diff --git a/src/compilerlib/pb_codegen_encode_binary.mli b/src/compilerlib/pb_codegen_encode_binary.mli
index 1efa4b3..c05ae8c 100644
--- a/src/compilerlib/pb_codegen_encode_binary.mli
+++ b/src/compilerlib/pb_codegen_encode_binary.mli
@@ -1,3 +1,5 @@
 (** Code generator for the [encode] function *)
 
-include Pb_codegen_sig.S
+include Pb_codegen_plugin.S
+
+val plugin : Pb_codegen_plugin.t
diff --git a/src/compilerlib/pb_codegen_encode_bs.ml b/src/compilerlib/pb_codegen_encode_bs.ml
index b29349d..5c30c9c 100644
--- a/src/compilerlib/pb_codegen_encode_bs.ml
+++ b/src/compilerlib/pb_codegen_encode_bs.ml
@@ -42,7 +42,7 @@ let setter_of_basic_type json_label basic_type pk =
   | x, y -> unsupported2 json_label x y
 
 let gen_field sc var_name json_label field_type pk =
-  (* reusable function for both basic type and wrapped type which 
+  (* reusable function for both basic type and wrapped type which
    * are both based upon basic type *)
   let basic_type_statement basic_type var_name pk =
     let setter, map_function = setter_of_basic_type json_label basic_type pk in
@@ -66,14 +66,12 @@ let gen_field sc var_name json_label field_type pk =
     let { Ot.udt_type; _ } = udt in
     let f_name =
       let function_prefix = "encode" in
-      let module_suffix = "bs" in
-      Pb_codegen_util.function_name_of_user_defined ~function_prefix
-        ~module_suffix udt
+      Pb_codegen_util.function_name_of_user_defined ~function_prefix udt
     in
     (match udt_type with
     | `Message ->
       F.linep sc "begin (* %s field *)" json_label;
-      F.scope sc (fun sc ->
+      F.sub_scope sc (fun sc ->
           F.linep sc "let json' = %s %s in" f_name var_name;
           F.linep sc "Js.Dict.set json \"%s\" (Js.Json.object_ json');"
             json_label);
@@ -84,7 +82,7 @@ let gen_field sc var_name json_label field_type pk =
   | Ot.Ft_wrapper_type wrapper_type, _ ->
     let { Ot.wt_type; Ot.wt_pk } = wrapper_type in
     F.line sc "begin";
-    F.scope sc (fun sc ->
+    F.sub_scope sc (fun sc ->
         F.linep sc "match %s with" var_name;
         F.linep sc "| None -> Js.Dict.set json \"%s\" Js.Json.null" json_label;
         let statement = basic_type_statement wt_type "__x__" wt_pk in
@@ -99,7 +97,7 @@ let gen_rft_optional sc var_name rf_label (field_type, _, pk, _) =
   F.linep sc "begin match %s with" var_name;
   F.line sc "| None -> ()";
   F.line sc "| Some v ->";
-  F.scope sc (fun sc ->
+  F.sub_scope sc (fun sc ->
       let json_label = Pb_codegen_util.camel_case_of_label rf_label in
       gen_field sc "v" json_label field_type pk);
   F.line sc "end;"
@@ -117,7 +115,7 @@ let gen_rft_repeated sc var_name rf_label repeated_field =
   F.linep sc "begin match %s with" var_name;
   F.line sc "| [] -> ()";
   F.linep sc "| __x__ -> (* %s *)" json_label;
-  F.scope sc (fun sc ->
+  F.sub_scope sc (fun sc ->
       match field_type, pk with
       | Ot.Ft_unit, _ -> unsupported json_label
       | Ot.Ft_basic_type basic_type, _ ->
@@ -139,16 +137,14 @@ let gen_rft_repeated sc var_name rf_label repeated_field =
       | Ot.Ft_user_defined_type udt, Ot.Pk_bytes ->
         let f_name =
           let function_prefix = "encode" in
-          let module_suffix = "bs" in
-          Pb_codegen_util.function_name_of_user_defined ~function_prefix
-            ~module_suffix udt
+          Pb_codegen_util.function_name_of_user_defined ~function_prefix udt
         in
         F.linep sc "let (%s':Js.Json.t) =" rf_label;
-        F.scope sc (fun sc ->
+        F.sub_scope sc (fun sc ->
             F.line sc "__x__";
             F.line sc "|> Array.of_list";
             F.linep sc "|> Array.map (fun v ->";
-            F.scope sc (fun sc ->
+            F.sub_scope sc (fun sc ->
                 F.linep sc "v |> %s |> Js.Json.object_" f_name);
             F.line sc ")";
             F.line sc "|> Js.Json.array");
@@ -159,7 +155,7 @@ let gen_rft_repeated sc var_name rf_label repeated_field =
 
 let gen_rft_variant sc var_name rf_label { Ot.v_constructors; _ } =
   F.linep sc "begin match %s with" var_name;
-  F.scope sc (fun sc ->
+  F.sub_scope sc (fun sc ->
       List.iter
         (fun { Ot.vc_constructor; vc_field_type; vc_payload_kind; _ } ->
           let var_name = "v" in
@@ -167,7 +163,7 @@ let gen_rft_variant sc var_name rf_label { Ot.v_constructors; _ } =
             Pb_codegen_util.camel_case_of_constructor vc_constructor
           in
           F.linep sc "| %s v ->" vc_constructor;
-          F.scope sc (fun sc ->
+          F.sub_scope sc (fun sc ->
               match vc_field_type with
               | Ot.Vct_nullary ->
                 F.linep sc "Js.Dict.set json \"%s\" Js.Json.null" json_label
@@ -176,17 +172,17 @@ let gen_rft_variant sc var_name rf_label { Ot.v_constructors; _ } =
         v_constructors);
   F.linep sc "end; (* match v.%s *)" rf_label
 
-let gen_record ?and_ module_prefix { Ot.r_name; r_fields } sc =
+let gen_record ?and_ { Ot.r_name; r_fields } sc =
   let rn = r_name in
-  F.linep sc "%s encode_%s (v:%s_types.%s) = "
+  F.linep sc "%s encode_%s (v:%s) = "
     (Pb_codegen_util.let_decl_of_and and_)
-    rn module_prefix rn;
-  F.scope sc (fun sc ->
+    rn rn;
+  F.sub_scope sc (fun sc ->
       F.line sc "let json = Js.Dict.empty () in";
       List.iter
         (fun record_field ->
           let { Ot.rf_label; rf_field_type; _ } = record_field in
-          let var_name = sp "v.%s_types.%s" module_prefix rf_label in
+          let var_name = sp "v.%s" rf_label in
 
           match rf_field_type with
           | Ot.Rft_nolabel nolabel_field ->
@@ -206,7 +202,7 @@ let gen_record ?and_ module_prefix { Ot.r_name; r_fields } sc =
         r_fields (* List.iter *);
       F.line sc "json")
 
-let gen_variant ?and_ module_prefix { Ot.v_name; v_constructors } sc =
+let gen_variant ?and_ { Ot.v_name; v_constructors } sc =
   let process_v_constructor sc v_constructor =
     let { Ot.vc_constructor; Ot.vc_field_type; Ot.vc_payload_kind; _ } =
       v_constructor
@@ -216,70 +212,67 @@ let gen_variant ?and_ module_prefix { Ot.v_name; v_constructors } sc =
 
     match vc_field_type with
     | Ot.Vct_nullary ->
-      F.linep sc "| %s_types.%s ->" module_prefix vc_constructor;
+      F.linep sc "| %s ->" vc_constructor;
       F.linep sc "  Js.Dict.set json \"%s\" Js.Json.null" json_label
     | Ot.Vct_non_nullary_constructor field_type ->
-      F.linep sc "| %s_types.%s v ->" module_prefix vc_constructor;
-      F.scope sc (fun sc ->
+      F.linep sc "| %s v ->" vc_constructor;
+      F.sub_scope sc (fun sc ->
           gen_field sc "v" json_label field_type vc_payload_kind)
   in
 
-  F.linep sc "%s encode_%s (v:%s_types.%s) = "
+  F.linep sc "%s encode_%s (v:%s) = "
     (Pb_codegen_util.let_decl_of_and and_)
-    v_name module_prefix v_name;
-  F.scope sc (fun sc ->
+    v_name v_name;
+  F.sub_scope sc (fun sc ->
       F.line sc "let json = Js.Dict.empty () in";
       F.line sc "begin match v with";
       List.iter (process_v_constructor sc) v_constructors;
       F.line sc "end;";
       F.line sc "json")
 
-let gen_const_variant ?and_ module_prefix { Ot.cv_name; Ot.cv_constructors } sc
-    =
-  F.linep sc "%s encode_%s (v:%s_types.%s) : string = "
+let gen_const_variant ?and_ { Ot.cv_name; Ot.cv_constructors } sc =
+  F.linep sc "%s encode_%s (v:%s) : string = "
     (Pb_codegen_util.let_decl_of_and and_)
-    cv_name module_prefix cv_name;
-  F.scope sc (fun sc ->
+    cv_name cv_name;
+  F.sub_scope sc (fun sc ->
       F.line sc "match v with";
       List.iter
         (fun { Ot.cvc_name; cvc_string_value; _ } ->
-          F.linep sc "| %s_types.%s -> \"%s\"" module_prefix cvc_name
-            cvc_string_value)
+          F.linep sc "| %s -> \"%s\"" cvc_name cvc_string_value)
         cv_constructors)
 
-let gen_unit ?and_ module_prefix { Ot.er_name } sc =
+let gen_unit ?and_ { Ot.er_name } sc =
   let rn = er_name in
-  F.linep sc "%s encode_%s (_v:%s_types.%s) = "
+  F.linep sc "%s encode_%s (_v:%s) = "
     (Pb_codegen_util.let_decl_of_and and_)
-    rn module_prefix rn;
+    rn rn;
   F.line sc "Js.Json.null"
 
-let gen_struct ?and_ t sc =
-  let { Ot.module_prefix; spec; _ } = t in
+let gen_struct ?and_ t sc : bool =
+  let { Ot.spec; _ } = t in
 
   let has_encoded =
     match spec with
     | Ot.Record r ->
-      gen_record ?and_ module_prefix r sc;
+      gen_record ?and_ r sc;
       true
     | Ot.Variant v ->
-      gen_variant ?and_ module_prefix v sc;
+      gen_variant ?and_ v sc;
       true
     | Ot.Const_variant v ->
-      gen_const_variant ?and_ module_prefix v sc;
+      gen_const_variant ?and_ v sc;
       true
     | Ot.Unit v ->
-      gen_unit ?and_ module_prefix v sc;
+      gen_unit ?and_ v sc;
       true
   in
   has_encoded
 
 let gen_sig ?and_ t sc =
   let _ = and_ in
-  let { Ot.module_prefix; spec; _ } = t in
+  let { Ot.spec; _ } = t in
   let f type_name =
-    F.linep sc "val encode_%s : %s_types.%s -> Js.Json.t Js.Dict.t" type_name
-      module_prefix type_name;
+    F.linep sc "val encode_%s : %s -> Js.Json.t Js.Dict.t" type_name type_name;
     F.linep sc
       ("(** [encode_%s v dict] encodes [v] int the " ^^ "given JSON [dict] *)")
       type_name
@@ -292,15 +285,22 @@ let gen_sig ?and_ t sc =
     f v.Ot.v_name;
     true
   | Ot.Const_variant { Ot.cv_name; _ } ->
-    F.linep sc "val encode_%s : %s_types.%s -> string" cv_name module_prefix
-      cv_name;
+    F.linep sc "val encode_%s : %s -> string" cv_name cv_name;
     F.linep sc "(** [encode_%s v] returns JSON string *)" cv_name;
     true
   | Ot.Unit { Ot.er_name } ->
-    F.linep sc "val encode_%s : %s_types.%s -> Js.Json.t" er_name module_prefix
-      er_name;
+    F.linep sc "val encode_%s : %s -> Js.Json.t" er_name er_name;
     F.linep sc "(** [encode_%s v] returns JSON null *)" er_name;
     true
 
 let ocamldoc_title = "Protobuf JSON Encoding"
-let file_suffix = Pb_codegen_decode_bs.file_suffix
+let requires_mutable_records = false
+
+let plugin : Pb_codegen_plugin.t =
+  let module P = struct
+    let gen_sig = gen_sig
+    let gen_struct = gen_struct
+    let ocamldoc_title = ocamldoc_title
+    let requires_mutable_records = requires_mutable_records
+  end in
+  (module P)
diff --git a/src/compilerlib/pb_codegen_encode_bs.mli b/src/compilerlib/pb_codegen_encode_bs.mli
new file mode 100644
index 0000000..3aa642c
--- /dev/null
+++ b/src/compilerlib/pb_codegen_encode_bs.mli
@@ -0,0 +1,3 @@
+include Pb_codegen_plugin.S
+
+val plugin : Pb_codegen_plugin.t
diff --git a/src/compilerlib/pb_codegen_encode_yojson.ml b/src/compilerlib/pb_codegen_encode_yojson.ml
index 3718820..c3e672c 100644
--- a/src/compilerlib/pb_codegen_encode_yojson.ml
+++ b/src/compilerlib/pb_codegen_encode_yojson.ml
@@ -2,7 +2,6 @@ module Ot = Pb_codegen_ocaml_type
 module F = Pb_codegen_formatting
 
 let sp = Pb_codegen_util.sp
-let file_suffix = Pb_codegen_decode_yojson.file_suffix
 
 let unsupported json_label =
   failwith (sp "Unsupported field type for field: %s" json_label)
@@ -57,10 +56,8 @@ let gen_field var_name json_label field_type pk =
   (* User defined *)
   | Ot.Ft_user_defined_type udt, _ ->
     let f_name =
-      let function_prefix = "encode" in
-      let module_suffix = file_suffix in
-      Pb_codegen_util.function_name_of_user_defined ~function_prefix
-        ~module_suffix udt
+      let function_prefix = "encode_json" in
+      Pb_codegen_util.function_name_of_user_defined ~function_prefix udt
     in
     Some (sp "(\"%s\", %s %s)" json_label f_name var_name)
   | _ -> assert false
@@ -74,7 +71,7 @@ let gen_rft_nolabel sc rf_label (field_type, _, pk) =
 
 let gen_rft_optional sc rf_label (field_type, _, pk, _) =
   F.linep sc "let assoc = match v.%s with" rf_label;
-  F.scope sc (fun sc ->
+  F.sub_scope sc (fun sc ->
       F.line sc "| None -> assoc";
       let json_label = Pb_codegen_util.camel_case_of_label rf_label in
       match gen_field "v" json_label field_type pk with
@@ -94,7 +91,7 @@ let gen_rft_repeated sc rf_label repeated_field =
   let json_label = Pb_codegen_util.camel_case_of_label rf_label in
 
   F.line sc "let assoc =";
-  F.scope sc (fun sc ->
+  F.sub_scope sc (fun sc ->
       (match field_type, pk with
       | Ot.Ft_unit, _ -> unsupported json_label
       | Ot.Ft_basic_type basic_type, _ ->
@@ -112,10 +109,8 @@ let gen_rft_repeated sc rf_label repeated_field =
       (* User defined *)
       | Ot.Ft_user_defined_type udt, _ ->
         let f_name =
-          let function_prefix = "encode" in
-          let module_suffix = file_suffix in
-          Pb_codegen_util.function_name_of_user_defined ~function_prefix
-            ~module_suffix udt
+          let function_prefix = "encode_json" in
+          Pb_codegen_util.function_name_of_user_defined ~function_prefix udt
         in
         F.linep sc "let l = %s |> List.map %s in" var_name f_name
       | _ -> unsupported json_label);
@@ -125,7 +120,7 @@ let gen_rft_repeated sc rf_label repeated_field =
 let gen_rft_variant sc rf_label { Ot.v_constructors; _ } =
   F.linep sc "let assoc = match v.%s with" rf_label;
 
-  F.scope sc (fun sc ->
+  F.sub_scope sc (fun sc ->
       List.iter
         (fun { Ot.vc_constructor; vc_field_type; vc_payload_kind; _ } ->
           let var_name = "v" in
@@ -133,7 +128,7 @@ let gen_rft_variant sc rf_label { Ot.v_constructors; _ } =
             Pb_codegen_util.camel_case_of_constructor vc_constructor
           in
           F.linep sc "| %s v ->" vc_constructor;
-          F.scope sc (fun sc ->
+          F.sub_scope sc (fun sc ->
               match vc_field_type with
               | Ot.Vct_nullary ->
                 F.linep sc "(\"%s\", `Null) :: assoc" json_label
@@ -147,13 +142,12 @@ let gen_rft_variant sc rf_label { Ot.v_constructors; _ } =
 
   F.linep sc "in (* match v.%s *)" rf_label
 
-let gen_record ?and_ module_prefix { Ot.r_name; r_fields } sc =
+let gen_record ?and_ { Ot.r_name; r_fields } sc =
   let rn = r_name in
-  F.linep sc "%s encode_%s (v:%s_types.%s) = "
+  F.linep sc "%s encode_json_%s (v:%s) = "
     (Pb_codegen_util.let_decl_of_and and_)
-    rn module_prefix rn;
-  F.scope sc (fun sc ->
-      F.linep sc "let open !%s_types in" module_prefix;
+    rn rn;
+  F.sub_scope sc (fun sc ->
       F.line sc "let assoc = [] in ";
       List.iter
         (fun record_field ->
@@ -177,14 +171,14 @@ let gen_record ?and_ module_prefix { Ot.r_name; r_fields } sc =
         r_fields (* List.iter *);
       F.line sc "`Assoc assoc")
 
-let gen_unit ?and_ module_prefix { Ot.er_name } sc =
+let gen_unit ?and_ { Ot.er_name } sc =
   let rn = er_name in
-  F.linep sc "%s encode_%s (v:%s_types.%s) = "
+  F.linep sc "%s encode_json_%s (v:%s) = "
     (Pb_codegen_util.let_decl_of_and and_)
-    rn module_prefix rn;
+    rn rn;
   F.line sc (sp "Pbrt_yojson.%s %s" "make_unit" "v")
 
-let gen_variant ?and_ module_prefix { Ot.v_name; v_constructors } sc =
+let gen_variant ?and_ { Ot.v_name; v_constructors } sc =
   let process_v_constructor sc v_constructor =
     let { Ot.vc_constructor; Ot.vc_field_type; Ot.vc_payload_kind; _ } =
       v_constructor
@@ -198,50 +192,45 @@ let gen_variant ?and_ module_prefix { Ot.v_name; v_constructors } sc =
     | Ot.Vct_non_nullary_constructor field_type ->
       (match gen_field "v" json_label field_type vc_payload_kind with
       | None ->
-        F.linep sc "| %s_types.%s -> `Assoc [(\"%s\", `Null)]" module_prefix
-          vc_constructor json_label
-      | Some exp ->
-        F.linep sc "| %s_types.%s v -> `Assoc [%s]" module_prefix vc_constructor
-          exp)
+        F.linep sc "| %s -> `Assoc [(\"%s\", `Null)]" vc_constructor json_label
+      | Some exp -> F.linep sc "| %s v -> `Assoc [%s]" vc_constructor exp)
   in
 
-  F.linep sc "%s encode_%s (v:%s_types.%s) = "
+  F.linep sc "%s encode_json_%s (v:%s) = "
     (Pb_codegen_util.let_decl_of_and and_)
-    v_name module_prefix v_name;
-  F.scope sc (fun sc ->
+    v_name v_name;
+  F.sub_scope sc (fun sc ->
       F.line sc "begin match v with";
       List.iter (process_v_constructor sc) v_constructors;
       F.line sc "end")
 
-let gen_const_variant ?and_ module_prefix { Ot.cv_name; Ot.cv_constructors } sc
-    =
-  F.linep sc "%s encode_%s (v:%s_types.%s) = "
+let gen_const_variant ?and_ { Ot.cv_name; Ot.cv_constructors } sc =
+  F.linep sc "%s encode_json_%s (v:%s) = "
     (Pb_codegen_util.let_decl_of_and and_)
-    cv_name module_prefix cv_name;
-  F.scope sc (fun sc ->
+    cv_name cv_name;
+  F.sub_scope sc (fun sc ->
       F.line sc "match v with";
       List.iter
         (fun { Ot.cvc_name; cvc_string_value; _ } ->
-          F.linep sc "| %s_types.%s -> `String \"%s\"" module_prefix cvc_name
-            cvc_string_value)
+          F.linep sc "| %s -> `String \"%s\"" cvc_name cvc_string_value)
         cv_constructors)
 
 let gen_struct ?and_ t sc =
-  let { Ot.spec; module_prefix; _ } = t in
+  let { Ot.spec; _ } = t in
 
   let has_encoded =
     match spec with
     | Ot.Record r ->
-      gen_record ?and_ module_prefix r sc;
+      gen_record ?and_ r sc;
       true
     | Ot.Variant v ->
-      gen_variant ?and_ module_prefix v sc;
+      gen_variant ?and_ v sc;
       true
     | Ot.Const_variant v ->
-      gen_const_variant ?and_ module_prefix v sc;
+      gen_const_variant ?and_ v sc;
       true
     | Ot.Unit u ->
-      gen_unit ?and_ module_prefix u sc;
+      gen_unit ?and_ u sc;
       true
   in
 
@@ -249,12 +238,10 @@ let gen_struct ?and_ t sc =
 
 let gen_sig ?and_ t sc =
   let _ = and_ in
-  let { Ot.module_prefix; _ } = t in
   let f type_name =
-    F.linep sc "val encode_%s : %s_types.%s -> Yojson.Basic.json" type_name
-      module_prefix type_name;
+    F.linep sc "val encode_json_%s : %s -> Yojson.Basic.t" type_name type_name;
     F.linep sc
-      ("(** [encode_%s v encoder] encodes [v] to " ^^ "to json*)")
+      ("(** [encode_json_%s v encoder] encodes [v] to " ^^ "to json *)")
       type_name
   in
   match t with
@@ -272,3 +259,13 @@ let gen_sig ?and_ t sc =
     true
 
 let ocamldoc_title = "Protobuf YoJson Encoding"
+let requires_mutable_records = false
+
+let plugin : Pb_codegen_plugin.t =
+  let module P = struct
+    let gen_sig = gen_sig
+    let gen_struct = gen_struct
+    let ocamldoc_title = ocamldoc_title
+    let requires_mutable_records = requires_mutable_records
+  end in
+  (module P)
diff --git a/src/compilerlib/pb_codegen_encode_yojson.mli b/src/compilerlib/pb_codegen_encode_yojson.mli
index 9ef9b24..1df8a2e 100644
--- a/src/compilerlib/pb_codegen_encode_yojson.mli
+++ b/src/compilerlib/pb_codegen_encode_yojson.mli
@@ -1,3 +1,5 @@
 (** Code generator for the [encode] JSON function *)
 
-include Pb_codegen_sig.S
+include Pb_codegen_plugin.S
+
+val plugin : Pb_codegen_plugin.t
diff --git a/src/compilerlib/pb_codegen_formatting.ml b/src/compilerlib/pb_codegen_formatting.ml
index c575cef..2be5a23 100644
--- a/src/compilerlib/pb_codegen_formatting.ml
+++ b/src/compilerlib/pb_codegen_formatting.ml
@@ -1,8 +1,8 @@
 (*
   The MIT License (MIT)
-  
+
   Copyright (c) 2016 Maxime Ransan <maxime.ransan@gmail.com>
-  
+
   Permission is hereby granted, free of charge, to any person obtaining a copy
   of this software and associated documentation files (the "Software"), to deal
   in the Software without restriction, including without limitation the rights
@@ -34,12 +34,12 @@ let line scope s = scope.items <- Line s :: scope.items
 let linep scope format = Printf.ksprintf (line scope) format
 let empty_line scope = line scope ""
 
-let scope scope f =
+let sub_scope scope f : unit =
   let sub_scope = empty_scope () in
   f sub_scope;
   scope.items <- Scope sub_scope :: scope.items
 
-let print scope =
+let to_string (scope : scope) : string =
   let rec loop acc i = function
     | Line s :: tl -> loop ((Pb_util.indentation_prefix i ^ s) :: acc) i tl
     | Scope { items } :: tl ->
@@ -49,3 +49,5 @@ let print scope =
   in
 
   String.concat "\n" @@ loop [] 0 scope.items
+
+let output oc (scope : scope) : unit = output_string oc @@ to_string scope
diff --git a/src/compilerlib/pb_codegen_formatting.mli b/src/compilerlib/pb_codegen_formatting.mli
index c5d49e0..9b6253c 100644
--- a/src/compilerlib/pb_codegen_formatting.mli
+++ b/src/compilerlib/pb_codegen_formatting.mli
@@ -1,8 +1,8 @@
 (*
   The MIT License (MIT)
-  
+
   Copyright (c) 2016 Maxime Ransan <maxime.ransan@gmail.com>
-  
+
   Permission is hereby granted, free of charge, to any person obtaining a copy
   of this software and associated documentation files (the "Software"), to deal
   in the Software without restriction, including without limitation the rights
@@ -25,13 +25,10 @@
 
 (** Formatting utilities for code generation *)
 
-(* TODO: add a little example of how to use this module as well
- * as a unit test for it *)
-
 (** {2 types} *)
 
 type scope
-(** A scope is a formatting container which can contains line of text as well 
+(** A scope is a formatting container which can contains line of text as well
     as other nested scopes
 
     In other word a scope define an indentation section.
@@ -50,12 +47,15 @@ val linep : scope -> ('a, unit, string, unit) format4 -> 'a
 val empty_line : scope -> unit
 (** [empty_line scope] adds an empty line to [scope] *)
 
-val scope : scope -> (scope -> unit) -> unit
-(** [scope scope f] adds a sub scope and apply [f] to it. *)
+val sub_scope : scope -> (scope -> unit) -> unit
+(** [sub_scope scope f] adds a sub scope and apply [f] to it. *)
 
 (** {2 Printing} *)
 
-val print : scope -> string
-(** [print scope] returns the formatted scops with a 2 character 
-    indentation for each scope. 
+val to_string : scope -> string
+(** [print scope] returns the formatted scops with a 2 character
+    indentation for each scope.
  *)
+
+val output : out_channel -> scope -> unit
+(** Write to output *)
diff --git a/src/compilerlib/pb_codegen_ocaml_type.ml b/src/compilerlib/pb_codegen_ocaml_type.ml
index 41dc55c..de7edec 100644
--- a/src/compilerlib/pb_codegen_ocaml_type.ml
+++ b/src/compilerlib/pb_codegen_ocaml_type.ml
@@ -34,15 +34,15 @@ type payload_kind =
   | Pk_bytes
 
 type user_defined_type = {
-  (* since code generated is split in multiple file (type, binary, json, ..)
+  udt_module_prefix: string option;
+      (** since code generated is split in multiple file (type, binary, json, ..)
      this defines the prefix for the given type, the suffix will
      be defined by each generator *)
-  udt_module_prefix: string option;
-  (* OCaml type name ie not the type name in proto file *)
   udt_type_name: string;
-  (* Need to keep track of this since encoding logic in binary
-     format is quite different *)
+      (** OCaml type name ie not the type name in proto file *)
   udt_type: [ `Message | `Enum ];
+      (** Need to keep track of this since encoding logic in binary
+     format is quite different *)
 }
 
 type basic_type =
diff --git a/src/compilerlib/pb_codegen_plugin.ml b/src/compilerlib/pb_codegen_plugin.ml
new file mode 100644
index 0000000..08a6050
--- /dev/null
+++ b/src/compilerlib/pb_codegen_plugin.ml
@@ -0,0 +1,31 @@
+(** Plugin architecture.
+
+    OCaml-protoc generates code based on a number of plugins,
+    each of which can contribute code to the output files (.ml and .mli).
+*)
+
+type codegen_f =
+  ?and_:unit ->
+  Pb_codegen_ocaml_type.type_ ->
+  Pb_codegen_formatting.scope ->
+  bool
+(** A code generation function *)
+
+module type S = sig
+  val gen_sig : codegen_f
+  (** Generate a signature file (.mli) *)
+
+  val gen_struct : codegen_f
+  (** Generate the implementation (.ml) *)
+
+  val ocamldoc_title : string
+  (** OCamldoc title *)
+
+  val requires_mutable_records : bool
+  (** Does this record depend on mutable records being defined? *)
+end
+
+type t = (module S)
+(** A plugin is a code-generator respecting the signature {!S}. *)
+
+let requires_mutable_records (module P : S) = P.requires_mutable_records
diff --git a/src/compilerlib/pb_codegen_pp.ml b/src/compilerlib/pb_codegen_pp.ml
index fb1fd51..fdbefc4 100644
--- a/src/compilerlib/pb_codegen_pp.ml
+++ b/src/compilerlib/pb_codegen_pp.ml
@@ -1,36 +1,30 @@
 module Ot = Pb_codegen_ocaml_type
 module F = Pb_codegen_formatting
 module L = Pb_logger
-
-let file_suffix = "pp"
-
 open Pb_codegen_util
 
 let gen_field field_type =
   match field_type with
   | Ot.Ft_user_defined_type udt ->
     let function_prefix = "pp" in
-    let module_suffix = file_suffix in
-    function_name_of_user_defined ~function_prefix ~module_suffix udt
+    function_name_of_user_defined ~function_prefix udt
   | Ot.Ft_wrapper_type { Ot.wt_type; _ } ->
     sp "Pbrt.Pp.pp_wrapper_%s" (string_of_basic_type ~for_pp:true wt_type)
   | _ -> sp "Pbrt.Pp.pp_%s" (string_of_field_type ~for_pp:true field_type)
 
-let gen_record ?and_ module_prefix { Ot.r_name; r_fields } sc =
+let gen_record ?and_ { Ot.r_name; r_fields } sc =
   L.log "gen_pp, record_name: %s\n" r_name;
 
-  F.line sc
-  @@ sp "%s pp_%s fmt (v:%s_types.%s) = " (let_decl_of_and and_) r_name
-       module_prefix r_name;
-  F.scope sc (fun sc ->
+  F.line sc @@ sp "%s pp_%s fmt (v:%s) = " (let_decl_of_and and_) r_name r_name;
+  F.sub_scope sc (fun sc ->
       F.line sc "let pp_i fmt () =";
-      F.scope sc (fun sc ->
+      F.sub_scope sc (fun sc ->
           List.iteri
             (fun i record_field ->
               let first = i = 0 in
               let { Ot.rf_label; rf_field_type; _ } = record_field in
 
-              let var_name = sp "v.%s_types.%s" module_prefix rf_label in
+              let var_name = sp "v.%s" rf_label in
               match rf_field_type with
               | Ot.Rft_nolabel (field_type, _, _)
               | Ot.Rft_required (field_type, _, _, _) ->
@@ -93,67 +87,57 @@ let gen_record ?and_ module_prefix { Ot.r_name; r_fields } sc =
       F.line sc "in";
       F.line sc "Pbrt.Pp.pp_brk pp_i fmt ()")
 
-let gen_unit ?and_ module_prefix { Ot.er_name } sc =
+let gen_unit ?and_ { Ot.er_name } sc : unit =
   F.line sc
-  @@ sp "%s pp_%s fmt (v:%s_types.%s) = " (let_decl_of_and and_) er_name
-       module_prefix er_name;
-  F.scope sc (fun sc ->
+  @@ sp "%s pp_%s fmt (v:%s) = " (let_decl_of_and and_) er_name er_name;
+  F.sub_scope sc (fun sc ->
       F.line sc "let pp_i fmt () =";
-      F.scope sc (fun sc -> F.line sc "Pbrt.Pp.pp_unit fmt ()");
+      F.sub_scope sc (fun sc -> F.line sc "Pbrt.Pp.pp_unit fmt ()");
       F.line sc "in";
       F.line sc "Pbrt.Pp.pp_brk pp_i fmt ()")
 
-let gen_variant ?and_ module_prefix { Ot.v_name; Ot.v_constructors } sc =
-  F.line sc
-  @@ sp "%s pp_%s fmt (v:%s_types.%s) =" (let_decl_of_and and_) v_name
-       module_prefix v_name;
-  F.scope sc (fun sc ->
+let gen_variant ?and_ { Ot.v_name; Ot.v_constructors } sc =
+  F.line sc @@ sp "%s pp_%s fmt (v:%s) =" (let_decl_of_and and_) v_name v_name;
+  F.sub_scope sc (fun sc ->
       F.line sc "match v with";
       List.iter
         (fun { Ot.vc_constructor; vc_field_type; _ } ->
           match vc_field_type with
           | Ot.Vct_nullary ->
             F.line sc
-            @@ sp "| %s_types.%s  -> Format.fprintf fmt \"%s\"" module_prefix
-                 vc_constructor vc_constructor
+            @@ sp "| %s  -> Format.fprintf fmt \"%s\"" vc_constructor
+                 vc_constructor
           | Ot.Vct_non_nullary_constructor field_type ->
             let field_string_of = gen_field field_type in
             F.line sc
-            @@ sp
-                 "| %s_types.%s x -> Format.fprintf fmt \"@[<hv2>%s(@,\
-                  %%a)@]\" %s x" module_prefix vc_constructor vc_constructor
-                 field_string_of)
+            @@ sp "| %s x -> Format.fprintf fmt \"@[<hv2>%s(@,%%a)@]\" %s x"
+                 vc_constructor vc_constructor field_string_of)
         v_constructors)
 
-let gen_const_variant ?and_ module_prefix { Ot.cv_name; cv_constructors } sc =
-  F.line sc
-  @@ sp "%s pp_%s fmt (v:%s_types.%s) =" (let_decl_of_and and_) cv_name
-       module_prefix cv_name;
-  F.scope sc (fun sc ->
+let gen_const_variant ?and_ { Ot.cv_name; cv_constructors } sc =
+  F.line sc @@ sp "%s pp_%s fmt (v:%s) =" (let_decl_of_and and_) cv_name cv_name;
+  F.sub_scope sc (fun sc ->
       F.line sc "match v with";
       List.iter
         (fun { Ot.cvc_name; _ } ->
-          F.line sc
-          @@ sp "| %s_types.%s -> Format.fprintf fmt \"%s\"" module_prefix
-               cvc_name cvc_name)
+          F.line sc @@ sp "| %s -> Format.fprintf fmt \"%s\"" cvc_name cvc_name)
         cv_constructors)
 
 let gen_struct ?and_ t sc =
-  let { Ot.module_prefix; spec; _ } = t in
+  let { Ot.spec; _ } = t in
   (match spec with
-  | Ot.Record r -> gen_record ?and_ module_prefix r sc
-  | Ot.Variant v -> gen_variant ?and_ module_prefix v sc
-  | Ot.Const_variant v -> gen_const_variant ?and_ module_prefix v sc
-  | Ot.Unit u -> gen_unit ?and_ module_prefix u sc);
+  | Ot.Record r -> gen_record ?and_ r sc
+  | Ot.Variant v -> gen_variant ?and_ v sc
+  | Ot.Const_variant v -> gen_const_variant ?and_ v sc
+  | Ot.Unit u -> gen_unit ?and_ u sc);
   true
 
 let gen_sig ?and_ t sc =
   let _ = and_ in
-  let { Ot.module_prefix; spec; _ } = t in
+  let { Ot.spec; _ } = t in
   let f type_name =
     F.line sc
-    @@ sp "val pp_%s : Format.formatter -> %s_types.%s -> unit " type_name
-         module_prefix type_name;
+    @@ sp "val pp_%s : Format.formatter -> %s -> unit " type_name type_name;
     F.line sc @@ sp "(** [pp_%s v] formats v *)" type_name
   in
   (match spec with
@@ -164,3 +148,13 @@ let gen_sig ?and_ t sc =
   true
 
 let ocamldoc_title = "Formatters"
+let requires_mutable_records = false
+
+let plugin : Pb_codegen_plugin.t =
+  let module P = struct
+    let gen_sig = gen_sig
+    let gen_struct = gen_struct
+    let ocamldoc_title = ocamldoc_title
+    let requires_mutable_records = requires_mutable_records
+  end in
+  (module P)
diff --git a/src/compilerlib/pb_codegen_pp.mli b/src/compilerlib/pb_codegen_pp.mli
index 8a2e281..f144559 100644
--- a/src/compilerlib/pb_codegen_pp.mli
+++ b/src/compilerlib/pb_codegen_pp.mli
@@ -1,3 +1,5 @@
 (** Code generator for the [pp] function *)
 
-include Pb_codegen_sig.S
+include Pb_codegen_plugin.S
+
+val plugin : Pb_codegen_plugin.t
diff --git a/src/compilerlib/pb_codegen_sig.ml b/src/compilerlib/pb_codegen_sig.ml
deleted file mode 100644
index d140422..0000000
--- a/src/compilerlib/pb_codegen_sig.ml
+++ /dev/null
@@ -1,19 +0,0 @@
-module type S = sig
-  val gen_sig :
-    ?and_:unit ->
-    Pb_codegen_ocaml_type.type_ ->
-    Pb_codegen_formatting.scope ->
-    bool
-  (** Generate a signature file (.mli) *)
-
-  val gen_struct :
-    ?and_:unit ->
-    Pb_codegen_ocaml_type.type_ ->
-    Pb_codegen_formatting.scope ->
-    bool
-  (** Generate the implementation (.ml) *)
-
-  val ocamldoc_title : string
-  val file_suffix : string
-end
-(* S *)
diff --git a/src/compilerlib/pb_codegen_sig.mli b/src/compilerlib/pb_codegen_sig.mli
deleted file mode 100644
index 84a673f..0000000
--- a/src/compilerlib/pb_codegen_sig.mli
+++ /dev/null
@@ -1,27 +0,0 @@
-(** Common module signature for all codegen modules *)
-
-(*type context = [| `Single_file | `Multi_file |] *)
-
-module type S = sig
-  val gen_sig :
-    ?and_:unit ->
-    Pb_codegen_ocaml_type.type_ ->
-    Pb_codegen_formatting.scope ->
-    bool
-  (** [gen_sig t] generates the code for the module signature *)
-
-  val gen_struct :
-    ?and_:unit ->
-    Pb_codegen_ocaml_type.type_ ->
-    Pb_codegen_formatting.scope ->
-    bool
-  (** [gen_struct t] generates the code for the module structure *)
-
-  val ocamldoc_title : string
-  (** OCamldoc title *)
-
-  val file_suffix : string
-  (** The suffix part of the generated file which will contain the 
-      struct and sig *)
-end
-(* S *)
diff --git a/src/compilerlib/pb_codegen_types.ml b/src/compilerlib/pb_codegen_types.ml
index c652158..cf94d19 100644
--- a/src/compilerlib/pb_codegen_types.ml
+++ b/src/compilerlib/pb_codegen_types.ml
@@ -16,7 +16,7 @@ let is_imperative_type = function
   | Ot.Rft_associative (Ot.At_hashtable, _, _, _) ->
     true
 
-let gen_record_mutable module_prefix { Ot.r_name; r_fields } sc =
+let gen_record_mutable { Ot.r_name; r_fields } sc : unit =
   let field_prefix field_type =
     if is_imperative_type field_type then
       ""
@@ -27,13 +27,12 @@ let gen_record_mutable module_prefix { Ot.r_name; r_fields } sc =
   let r_name = Pb_codegen_util.mutable_record_name r_name in
 
   F.linep sc "type %s = {" r_name;
-  F.scope sc (fun sc ->
+  F.sub_scope sc (fun sc ->
       List.iter
         (fun { Ot.rf_label; rf_field_type; _ } ->
           let prefix = field_prefix rf_field_type in
           let type_ =
-            Pb_codegen_util.string_of_record_field_type ~module_prefix
-              rf_field_type
+            Pb_codegen_util.string_of_record_field_type rf_field_type
           in
           F.linep sc "%s%s : %s;" prefix rf_label type_)
         r_fields);
@@ -48,7 +47,7 @@ let gen_record ?and_ { Ot.r_name; r_fields } sc =
   in
 
   F.linep sc "%s %s = {" (type_decl_of_and and_) r_name;
-  F.scope sc (fun sc ->
+  F.sub_scope sc (fun sc ->
       List.iter
         (fun { Ot.rf_label; rf_field_type; rf_mutable } ->
           let prefix = field_prefix rf_mutable in
@@ -64,7 +63,7 @@ let gen_variant ?and_ variant sc =
 
   F.linep sc "%s %s =" (type_decl_of_and and_) v_name;
 
-  F.scope sc (fun sc ->
+  F.sub_scope sc (fun sc ->
       List.iter
         (fun { Ot.vc_constructor; vc_field_type; _ } ->
           match vc_field_type with
@@ -76,7 +75,7 @@ let gen_variant ?and_ variant sc =
 
 let gen_const_variant ?and_ { Ot.cv_name; cv_constructors } sc =
   F.linep sc "%s %s =" (type_decl_of_and and_) cv_name;
-  F.scope sc (fun sc ->
+  F.sub_scope sc (fun sc ->
       List.iter
         (fun { Ot.cvc_name; _ } -> F.linep sc "| %s " cvc_name)
         cv_constructors)
@@ -89,7 +88,7 @@ let print_ppx_extension { Ot.type_level_ppx_extension; _ } sc =
   | None -> ()
   | Some ppx_content -> F.linep sc "[@@%s]" ppx_content
 
-let gen_struct ?and_ t scope =
+let gen_struct_full ~with_mutable_records ?and_ t scope =
   let { Ot.spec; _ } = t in
   (match spec with
   | Ot.Record r -> gen_record ?and_ r scope
@@ -97,8 +96,16 @@ let gen_struct ?and_ t scope =
   | Ot.Const_variant v -> gen_const_variant ?and_ v scope
   | Ot.Unit v -> gen_unit ?and_ v scope);
   print_ppx_extension t scope;
+
+  (match spec with
+  | Ot.Record r when with_mutable_records -> gen_record_mutable r scope
+  | _ -> ());
+
   true
 
+let gen_struct ?and_ t sc =
+  gen_struct_full ?and_ ~with_mutable_records:false t sc
+
 let gen_sig ?and_ t scope =
   let { Ot.spec; _ } = t in
   (match spec with
@@ -110,4 +117,4 @@ let gen_sig ?and_ t scope =
   true
 
 let ocamldoc_title = "Types"
-let file_suffix = "types"
+let requires_mutable_records = false
diff --git a/src/compilerlib/pb_codegen_types.mli b/src/compilerlib/pb_codegen_types.mli
index a1fbf60..f191201 100644
--- a/src/compilerlib/pb_codegen_types.mli
+++ b/src/compilerlib/pb_codegen_types.mli
@@ -1,10 +1,6 @@
 (** Code generator for the OCaml type *)
 
-include Pb_codegen_sig.S
+include Pb_codegen_plugin.S
 
 val gen_record_mutable :
-  string ->
-  (* module_prefix *)
-  Pb_codegen_ocaml_type.record ->
-  Pb_codegen_formatting.scope ->
-  unit
+  Pb_codegen_ocaml_type.record -> Pb_codegen_formatting.scope -> unit
diff --git a/src/compilerlib/pb_codegen_util.ml b/src/compilerlib/pb_codegen_util.ml
index 964e53b..d11a318 100644
--- a/src/compilerlib/pb_codegen_util.ml
+++ b/src/compilerlib/pb_codegen_util.ml
@@ -29,9 +29,9 @@ let string_of_user_defined ?module_prefix = function
   | { Ot.udt_module_prefix = None; Ot.udt_type_name; _ } ->
     (match module_prefix with
     | None -> udt_type_name
-    | Some module_prefix -> module_prefix ^ "_types." ^ udt_type_name)
+    | Some module_prefix -> module_prefix ^ "." ^ udt_type_name)
   | { Ot.udt_module_prefix = Some module_prefix; Ot.udt_type_name; _ } ->
-    module_prefix ^ "_types." ^ udt_type_name
+    module_prefix ^ "." ^ udt_type_name
 
 let string_of_field_type ?for_pp ?module_prefix = function
   | Ot.Ft_unit -> "unit"
@@ -69,7 +69,7 @@ let string_of_record_field_type ?module_prefix = function
   | Ot.Rft_variant { Ot.v_name; _ } ->
     (match module_prefix with
     | None -> v_name
-    | Some module_prefix -> module_prefix ^ "_types." ^ v_name)
+    | Some module_prefix -> module_prefix ^ "." ^ v_name)
 
 (** [function_name_of_user_defined prefix user_defined] returns the function
     name of the form `(module'.'?)prefix_(type_name)`.
@@ -78,18 +78,18 @@ let string_of_record_field_type ?module_prefix = function
     (encode/decode/to_string) will call the same generated function for each
     user defined field type.
  *)
-let function_name_of_user_defined ~function_prefix ~module_suffix = function
+let function_name_of_user_defined ~function_prefix = function
   | { Ot.udt_module_prefix = Some module_prefix; Ot.udt_type_name; _ } ->
-    sp "%s_%s.%s_%s" module_prefix module_suffix function_prefix udt_type_name
+    sp "%s.%s_%s" module_prefix function_prefix udt_type_name
   | { Ot.udt_module_prefix = None; Ot.udt_type_name; _ } ->
     sp "%s_%s" function_prefix udt_type_name
 
-let caml_file_name_of_proto_file_name ~proto_file_name ~file_suffix =
+let caml_file_name_of_proto_file_name ~proto_file_name =
   let splitted = Pb_util.rev_split_by_char '.' proto_file_name in
   if List.length splitted < 2 || List.hd splitted <> "proto" then
     failwith "Proto file has no valid extension"
   else
-    String.concat "_" @@ List.rev @@ (file_suffix :: List.tl splitted)
+    String.concat "_" @@ List.rev @@ List.tl splitted
 
 let mutable_record_name s = s ^ "_mutable"
 
diff --git a/src/compilerlib/pb_codegen_util.mli b/src/compilerlib/pb_codegen_util.mli
index 2e2b00c..fa52b10 100644
--- a/src/compilerlib/pb_codegen_util.mli
+++ b/src/compilerlib/pb_codegen_util.mli
@@ -21,10 +21,7 @@ val string_of_field_type :
   string
 
 val function_name_of_user_defined :
-  function_prefix:string ->
-  module_suffix:string ->
-  Pb_codegen_ocaml_type.user_defined_type ->
-  string
+  function_prefix:string -> Pb_codegen_ocaml_type.user_defined_type -> string
 (** [function_name_of_user_defined prefix user_defined] returns the function
     name of the form `(module'.'?)<prefix>_(type_name)`.
 
@@ -33,8 +30,7 @@ val function_name_of_user_defined :
     user defined field type.
  *)
 
-val caml_file_name_of_proto_file_name :
-  proto_file_name:string -> file_suffix:string -> string
+val caml_file_name_of_proto_file_name : proto_file_name:string -> string
 (** [caml_file_name_of_proto_file_name filename] returns the OCaml file name from
     the protobuf file name
  *)
diff --git a/src/compilerlib/pb_exception.ml b/src/compilerlib/pb_exception.ml
index 2070607..8cf2de4 100644
--- a/src/compilerlib/pb_exception.ml
+++ b/src/compilerlib/pb_exception.ml
@@ -1,8 +1,8 @@
 (*
   The MIT License (MIT)
-  
+
   Copyright (c) 2016 Maxime Ransan <maxime.ransan@gmail.com>
-  
+
   Permission is hereby granted, free of charge, to any person obtaining a copy
   of this software and associated documentation files (the "Software"), to deal
   in the Software without restriction, including without limitation the rights
diff --git a/src/compilerlib/pb_exception.mli b/src/compilerlib/pb_exception.mli
index 64a8b9a..1faa82c 100644
--- a/src/compilerlib/pb_exception.mli
+++ b/src/compilerlib/pb_exception.mli
@@ -1,8 +1,8 @@
 (*
   The MIT License (MIT)
-  
+
   Copyright (c) 2016 Maxime Ransan <maxime.ransan@gmail.com>
-  
+
   Permission is hereby granted, free of charge, to any person obtaining a copy
   of this software and associated documentation files (the "Software"), to deal
   in the Software without restriction, including without limitation the rights
diff --git a/src/compilerlib/pb_field_type.ml b/src/compilerlib/pb_field_type.ml
index 1ff26fe..c8ba32f 100644
--- a/src/compilerlib/pb_field_type.ml
+++ b/src/compilerlib/pb_field_type.ml
@@ -1,8 +1,8 @@
 (*
   The MIT License (MIT)
-  
+
   Copyright (c) 2016 Maxime Ransan <maxime.ransan@gmail.com>
-  
+
   Permission is hereby granted, free of charge, to any person obtaining a copy
   of this software and associated documentation files (the "Software"), to deal
   in the Software without restriction, including without limitation the rights
diff --git a/src/compilerlib/pb_field_type.mli b/src/compilerlib/pb_field_type.mli
index b5176ba..1ee31bb 100644
--- a/src/compilerlib/pb_field_type.mli
+++ b/src/compilerlib/pb_field_type.mli
@@ -1,8 +1,8 @@
 (*
   The MIT License (MIT)
-  
+
   Copyright (c) 2016 Maxime Ransan <maxime.ransan@gmail.com>
-  
+
   Permission is hereby granted, free of charge, to any person obtaining a copy
   of this software and associated documentation files (the "Software"), to deal
   in the Software without restriction, including without limitation the rights
@@ -27,12 +27,11 @@
 
 type type_path = string list
 (** Scope path of a type used for a message field.
-    
-    For instance in the following field defintion:
 
+    For instance in the following field definition:
     [required foo.bar.Msg1 f = 1]
 
-    The [type_path] would be [\["foo"; "bar"\]]
+    The [type_path] would be [["foo"; "bar"]]
   *)
 
 type unresolved = {
@@ -40,29 +39,27 @@ type unresolved = {
   type_name: string;
   from_root: bool;
       (** from_root indicates that the scope for the type is
-                         from the root of the type system. (ie starts with '.')
-                      *)
+          from the root of the type system. (ie starts with '.') *)
 }
-(** In the first phase of the compilation 
-    the field of message type are not resolved but only 
-    properly parsed. 
-    
+(** In the first phase of the compilation
+    the field of message type are not resolved but only
+    properly parsed.
+
     The following type summarizes the information of a field
-    type. 
+    type.
 
     In the following field definition:
-    
-    [required foo.bar.Msg1 f = 1] 
+    [required foo.bar.Msg1 f = 1]
 
     The unresolved type would be: [{
-      scope=\["foo";"bar"\]; 
-      type_name="Msg1"; 
+      scope=["foo";"bar"];
+      type_name="Msg1";
       from_root = false
     }]
  *)
 
 type resolved = int
-(** After phase 2 compilation the field type is resolved to an 
+(** After phase 2 compilation the field type is resolved to a
     known message which can be uniquely identified by its id.
   *)
 
@@ -115,11 +112,11 @@ type 'a t =
   [ builtin_type
   | `User_defined of 'a (* Message or Enum type *)
   ]
-(** field type. 
-    
-    The ['a] type is for re-using the same type 
-    definition for the 2 compilation phases. 
-    
+(** field type.
+
+    The ['a] type is for re-using the same type
+    definition for the 2 compilation phases.
+
     After Phase 1 ['a] is [unresolved] while after Phase2
     ['a] is [resolved].
   *)
diff --git a/src/compilerlib/pb_location.ml b/src/compilerlib/pb_location.ml
index 8976340..4e1a9b9 100644
--- a/src/compilerlib/pb_location.ml
+++ b/src/compilerlib/pb_location.ml
@@ -1,8 +1,8 @@
 (*
   The MIT License (MIT)
-  
+
   Copyright (c) 2016 Maxime Ransan <maxime.ransan@gmail.com>
-  
+
   Permission is hereby granted, free of charge, to any person obtaining a copy
   of this software and associated documentation files (the "Software"), to deal
   in the Software without restriction, including without limitation the rights
diff --git a/src/compilerlib/pb_location.mli b/src/compilerlib/pb_location.mli
index fd209ca..66dfa06 100644
--- a/src/compilerlib/pb_location.mli
+++ b/src/compilerlib/pb_location.mli
@@ -1,8 +1,8 @@
 (*
   The MIT License (MIT)
-  
+
   Copyright (c) 2016 Maxime Ransan <maxime.ransan@gmail.com>
-  
+
   Permission is hereby granted, free of charge, to any person obtaining a copy
   of this software and associated documentation files (the "Software"), to deal
   in the Software without restriction, including without limitation the rights
diff --git a/src/compilerlib/pb_logger.ml b/src/compilerlib/pb_logger.ml
index 36176e9..bdb08de 100644
--- a/src/compilerlib/pb_logger.ml
+++ b/src/compilerlib/pb_logger.ml
@@ -1,8 +1,8 @@
 (*
   The MIT License (MIT)
-  
+
   Copyright (c) 2016 Maxime Ransan <maxime.ransan@gmail.com>
-  
+
   Permission is hereby granted, free of charge, to any person obtaining a copy
   of this software and associated documentation files (the "Software"), to deal
   in the Software without restriction, including without limitation the rights
diff --git a/src/compilerlib/pb_logger.mli b/src/compilerlib/pb_logger.mli
index 135b881..2baf385 100644
--- a/src/compilerlib/pb_logger.mli
+++ b/src/compilerlib/pb_logger.mli
@@ -1,8 +1,8 @@
 (*
   The MIT License (MIT)
-  
+
   Copyright (c) 2016 Maxime Ransan <maxime.ransan@gmail.com>
-  
+
   Permission is hereby granted, free of charge, to any person obtaining a copy
   of this software and associated documentation files (the "Software"), to deal
   in the Software without restriction, including without limitation the rights
@@ -37,7 +37,7 @@ val setup_from_out_channel : out_channel -> unit
 val log : ('a, out_channel, unit) format -> 'a
 (** [log format x y ] same as [Printf.printf] except that the message will
     get printed to the previously setup [out_channel] in the
-    [setup_from_out_channel] function. 
+    [setup_from_out_channel] function.
  *)
 
 val endline : string -> unit
diff --git a/src/compilerlib/pb_option.mli b/src/compilerlib/pb_option.mli
index 9fcdb0a..f282d6a 100644
--- a/src/compilerlib/pb_option.mli
+++ b/src/compilerlib/pb_option.mli
@@ -25,7 +25,7 @@ type t = option_name * value
 
 type set
 (** Collection of options
- 
+
     Can be used for field/message or file options *)
 
 val empty : set
diff --git a/src/compilerlib/pb_parsing.ml b/src/compilerlib/pb_parsing.ml
index 2882b3c..0d5a0ca 100644
--- a/src/compilerlib/pb_parsing.ml
+++ b/src/compilerlib/pb_parsing.ml
@@ -1,8 +1,8 @@
 (*
   The MIT License (MIT)
-  
+
   Copyright (c) 2016 Maxime Ransan <maxime.ransan@gmail.com>
-  
+
   Permission is hereby granted, free of charge, to any person obtaining a copy
   of this software and associated documentation files (the "Software"), to deal
   in the Software without restriction, including without limitation the rights
diff --git a/src/compilerlib/pb_parsing.mli b/src/compilerlib/pb_parsing.mli
index 087b6c8..aa43847 100644
--- a/src/compilerlib/pb_parsing.mli
+++ b/src/compilerlib/pb_parsing.mli
@@ -1,8 +1,8 @@
 (*
   The MIT License (MIT)
-  
+
   Copyright (c) 2016 Maxime Ransan <maxime.ransan@gmail.com>
-  
+
   Permission is hereby granted, free of charge, to any person obtaining a copy
   of this software and associated documentation files (the "Software"), to deal
   in the Software without restriction, including without limitation the rights
@@ -25,7 +25,7 @@
 
 (** Parsing compilation step *)
 
-(** This module perform the parsing of the given protobuf file along with 
+(** This module perform the parsing of the given protobuf file along with
     all the imported files. *)
 
 type file_loader = string -> string * string
@@ -33,5 +33,5 @@ type file_loader = string -> string * string
     content *)
 
 val parse_file : file_loader -> string -> Pb_parsing_parse_tree.proto list
-(** [parse_file file_loader file_name] reutrns the protobuf parse tree 
+(** [parse_file file_loader file_name] reutrns the protobuf parse tree
     of the given [file_name] and all its dependent files (ie imports). *)
diff --git a/src/compilerlib/pb_parsing_parse_tree.ml b/src/compilerlib/pb_parsing_parse_tree.ml
index 1739ffe..ba9f39e 100644
--- a/src/compilerlib/pb_parsing_parse_tree.ml
+++ b/src/compilerlib/pb_parsing_parse_tree.ml
@@ -35,9 +35,9 @@ type message_field_label =
 *)
 
 type oneof_field_label = unit
-(** Oneof field fields label 
+(** Oneof field fields label
 
-    Oneof fields have no label, they are simply choices for the 
+    Oneof fields have no label, they are simply choices for the
     oneof fiel they belong to. *)
 
 type 'a field = {
@@ -47,11 +47,11 @@ type 'a field = {
   field_type: Pb_field_type.unresolved_t;
   field_options: Pb_option.set;
 }
-(** message field. 
+(** message field.
 
-    Note this field is parametrized with the label type 
-    so that it can be used both by normal field and one of 
-    field since the only difference between the 2 is 
+    Note this field is parametrized with the label type
+    so that it can be used both by normal field and one of
+    field since the only difference between the 2 is
     the label.
 *)
 
@@ -101,8 +101,8 @@ type extension_range =
   | Extension_single_number of int
   | Extension_range of extension_range_from * extension_range_to
 
-(** Body content defines all the possible consituant 
-    of a message. 
+(** Body content defines all the possible consituant
+    of a message.
 *)
 type message_body_content =
   | Message_field of message_field
@@ -119,10 +119,10 @@ and message = {
   message_name: string;
   message_body: message_body_content list;
 }
-(** Message entity. 
+(** Message entity.
 
     Note the ID is simply for uniquely (and easily) identifying a type. It is
-    expected to be generated by a parser. The later compilation 
+    expected to be generated by a parser. The later compilation
     functions expects this id to be unique.
 *)
 
@@ -166,7 +166,7 @@ type proto = {
   enums: enum list;
   extends: extend list;
 }
-(** Definition of a protobuffer message file. 
+(** Definition of a protobuffer message file.
 *)
 
 [@@@warning "-44"]
diff --git a/src/compilerlib/pb_parsing_parser.mly b/src/compilerlib/pb_parsing_parser.mly
index 949d250..4737f19 100644
--- a/src/compilerlib/pb_parsing_parser.mly
+++ b/src/compilerlib/pb_parsing_parser.mly
@@ -93,6 +93,8 @@
 %type <Pb_parsing_parse_tree.extension_range list> extension_
 %start reserved_
 %type <Pb_parsing_parse_tree.extension_range list> reserved_
+%start extend_
+%type <Pb_parsing_parse_tree.extend> extend_
 
 %%
 
@@ -113,6 +115,7 @@ option_          : option        T_eof {$1}
 extension_range_list_ : extension_range_list T_eof {$1}
 extension_       : extension     T_eof {$1}
 reserved_        : reserved      T_eof {$1}
+extend_          : extend T_eof {$1}
 
 /* (* Main protobuf symbol *) */
 
diff --git a/src/compilerlib/pb_parsing_util.mli b/src/compilerlib/pb_parsing_util.mli
index 8222107..2f3bb84 100644
--- a/src/compilerlib/pb_parsing_util.mli
+++ b/src/compilerlib/pb_parsing_util.mli
@@ -104,7 +104,7 @@ val proto :
   ?extend:Pt.extend ->
   unit ->
   Pt.proto
-(** [proto ~message ~proto ()] will add [message] to [proto] and return a 
+(** [proto ~message ~proto ()] will add [message] to [proto] and return a
     the updated proto value. *)
 
 val finalize_proto_value : Pt.proto -> Pt.proto
diff --git a/src/compilerlib/pb_typing.ml b/src/compilerlib/pb_typing.ml
index 3223af9..c572574 100644
--- a/src/compilerlib/pb_typing.ml
+++ b/src/compilerlib/pb_typing.ml
@@ -1,8 +1,8 @@
 (*
   The MIT License (MIT)
-  
+
   Copyright (c) 2016 Maxime Ransan <maxime.ransan@gmail.com>
-  
+
   Permission is hereby granted, free of charge, to any person obtaining a copy
   of this software and associated documentation files (the "Software"), to deal
   in the Software without restriction, including without limitation the rights
diff --git a/src/compilerlib/pb_typing.mli b/src/compilerlib/pb_typing.mli
index 216105d..a84db4a 100644
--- a/src/compilerlib/pb_typing.mli
+++ b/src/compilerlib/pb_typing.mli
@@ -1,8 +1,8 @@
 (*
   The MIT License (MIT)
-  
+
   Copyright (c) 2016 Maxime Ransan <maxime.ransan@gmail.com>
-  
+
   Permission is hereby granted, free of charge, to any person obtaining a copy
   of this software and associated documentation files (the "Software"), to deal
   in the Software without restriction, including without limitation the rights
@@ -25,7 +25,7 @@
 
 (** Typing compilation step *)
 
-(** This module performs the typing step in the compilation process by 
+(** This module performs the typing step in the compilation process by
     doing the following:
     {ul
     {- Type tree construction and validation}
@@ -39,6 +39,6 @@ module Tt = Pb_typing_type_tree
 val perform_typing :
   Pb_parsing_parse_tree.proto list ->
   Pb_field_type.resolved Tt.proto_type list list
-(** [perform_typing parsed_tree] returned the type tree organized in groups 
-    of fully resolved types. Each group contains all the mutually recursive 
+(** [perform_typing parsed_tree] returned the type tree organized in groups
+    of fully resolved types. Each group contains all the mutually recursive
     types and the type group by reverse dependency order. *)
diff --git a/src/compilerlib/pb_typing_graph.ml b/src/compilerlib/pb_typing_graph.ml
index 8b5123c..23688b9 100644
--- a/src/compilerlib/pb_typing_graph.ml
+++ b/src/compilerlib/pb_typing_graph.ml
@@ -1,8 +1,8 @@
 (*
   The MIT License (MIT)
-  
+
   Copyright (c) 2016 Maxime Ransan <maxime.ransan@gmail.com>
-  
+
   Permission is hereby granted, free of charge, to any person obtaining a copy
   of this software and associated documentation files (the "Software"), to deal
   in the Software without restriction, including without limitation the rights
@@ -23,15 +23,13 @@
 
 *)
 
-module Int_map = Map.Make (struct
-  type t = int
+module Int_map = Pb_util.Int_map
 
-  let compare (x : int) (y : int) = Stdlib.compare x y
-end)
+type id = int
 
 type node = {
-  id: int;
-  sub: int list;
+  id: id;
+  sub: id list;
 }
 
 let create_node id sub = { id; sub }
@@ -51,7 +49,7 @@ module Tarjan = struct
 
   type tgraph = tnode Int_map.t
 
-  let reset g =
+  let new_tgraph (g : graph) : tgraph =
     Int_map.map
       (fun core -> { core; index = None; lowlink = None; on_stack = false })
       g
@@ -115,8 +113,8 @@ module Tarjan = struct
     ) else
       sccs, stack, index
 
-  let tarjan g =
-    let g = reset g in
+  let tarjan g : id list list =
+    let g = new_tgraph g in
     let sccs, _, _ =
       Int_map.fold
         (fun _ n (sccs, stack, index) ->
diff --git a/src/compilerlib/pb_typing_graph.mli b/src/compilerlib/pb_typing_graph.mli
index 8ea693a..c6119b4 100644
--- a/src/compilerlib/pb_typing_graph.mli
+++ b/src/compilerlib/pb_typing_graph.mli
@@ -1,8 +1,8 @@
 (*
   The MIT License (MIT)
-  
+
   Copyright (c) 2016 Maxime Ransan <maxime.ransan@gmail.com>
-  
+
   Permission is hereby granted, free of charge, to any person obtaining a copy
   of this software and associated documentation files (the "Software"), to deal
   in the Software without restriction, including without limitation the rights
@@ -27,6 +27,9 @@
 
 (** {2 Types} *)
 
+type id = int
+(** Identifier of a node in the graph *)
+
 type node
 (** Node in a graph, each node is identified using a unique [int] id *)
 
@@ -35,28 +38,28 @@ type graph
 
 (** {2 Creators } *)
 
-val create_node : int -> int list -> node
-(** [create_node id sub] create a node uniquely identified with [id] 
-    and connections to other nodes in [sub]. 
+val create_node : id -> id list -> node
+(** [create_node id sub] create a node uniquely identified with [id]
+    and connections to other nodes in [sub].
 
     The client application is responsible to ensure that the graph is
-    consistent, by adding all nodes identified in [sub] to the same 
-    graph later. 
+    consistent, by adding all nodes identified in [sub] to the same
+    graph later.
   *)
 
 val empty_graph : graph
 (** [empty_graph ()] create a new empty graph.*)
 
 val add_node : node -> graph -> graph
-(** [add_node node graph] add [node] to [graph] 
+(** [add_node node graph] add [node] to [graph]
  *)
 
 (** {2 Algorithms} *)
 
-val tarjan : graph -> int list list
+val tarjan : graph -> id list list
 (** [tarjan graph] compute the ordered list of strongly connected components of
-    a graph. 
+    a graph.
 
     The returned list is order in decreasing order of dependencies. This means
-    the last component of the list does not link to any other components. 
+    the last component of the list does not link to any other components.
  *)
diff --git a/src/compilerlib/pb_typing_recursion.ml b/src/compilerlib/pb_typing_recursion.ml
index 24c9635..0323c46 100644
--- a/src/compilerlib/pb_typing_recursion.ml
+++ b/src/compilerlib/pb_typing_recursion.ml
@@ -1,8 +1,8 @@
 (*
   The MIT License (MIT)
-  
+
   Copyright (c) 2016 Maxime Ransan <maxime.ransan@gmail.com>
-  
+
   Permission is hereby granted, free of charge, to any person obtaining a copy
   of this software and associated documentation files (the "Software"), to deal
   in the Software without restriction, including without limitation the rights
diff --git a/src/compilerlib/pb_typing_recursion.mli b/src/compilerlib/pb_typing_recursion.mli
index 505eede..6e1891e 100644
--- a/src/compilerlib/pb_typing_recursion.mli
+++ b/src/compilerlib/pb_typing_recursion.mli
@@ -1,8 +1,8 @@
 (*
   The MIT License (MIT)
-  
+
   Copyright (c) 2016 Maxime Ransan <maxime.ransan@gmail.com>
-  
+
   Permission is hereby granted, free of charge, to any person obtaining a copy
   of this software and associated documentation files (the "Software"), to deal
   in the Software without restriction, including without limitation the rights
@@ -27,19 +27,19 @@
 
 (** The protobuf syntax allows the definition of mutually recursive
     types, however this declaration is implicit (ie no dedicated syntax
-    to group together types which are mutually recursive). 
+    to group together types which are mutually recursive).
+
+    OCaml requires mutually recursive types to be defined with an explicit
+    syntax (using the [and] keyword), therefore it is necessary to
+    find all the mutually recursive protobuf types.
 
-    OCaml requires mutually recursive types to be defined with an explicit 
-    syntax (using the [and] keyword), therefore it is necessary to 
-    find all the mutually recursive protobuf types. 
-    
-    This module performs the recursion analysis using the Tarjan graph 
+    This module performs the recursion analysis using the Tarjan graph
     algorithm to find all the Strongly Connnected Components. *)
 
 module Tt = Pb_typing_type_tree
 
 val group :
   Pb_field_type.resolved Tt.proto -> Pb_field_type.resolved Tt.proto list
-(** [group types] returns the list of all the mutually recursive group 
-      of types in reverse order of dependency. In other the last group of 
+(** [group types] returns the list of all the mutually recursive group
+      of types in reverse order of dependency. In other the last group of
       types of the returned list don't depend on any other types. *)
diff --git a/src/compilerlib/pb_typing_resolution.ml b/src/compilerlib/pb_typing_resolution.ml
index 392405a..370460e 100644
--- a/src/compilerlib/pb_typing_resolution.ml
+++ b/src/compilerlib/pb_typing_resolution.ml
@@ -1,8 +1,8 @@
 (*
   The MIT License (MIT)
-  
+
   Copyright (c) 2016 Maxime Ransan <maxime.ransan@gmail.com>
-  
+
   Permission is hereby granted, free of charge, to any person obtaining a copy
   of this software and associated documentation files (the "Software"), to deal
   in the Software without restriction, including without limitation the rights
@@ -26,6 +26,7 @@
 module E = Pb_exception
 module Pt = Pb_parsing_parse_tree
 module Tt = Pb_typing_type_tree
+module String_map = Pb_util.Str_map
 
 module Types_by_scope = struct
   type type_ = Pb_field_type.unresolved Tt.proto_type
@@ -35,12 +36,6 @@ module Types_by_scope = struct
     | Tt.Enum { Tt.enum_name; _ } -> enum_name
     | Tt.Message { Tt.message_name; _ } -> message_name
 
-  module String_map = Map.Make (struct
-    type t = string
-
-    let compare (x : string) (y : string) = Stdlib.compare x y
-  end)
-
   type t = {
     name: string;
     types: type_ String_map.t;
@@ -78,25 +73,26 @@ module Types_by_scope = struct
     in
     aux t path
 
-  let rec print ?(level = 0) { types; subs; name } =
-    let pr = Printf.printf in
-    pr "%s %s\n" (Pb_util.indentation_prefix level) name;
+  let fpf = Printf.fprintf
+
+  let rec print ?(level = 0) oc { types; subs; name } : unit =
+    fpf oc "%s %s\n" (Pb_util.indentation_prefix level) name;
     String_map.iter
       (fun _ t ->
-        pr "%s +-%s \n" (Pb_util.indentation_prefix level) (name_of_type t))
+        fpf oc "%s +-%s \n" (Pb_util.indentation_prefix level) (name_of_type t))
       types;
-    String_map.iter (fun _ sub -> print ~level:(level + 1) sub) subs
+    String_map.iter (fun _ sub -> print oc ~level:(level + 1) sub) subs
 
-  let print t = print t
+  let print oc t = print oc t
   let empty = empty ()
 end
 (* Types_by_scope *)
 
-(* this function returns the type path of a message which is the 
- * packages followed by the enclosing message names and eventually 
- * the message name of the given type. 
+(* this function returns the type path of a message which is the
+ * packages followed by the enclosing message names and eventually
+ * the message name of the given type.
  *
- * If the type is an enum then [Failure] is raised. 
+ * If the type is an enum then [Failure] is raised.
  * TODO: change [Failure] to a [Pb_exception.Compilation_error] *)
 let type_path_of_type { Tt.scope; spec; _ } =
   match spec with
@@ -105,10 +101,10 @@ let type_path_of_type { Tt.scope; spec; _ } =
     let { Tt.packages; message_names } = scope in
     packages @ message_names @ [ message_name ]
 
-(* this function returns all the scope to search for a type starting 
- * by the most innner one first. 
+(* this function returns all the scope to search for a type starting
+ * by the most innner one first.
  *
- * If [message_scope] = ['Msg1'; 'Msg2'] and [field_scope] = ['Msg3'] then 
+ * If [message_scope] = ['Msg1'; 'Msg2'] and [field_scope] = ['Msg3'] then
  * the following scopes will be returned:
  * [
  *   ['Msg1'; 'Msg2'; 'Msg3'];  // This would be the scope of the current msg
@@ -130,13 +126,13 @@ let compute_search_type_paths unresolved_field_type message_type_path =
     List.rev @@ loop [] message_type_path
   )
 
-(* this function ensure that the default value of the field is correct 
- * with respect to its type when this latter is a builtin one. 
+(* this function ensure that the default value of the field is correct
+ * with respect to its type when this latter is a builtin one.
  *
- * in case the default value is invalid then an 
- * [Pb_exception.Compilation_error] is raised. 
+ * in case the default value is invalid then an
+ * [Pb_exception.Compilation_error] is raised.
  *
- * Note that this function also does type coersion when the default value 
+ * Note that this function also does type coersion when the default value
  * is an int and the builtin type is a float or double. *)
 let resolve_builtin_type_field_default field_name builtin_type field_default =
   match field_default with
@@ -184,14 +180,14 @@ let resolve_builtin_type_field_default field_name builtin_type field_default =
       E.invalid_default_value ~field_name
         ~info:"default value not supported for bytes" ())
 
-(* This function verifies that the default value for a used defined 
- * field is correct. 
+(* This function verifies that the default value for a used defined
+ * field is correct.
  *
- * In protobuf, only field which type is [enum] can have a default 
- * value. Field of type [message] can't. 
+ * In protobuf, only field which type is [enum] can have a default
+ * value. Field of type [message] can't.
  *
  * In the case the field is an enum then the default value must be
- * a litteral value which is one of the enum value. 
+ * a litteral value which is one of the enum value.
  *
  * If the validation fails then [Pb_exception.Compilation_error] is raised *)
 let resolve_enum_field_default field_name type_ field_default =
@@ -225,13 +221,13 @@ let resolve_enum_field_default field_name type_ field_default =
     E.invalid_default_value ~field_name
       ~info:"default value not supported for message" ()
 
-(* this function resolves both the type and the defaut value of a field 
+(* this function resolves both the type and the defaut value of a field
  * type. Note that it is necessary to verify both at the same time since
  * the default value must be of the same type as the field type in order
- * to be valid. 
+ * to be valid.
  *
- * For builtin the type the validation is trivial while for user defined 
- * type a search must be done for all the possible scopes the type 
+ * For builtin the type the validation is trivial while for user defined
+ * type a search must be done for all the possible scopes the type
  * might be in. *)
 let resolve_field_type_and_default t field_name field_type field_default
     message_type_path =
@@ -257,8 +253,8 @@ let resolve_field_type_and_default t field_name field_type field_default
     in
     aux (compute_search_type_paths unresolved_field_type message_type_path)
 
-(* this function resolves all the field type of the given type *)
-let resolve_type t type_ =
+(** this function resolves all the field type of the given type *)
+let resolve_type t type_ : int Tt.proto_type =
   let { Tt.scope; id; file_name; file_options; spec } = type_ in
 
   match spec with
diff --git a/src/compilerlib/pb_typing_resolution.mli b/src/compilerlib/pb_typing_resolution.mli
index 01ee72f..ca292a3 100644
--- a/src/compilerlib/pb_typing_resolution.mli
+++ b/src/compilerlib/pb_typing_resolution.mli
@@ -1,8 +1,8 @@
 (*
   The MIT License (MIT)
-  
+
   Copyright (c) 2016 Maxime Ransan <maxime.ransan@gmail.com>
-  
+
   Permission is hereby granted, free of charge, to any person obtaining a copy
   of this software and associated documentation files (the "Software"), to deal
   in the Software without restriction, including without limitation the rights
@@ -23,26 +23,26 @@
 
 *)
 
-(** Type resolution *)
+(** Type resolution.
 
-(** This module resolves all the user defined for the protobuf message 
-    fields; while built-in types were previously valided in 
-    [Pb_typing_validation] module. 
+    This module resolves all the user defined for the protobuf message
+    fields; while built-in types were previously valided in
+    [Pb_typing_validation] module.
 
     After this resolution is complete each user defined field type will
-    be a link to a unique message type identify by its id; in other word 
-    the type graph will is complete. 
+    be a link to a unique message type identify by its id; in other word
+    the type graph will is complete.
 
-    If the resolution fails then [Pb_exception.Compilation_error] is 
+    If the resolution fails then [Pb_exception.Compilation_error] is
     raised.
   *)
 
 module Tt = Pb_typing_type_tree
 
-(** Custom container for all the types (message or enums) which are 
-    organized by their scope. This allow efficient search of a type given 
-    its type path 
-    
+(** Custom container for all the types (message or enums) which are
+    organized by their scope. This allow efficient search of a type given
+    its type path
+
     The construction of this container is the first step in the type
     resolution. *)
 module Types_by_scope : sig
@@ -62,7 +62,7 @@ module Types_by_scope : sig
     Pb_field_type.unresolved Tt.proto_type
   (** find a protobuf type given its type path *)
 
-  val print : t -> unit
+  val print : out_channel -> t -> unit
   (** pretty print of the container *)
 end
 (* Types_by_scope *)
@@ -70,6 +70,6 @@ end
 val resolve_types :
   Pb_field_type.unresolved Tt.proto_type list ->
   Pb_field_type.resolved Tt.proto_type list
-(** [resolve_types types] resolves all the field types for all the [types]. 
-    If a field cannot be resolved then [Pb_exception.Compilation_error] is 
+(** [resolve_types types] resolves all the field types for all the [types].
+    If a field cannot be resolved then [Pb_exception.Compilation_error] is
     raised. *)
diff --git a/src/compilerlib/pb_typing_type_tree.ml b/src/compilerlib/pb_typing_type_tree.ml
index 9b7c225..97cd4ed 100644
--- a/src/compilerlib/pb_typing_type_tree.ml
+++ b/src/compilerlib/pb_typing_type_tree.ml
@@ -1,8 +1,8 @@
 (*
   The MIT License (MIT)
-  
+
   Copyright (c) 2016 Maxime Ransan <maxime.ransan@gmail.com>
-  
+
   Permission is hereby granted, free of charge, to any person obtaining a copy
   of this software and associated documentation files (the "Software"), to deal
   in the Software without restriction, including without limitation the rights
@@ -25,9 +25,9 @@
 
 module Pt = Pb_parsing_parse_tree
 
-(** Protobuf typed tree. 
+(** Protobuf typed tree.
 
-    The typetree type is parametrized to allow for 2 phase compilation. 
+    The typetree type is parametrized to allow for 2 phase compilation.
  *)
 
 type ('a, 'b) field = {
@@ -36,8 +36,8 @@ type ('a, 'b) field = {
   field_default: Pb_option.constant option;
   field_options: Pb_option.set;
 }
-(** Field definition. 
-    
+(** Field definition.
+
     {ul
     {- ['a] is for [unresolved] or [resolved]}
     {- ['b] is for [field_label] to account for both normal and one of fields.}
@@ -66,10 +66,10 @@ type type_scope = {
   packages: string list;
   message_names: string list;
 }
-(** Type scope 
-      
-    The scope of a type (message or enum) is defined by the package 
-    (defined in the top of the proto file as well as the messages above 
+(** Type scope
+
+    The scope of a type (message or enum) is defined by the package
+    (defined in the top of the proto file as well as the messages above
     it since a message definition can be nested *)
 
 (** item for the message body *)
diff --git a/src/compilerlib/pb_typing_util.ml b/src/compilerlib/pb_typing_util.ml
index 72c0bd5..eab918b 100644
--- a/src/compilerlib/pb_typing_util.ml
+++ b/src/compilerlib/pb_typing_util.ml
@@ -1,8 +1,8 @@
 (*
   The MIT License (MIT)
-  
+
   Copyright (c) 2016 Maxime Ransan <maxime.ransan@gmail.com>
-  
+
   Permission is hereby granted, free of charge, to any person obtaining a copy
   of this software and associated documentation files (the "Software"), to deal
   in the Software without restriction, including without limitation the rights
diff --git a/src/compilerlib/pb_typing_util.mli b/src/compilerlib/pb_typing_util.mli
index ce052f5..2f79c8f 100644
--- a/src/compilerlib/pb_typing_util.mli
+++ b/src/compilerlib/pb_typing_util.mli
@@ -1,8 +1,8 @@
 (*
   The MIT License (MIT)
-  
+
   Copyright (c) 2016 Maxime Ransan <maxime.ransan@gmail.com>
-  
+
   Permission is hereby granted, free of charge, to any person obtaining a copy
   of this software and associated documentation files (the "Software"), to deal
   in the Software without restriction, including without limitation the rights
@@ -25,7 +25,7 @@
 
 (** Typed tree utilities *)
 
-(** This module provides accessor/creator function to facilitate 
+(** This module provides accessor/creator function to facilitate
     the manipulation of the type tree. *)
 
 module Pt = Pb_parsing_parse_tree
@@ -51,14 +51,14 @@ val field_default : ('a, 'b) Tt.field -> Pb_option.constant option
 val field_options : ('a, 'b) Tt.field -> Pb_option.set
 
 val field_option : ('a, 'b) Tt.field -> string -> Pb_option.value option
-(** [field_option field option_name] returns the constant associated with 
+(** [field_option field option_name] returns the constant associated with
     [option_name]. If the fields options does not contain [option_name] [None]
     is returned.
   *)
 
 val type_of_id : 'a Tt.proto -> int -> 'a Tt.proto_type
-(** [type_of_id all_types id] returns the type associated with the given id, 
-    raise [Not_found] if the type is not in the all_types. 
+(** [type_of_id all_types id] returns the type associated with the given id,
+    raise [Not_found] if the type is not in the all_types.
   *)
 
 val string_of_message : int -> Tt.type_scope -> 'a Tt.message -> string
@@ -68,7 +68,7 @@ val enum_option : Tt.enum -> string -> Pb_option.value option
 (** {2 Accessor for Tt.type} *)
 
 val type_name_of_type : 'a Tt.proto_type -> string
-(** [type_name_of_type t] returns the type name (as defined in the 
+(** [type_name_of_type t] returns the type name (as defined in the
     message file) of [t].
  *)
 
@@ -76,8 +76,8 @@ val type_scope_of_type : 'a Tt.proto_type -> Tt.type_scope
 (** [type_scope_of_type t] returns the scope of type [t]. *)
 
 val is_empty_message : 'a Tt.proto_type -> bool
-(** [is_empty_message t] returns true if [t] is a message type and 
-    has no fields defined. 
+(** [is_empty_message t] returns true if [t] is a message type and
+    has no fields defined.
  *)
 
 (** {2 Creator} *)
diff --git a/src/compilerlib/pb_typing_validation.ml b/src/compilerlib/pb_typing_validation.ml
index 9e643c1..4dd017f 100644
--- a/src/compilerlib/pb_typing_validation.ml
+++ b/src/compilerlib/pb_typing_validation.ml
@@ -1,8 +1,8 @@
 (*
   The MIT License (MIT)
-  
+
   Copyright (c) 2016 Maxime Ransan <maxime.ransan@gmail.com>
-  
+
   Permission is hereby granted, free of charge, to any person obtaining a copy
   of this software and associated documentation files (the "Software"), to deal
   in the Software without restriction, including without limitation the rights
@@ -28,7 +28,7 @@ module Pt = Pb_parsing_parse_tree
 module Tt = Pb_typing_type_tree
 module Typing_util = Pb_typing_util
 
-let scope_of_package = function
+let scope_of_package : string option -> Tt.type_scope = function
   | Some s ->
     {
       Typing_util.empty_scope with
@@ -36,21 +36,24 @@ let scope_of_package = function
     }
   | None -> Typing_util.empty_scope
 
-let get_default field_name field_options =
+let get_default field_name field_options : Pb_option.constant option =
   match Pb_option.get field_options "default" with
   | Some (Pb_option.Scalar_value constant) -> Some constant
   | Some (Pb_option.Message_literal _) ->
     E.invalid_default_value ~field_name
       ~info:"message literals are unsupported for default values" ()
+  | Some (Pb_option.List_literal _) ->
+    E.invalid_default_value ~field_name
+      ~info:"list literals are unsupported for default values" ()
   | None -> None
 
-let compile_field_p1 field_parsed =
+let compile_field_p1 field_parsed : _ Tt.field =
   let { Pt.field_type; Pt.field_options; Pt.field_name; _ } = field_parsed in
 
   let field_default = get_default field_name field_options in
   { Tt.field_parsed; Tt.field_type; Tt.field_default; Tt.field_options }
 
-let compile_map_p1 map_parsed =
+let compile_map_p1 map_parsed : _ Tt.map_field =
   let {
     Pt.map_name;
     Pt.map_number;
@@ -63,7 +66,7 @@ let compile_map_p1 map_parsed =
 
   Tt.{ map_name; map_number; map_key_type; map_value_type; map_options }
 
-let compile_oneof_p1 oneof_parsed =
+let compile_oneof_p1 oneof_parsed : _ Tt.oneof =
   let init =
     {
       Tt.oneof_name = oneof_parsed.Pt.oneof_name;
@@ -96,10 +99,11 @@ let rec list_assoc2 x = function
     else
       list_assoc2 x l
 
-let make_proto_type ~file_name ~file_options ~id ~scope ~spec =
+let make_proto_type ~file_name ~file_options ~id ~scope ~spec : _ Tt.proto_type
+    =
   { Tt.id; Tt.scope; Tt.file_name; Tt.file_options; Tt.spec }
 
-(* compile a [Pbpt] enum to a [Pbtt] type *)
+(** compile a [Pbpt] enum to a [Pbtt] type *)
 let compile_enum_p1 ?(parent_options = Pb_option.empty) file_name file_options
     scope parsed_enum =
   let { Pt.enum_id; enum_name; enum_body } = parsed_enum in
@@ -135,10 +139,10 @@ let compile_enum_p1 ?(parent_options = Pb_option.empty) file_name file_options
 
   make_proto_type ~file_name ~file_options ~id:enum_id ~scope ~spec
 
-(* compile a [Pbpt] message a list of [Pbtt] types (ie messages can 
- * defined more than one type). *)
+(** compile a [Pbpt] message a list of [Pbtt] types (ie messages can
+    defined more than one type). *)
 let rec validate_message ?(parent_options = Pb_option.empty) file_name
-    file_options message_scope parsed_message =
+    file_options message_scope parsed_message : _ Tt.proto_type list =
   let { Pt.id; Pt.message_name; Pt.message_body } = parsed_message in
 
   let { Tt.message_names; _ } = message_scope in
@@ -212,13 +216,13 @@ let rec validate_message ?(parent_options = Pb_option.empty) file_name
      [Pb_parsing_util.verify_message] along with
      the proto3 invariant. *)
 
-  (* Both field name and field number must be unique 
-   * within a message scope. This includes the field in a 
-   * oneof field inside the message. 
-   *
-   * This function verifies this constrain and raises
-   * the corresponding Duplicated_field_number exception in 
-   * case it is violated. *)
+  (* Both field name and field number must be unique
+     within a message scope. This includes the field in a
+     oneof field inside the message.
+
+     This function verifies this constrain and raises
+     the corresponding Duplicated_field_number exception in
+     case it is violated. *)
   let validate_duplicate (number_index : (int * string) list) name number =
     if
       not_found (fun () -> ignore @@ List.assoc number number_index)
@@ -237,7 +241,7 @@ let rec validate_message ?(parent_options = Pb_option.empty) file_name
   in
 
   ignore
-  @@ List.fold_left
+    (List.fold_left
        (fun number_index -> function
          | Tt.Message_field field -> validate_duplicate_field number_index field
          | Tt.Message_oneof_field { Tt.oneof_fields; _ } ->
@@ -245,23 +249,23 @@ let rec validate_message ?(parent_options = Pb_option.empty) file_name
          | Tt.Message_map_field m ->
            let { Tt.map_name; map_number; _ } = m in
            validate_duplicate number_index map_name map_number)
-       [] message_body;
+       [] message_body
+      : _ list);
 
   let spec =
-    Tt.(
-      Message
-        {
-          extensions = acc.Acc.extensions;
-          message_options = acc.Acc.options;
-          message_name;
-          message_body;
-        })
+    Tt.Message
+      {
+        Tt.extensions = acc.Acc.extensions;
+        message_options = acc.Acc.options;
+        message_name;
+        message_body;
+      }
   in
 
   acc.Acc.all_types
   @ [ make_proto_type ~file_name ~file_options ~id ~scope:message_scope ~spec ]
 
-let validate proto =
+let validate (proto : Pt.proto) : _ Tt.proto_type list =
   let { Pt.package; Pt.proto_file_name; messages; enums; file_options; _ } =
     proto
   in
diff --git a/src/compilerlib/pb_typing_validation.mli b/src/compilerlib/pb_typing_validation.mli
index e161615..b2b53dc 100644
--- a/src/compilerlib/pb_typing_validation.mli
+++ b/src/compilerlib/pb_typing_validation.mli
@@ -1,8 +1,8 @@
 (*
   The MIT License (MIT)
-  
+
   Copyright (c) 2016 Maxime Ransan <maxime.ransan@gmail.com>
-  
+
   Permission is hereby granted, free of charge, to any person obtaining a copy
   of this software and associated documentation files (the "Software"), to deal
   in the Software without restriction, including without limitation the rights
@@ -26,10 +26,10 @@
 (** Typed tree construction and validation *)
 
 (** This module focuses on building the Type tree from the parse tree
-    and provides a number of validation in doing so. 
+    and provides a number of validation in doing so.
 
     Not that the field types are not resolved, and the returned type tree is
-    hence of type [unresolved_field_type Pb_typing_type_tree.proto_type]. 
+    hence of type [unresolved_field_type Pb_typing_type_tree.proto_type].
     The type resolution is done in [Pb_typing_resolution] *)
 
 module Pt = Pb_parsing_parse_tree
@@ -38,7 +38,7 @@ module Tt = Pb_typing_type_tree
 (** {2 Type tree construction} *)
 
 val validate : Pt.proto -> Pb_field_type.unresolved Tt.proto
-(** [validate file_name proto] makes a first phase compilation of the 
+(** [validate file_name proto] makes a first phase compilation of the
     parsed tree.  *)
 
 (** {2 Testing Only} *)
diff --git a/src/compilerlib/pb_util.ml b/src/compilerlib/pb_util.ml
index 25f9e26..e45840e 100644
--- a/src/compilerlib/pb_util.ml
+++ b/src/compilerlib/pb_util.ml
@@ -1,8 +1,8 @@
 (*
   The MIT License (MIT)
-  
+
   Copyright (c) 2016 Maxime Ransan <maxime.ransan@gmail.com>
-  
+
   Permission is hereby granted, free of charge, to any person obtaining a copy
   of this software and associated documentation files (the "Software"), to deal
   in the Software without restriction, including without limitation the rights
@@ -70,7 +70,6 @@ module Option = struct
     | None -> "None"
     | Some x -> Printf.sprintf "Some(%s)" (f x)
 end
-(* Option *)
 
 let indentation_prefix =
   let h = Hashtbl.create 16 in
@@ -129,4 +128,11 @@ module List = struct
       | None -> filter_map f tl
       | Some x -> x :: filter_map f tl)
 end
-(* List *)
+
+module Int_map = Map.Make (struct
+  type t = int
+
+  let compare (x : int) (y : int) = Stdlib.compare x y
+end)
+
+module Str_map = Map.Make (String)
diff --git a/src/compilerlib/pb_util.mli b/src/compilerlib/pb_util.mli
index e1d0208..ac88b69 100644
--- a/src/compilerlib/pb_util.mli
+++ b/src/compilerlib/pb_util.mli
@@ -1,8 +1,8 @@
 (*
   The MIT License (MIT)
-  
+
   Copyright (c) 2016 Maxime Ransan <maxime.ransan@gmail.com>
-  
+
   Permission is hereby granted, free of charge, to any person obtaining a copy
   of this software and associated documentation files (the "Software"), to deal
   in the Software without restriction, including without limitation the rights
@@ -27,10 +27,10 @@
 
 val rev_split_by_char : char -> string -> string list
 (** [rev_split_by_char c s] will split the string [s] using the delimieter [c]
-    and return the component in reverse order (ie from right to left). 
-    
+    and return the component in reverse order (ie from right to left).
+
     For instance when splitting a filename with the '.' character, the file
-    extension will be the head of the returned list. 
+    extension will be the head of the returned list.
  *)
 
 val string_of_string_list : string list -> string
@@ -44,25 +44,24 @@ val indentation_prefix : int -> string
 
 module Option : sig
   val default : 'a -> 'a option -> 'a
-  (** [option_default x o] returns [x] is [o] is [None] otherwise [y] 
-      when [o] is [Some y]. 
+  (** [option_default x o] returns [x] is [o] is [None] otherwise [y]
+      when [o] is [Some y].
    *)
 
   val min_value : 'a option -> 'a option -> 'a option
-  (** [min_value x y] returns the min value of the x and y if x and y are of 
+  (** [min_value x y] returns the min value of the x and y if x and y are of
       the form [Some _]. If [x] or [y] is [None] then [Invalid_argument]
       exception is raised *)
 
   val eq_value : 'a option -> 'a option -> bool
-  (** [eq_value x y] returns [true] if both [x] and [y] are of 
+  (** [eq_value x y] returns [true] if both [x] and [y] are of
       the form [Some _]. If [x] or [y] is [None] then [Invalid_argument]
       exception is raised *)
 
   val string_of_option : ('a -> string) -> 'a option -> string
-  (** [string_of_option f x] returns string representation of [x] using 
+  (** [string_of_option f x] returns string representation of [x] using
       [f] when [x] is of the form [Some _]. *)
 end
-(* Option *)
 
 val read_file : string -> string
 val protect : finally:(unit -> unit) -> (unit -> 'a) -> 'a
@@ -72,15 +71,17 @@ module List : sig
   (** [pop_last l] removes the last element from the list *)
 
   val apply_until : ('a -> 'b option) -> 'a list -> 'b option
-  (** [apply_until f l] applies [f ei] until it returns [Some x] 
-      
-      If the end of the list is reached without [f] returning [Some _] then 
-      [None] is returned. 
+  (** [apply_until f l] applies [f ei] until it returns [Some x]
+
+      If the end of the list is reached without [f] returning [Some _] then
+      [None] is returned.
    *)
 
   val filter_map : ('a -> 'b option) -> 'a list -> 'b list
   (** [filter_map f l] returns the list of element [x] for which [f] returned
-      [Some x]. The length of the returned list will be less or equal than 
+      [Some x]. The length of the returned list will be less or equal than
       the length of the input list [l]. *)
 end
-(* List *)
+
+module Str_map : Map.S with type key = string
+module Int_map : Map.S with type key = int
diff --git a/src/dune b/src/dune
new file mode 100644
index 0000000..f6dc452
--- /dev/null
+++ b/src/dune
@@ -0,0 +1,3 @@
+
+(env
+  (_ (flags :standard -warn-error -a+8 -w +a-4-40-41-42-44-48-70)))
diff --git a/src/examples/dune b/src/examples/dune
index f9d412c..882a7fa 100644
--- a/src/examples/dune
+++ b/src/examples/dune
@@ -1,47 +1,43 @@
 (rule
- (targets example01_types.ml example01_pb.ml example01_pp.ml
-   example01_types.mli example01_pb.mli example01_pp.mli)
+ (targets example01.ml example01.mli)
  (deps example01.proto)
  (action
-  (run ocaml-protoc -binary -pp -ml_out ./ %{deps})))
+  (run ocaml-protoc --binary --pp --ml_out ./ %{deps})))
 
 (test
  (name example01)
- (modules example01 example01_types example01_pp example01_pb)
+ (modules t_example01 example01)
  (libraries pbrt))
 
 (rule
- (targets example03_types.ml example03_pp.ml example03_types.mli
-   example03_pp.mli)
+ (targets example03.ml example03.mli)
  (deps example03.proto)
  (action
-  (run ocaml-protoc -pp -ml_out ./ %{deps})))
+  (run ocaml-protoc --pp --ml_out ./ %{deps})))
 
 (test
  (name example03)
- (modules example03 example03_types example03_pp)
+ (modules t_example03 example03)
  (libraries pbrt))
 
 (rule
- (targets example04_types.ml example04_pp.ml example04_types.mli
-   example04_pp.mli)
+ (targets example04.ml example04.mli)
  (deps example04.proto)
  (action
-  (run ocaml-protoc -pp -ml_out ./ %{deps})))
+  (run ocaml-protoc --pp --ml_out ./ %{deps})))
 
 (test
  (name example04)
- (modules example04 example04_types example04_pp)
+ (modules t_example04 example04)
  (libraries pbrt))
 
 (rule
- (targets example05_types.ml example05_pb.ml example05_pp.ml
-   example05_types.mli example05_pb.mli example05_pp.mli)
+ (targets example05.ml example05.mli)
  (deps example05.proto)
  (action
-  (run ocaml-protoc -binary -pp -ml_out ./ %{deps})))
+  (run ocaml-protoc --binary --pp --ml_out ./ %{deps})))
 
 (test
  (name example05)
- (modules example05 example05_types example05_pp example05_pb)
+ (modules t_example05 example05)
  (libraries pbrt))
diff --git a/src/examples/example03.ml b/src/examples/example03.ml
deleted file mode 100644
index be58b15..0000000
--- a/src/examples/example03.ml
+++ /dev/null
@@ -1,3 +0,0 @@
-let () =
-  let v = Example03_types.Some "This is an example" in
-  Format.fprintf Format.std_formatter "%a\n" Example03_pp.pp_string_some v
diff --git a/src/examples/example01.ml b/src/examples/t_example01.ml
similarity index 78%
rename from src/examples/example01.ml
rename to src/examples/t_example01.ml
index dfdd470..ba1d454 100644
--- a/src/examples/example01.ml
+++ b/src/examples/t_example01.ml
@@ -1,7 +1,7 @@
 let () =
   (* Create OCaml value of generated type *)
   let person =
-    Example01_types.
+    Example01.
       {
         name = "John Doe";
         id = 1234l;
@@ -10,7 +10,7 @@ let () =
       }
   in
 
-  print_endline @@ Format.asprintf "person: %a" Example01_pp.pp_person person;
+  print_endline @@ Format.asprintf "person: %a" Example01.pp_person person;
 
   (*
   let json = 
@@ -30,13 +30,13 @@ let () =
   *)
   let binary =
     let encoder = Pbrt.Encoder.create () in
-    Example01_pb.encode_person person encoder;
+    Example01.encode_pb_person person encoder;
     Pbrt.Encoder.to_bytes encoder
   in
 
   let person' =
     let decoder = Pbrt.Decoder.of_bytes binary in
-    Example01_pb.decode_person decoder
+    Example01.decode_pb_person decoder
   in
 
   assert (person = person')
diff --git a/src/examples/t_example03.ml b/src/examples/t_example03.ml
new file mode 100644
index 0000000..c2faf36
--- /dev/null
+++ b/src/examples/t_example03.ml
@@ -0,0 +1,3 @@
+let () =
+  let v = Example03.Some "This is an example" in
+  Format.fprintf Format.std_formatter "%a\n" Example03.pp_string_some v
diff --git a/src/examples/example04.ml b/src/examples/t_example04.ml
similarity index 80%
rename from src/examples/example04.ml
rename to src/examples/t_example04.ml
index 34f8370..6710c9c 100644
--- a/src/examples/example04.ml
+++ b/src/examples/t_example04.ml
@@ -1,5 +1,4 @@
-open Example04_types
-open Example04_pp
+open Example04
 
 let ( @+ ) value next = Cons { value; next }
 
diff --git a/src/examples/example05.ml b/src/examples/t_example05.ml
similarity index 79%
rename from src/examples/example05.ml
rename to src/examples/t_example05.ml
index e255b9f..040068b 100644
--- a/src/examples/example05.ml
+++ b/src/examples/t_example05.ml
@@ -5,7 +5,7 @@ let () =
   Hashtbl.add details "blah" "blah";
 
   let person =
-    Example05_types.
+    Example05.
       {
         name = "John Doe";
         id = 1234l;
@@ -17,7 +17,7 @@ let () =
 
   (* Create a Protobuf encoder and encode value *)
   let encoder = Pbrt.Encoder.create () in
-  Example05_pb.encode_person person encoder;
+  Example05.encode_pb_person person encoder;
 
   (* Output the protobuf message to a file *)
   let oc = open_out "myfile" in
@@ -33,5 +33,5 @@ let () =
     close_in ic;
     bytes
   in
-  let person = Example05_pb.decode_person (Pbrt.Decoder.of_bytes bytes) in
-  Format.fprintf Format.std_formatter "%a" Example05_pp.pp_person person
+  let person = Example05.decode_pb_person (Pbrt.Decoder.of_bytes bytes) in
+  Format.fprintf Format.std_formatter "%a" Example05.pp_person person
diff --git a/src/ocaml-protoc/ocaml_protoc.ml b/src/ocaml-protoc/ocaml_protoc.ml
index 1570f68..09deb01 100644
--- a/src/ocaml-protoc/ocaml_protoc.ml
+++ b/src/ocaml-protoc/ocaml_protoc.ml
@@ -50,7 +50,7 @@ let () =
       Compilation.compile cmdline cmd_line_file_options
     in
 
-    Generation.generate_code ocaml_types proto_file_options cmdline
+    Generation.generate_code ocaml_types ~proto_file_options cmdline
   with exn ->
     Printf.eprintf "%s\n" (Printexc.to_string exn);
     exit 1
diff --git a/src/ocaml-protoc/ocaml_protoc_cmdline.ml b/src/ocaml-protoc/ocaml_protoc_cmdline.ml
index f5c19c2..4002a6e 100644
--- a/src/ocaml-protoc/ocaml_protoc_cmdline.ml
+++ b/src/ocaml-protoc/ocaml_protoc_cmdline.ml
@@ -52,25 +52,25 @@ module File_options = struct
       ocaml_all_types_ppx = None;
     }
 
-  (** Compute the command line arguments for be used with the Arg module.  *)
+  (** Compute the command line arguments be used with the {!Arg} module.  *)
   let cmd_line_args t =
     [
-      ( "-int32_type",
+      ( "--int32_type",
         Arg.String
           (function
           | "int_t" -> t.int32_type <- Some "int_t"
           | x -> failwith @@ Printf.sprintf "Invalid int32_type value %s" x),
         " int32_type file option" );
-      ( "-int64_type",
+      ( "--int64_type",
         Arg.String
           (function
           | "int_t" -> t.int64_type <- Some "int_t"
           | x -> failwith @@ Printf.sprintf "Invalid int64_type value %s" x),
         " int64_type file option" );
-      ( "-ocaml_file_ppx",
+      ( "--ocaml_file_ppx",
         Arg.String (fun s -> t.ocaml_file_ppx <- Some s),
         " ocaml_file_ppx file option" );
-      ( "-ocaml_all_types_ppx",
+      ( "--ocaml_all_types_ppx",
         Arg.String (fun s -> t.ocaml_all_types_ppx <- Some s),
         " ocaml_all_types_ppx file option" );
     ]
@@ -133,15 +133,18 @@ module Cmdline = struct
 
   let cmd_line_args t =
     [
-      "-yojson", Arg.Set t.yojson, " generate yojson encoding";
-      "-bs", Arg.Set t.bs, " generate BuckleScript encoding";
-      "-binary", Arg.Set t.binary, " generate binary encoding";
-      "-pp", Arg.Set t.pp, " generate pretty print functions";
+      "--yojson", Arg.Set t.yojson, " generate yojson encoding";
+      "--bs", Arg.Set t.bs, " generate BuckleScript encoding";
+      "--binary", Arg.Set t.binary, " generate binary encoding";
+      "--pp", Arg.Set t.pp, " generate pretty print functions";
       ( "-I",
         Arg.String (fun s -> t.include_dirs <- s :: t.include_dirs),
         " include directories" );
-      "-ml_out", Arg.String (fun s -> t.ml_out <- s), " output directory";
-      ( "-unsigned",
+      "--ml_out", Arg.String (fun s -> t.ml_out <- s), " output directory";
+      ( "--debug",
+        Arg.Unit (fun () -> Pb_logger.setup_from_out_channel stderr),
+        " print logs on stderr" );
+      ( "--unsigned",
         Arg.Set t.unsigned_tag,
         " tag uint32 and uint64 types with `unsigned" );
     ]
diff --git a/src/ocaml-protoc/ocaml_protoc_generation.ml b/src/ocaml-protoc/ocaml_protoc_generation.ml
index ea92165..55ba78b 100644
--- a/src/ocaml-protoc/ocaml_protoc_generation.ml
+++ b/src/ocaml-protoc/ocaml_protoc_generation.ml
@@ -23,55 +23,25 @@
 
 *)
 
-module E = Pb_exception
 module Ot = Pb_codegen_ocaml_type
 module F = Pb_codegen_formatting
 module Codegen_util = Pb_codegen_util
 module Cmdline = Ocaml_protoc_cmdline.Cmdline
+module CG_all = Pb_codegen_all
 
-type codegen_f = ?and_:unit -> Ot.type_ -> F.scope -> bool
-
-let generate_for_all_types ocaml_types sc (f : codegen_f) ocamldoc_title : unit
-    =
-  (match ocamldoc_title with
-  | None -> ()
-  | Some ocamldoc_title ->
-    F.empty_line sc;
-    F.linep sc "(** {2 %s} *)" ocamldoc_title;
-    F.empty_line sc);
-
-  List.iter
-    (fun types ->
-      let (_ : bool) =
-        List.fold_left
-          (fun first type_ ->
-            let has_encoded =
-              if first then
-                f type_ sc
-              else
-                f ~and_:() type_ sc
-            in
-            F.empty_line sc;
-            first && not has_encoded)
-          true types
-      in
-      ())
-    ocaml_types
-
-let open_files cmdline file_suffix f =
+let open_files cmdline (f : ml:out_channel -> mli:out_channel -> 'a) : 'a =
   let { Cmdline.ml_out; Cmdline.proto_file_name; _ } = cmdline in
 
   let out_file_name =
     let proto_file_name = Filename.basename proto_file_name in
     let out_basename =
       Codegen_util.caml_file_name_of_proto_file_name ~proto_file_name
-        ~file_suffix
     in
     Filename.concat ml_out out_basename
   in
 
-  Printf.printf "Generating %s.mli\n%!" out_file_name;
-  Printf.printf "Generating %s.ml\n%!" out_file_name;
+  Pb_logger.log "Generating %s.mli\n%!" out_file_name;
+  Pb_logger.log "Generating %s.ml\n%!" out_file_name;
 
   let oc_mli = open_out @@ out_file_name ^ ".mli" in
   let oc_ml = open_out @@ out_file_name ^ ".ml" in
@@ -80,169 +50,38 @@ let open_files cmdline file_suffix f =
     ~finally:(fun () ->
       close_out oc_mli;
       close_out oc_ml)
-    (fun () -> f (oc_mli, oc_ml))
-
-let generate_type_and_default ocaml_types proto_file_options cmdline : unit =
-  let file_suffix = "types" in
-
-  let print_ppx sc =
-    match Pb_option.get proto_file_options "ocaml_file_ppx" with
-    | None -> ()
-    | Some Pb_option.(Scalar_value (Constant_string s)) ->
-      F.linep sc "[@@@%s]" s
-    | _ -> E.invalid_ppx_extension_option cmdline.Cmdline.proto_file_name
+    (fun () -> f ~ml:oc_ml ~mli:oc_mli)
+
+let generate_code ocaml_types ~proto_file_options cmdline : unit =
+  let plugins =
+    List.flatten
+      [
+        (if !(cmdline.Cmdline.yojson) then
+          [ Pb_codegen_encode_yojson.plugin; Pb_codegen_decode_yojson.plugin ]
+        else
+          []);
+        (if !(cmdline.Cmdline.binary) then
+          [ Pb_codegen_encode_binary.plugin; Pb_codegen_decode_binary.plugin ]
+        else
+          []);
+        (if !(cmdline.Cmdline.pp) then
+          [ Pb_codegen_pp.plugin ]
+        else
+          []);
+        (if !(cmdline.Cmdline.bs) then
+          [ Pb_codegen_encode_bs.plugin; Pb_codegen_decode_bs.plugin ]
+        else
+          []);
+      ]
   in
 
-  (* .ml file *)
-  let ml_sc = F.empty_scope () in
-  F.line ml_sc "[@@@ocaml.warning \"-27-30-39\"]";
-  F.empty_line ml_sc;
-  print_ppx ml_sc;
-  F.empty_line ml_sc;
-  generate_for_all_types ocaml_types ml_sc Pb_codegen_types.gen_struct None;
-  generate_for_all_types ocaml_types ml_sc Pb_codegen_default.gen_struct None;
-
-  (* .mli file *)
-  let mli_sc = F.empty_scope () in
-  F.linep mli_sc "(** %s Types *)"
-    (Filename.basename cmdline.Cmdline.proto_file_name);
-  F.empty_line mli_sc;
-  print_ppx mli_sc;
-  F.empty_line mli_sc;
-  generate_for_all_types ocaml_types mli_sc Pb_codegen_types.gen_sig
-    (Some Pb_codegen_types.ocamldoc_title);
-  generate_for_all_types ocaml_types mli_sc Pb_codegen_default.gen_sig
-    (Some Pb_codegen_default.ocamldoc_title);
-
-  open_files cmdline file_suffix @@ fun (sig_oc, struct_oc) ->
-  output_string struct_oc (F.print ml_sc);
-  output_string sig_oc (F.print mli_sc)
-
-let generate_mutable_records gen_file_suffix ocaml_types (sc : F.scope) : unit =
-  let ocaml_types = List.flatten ocaml_types in
-  List.iter
-    (fun { Ot.spec; module_prefix; _ } ->
-      match spec with
-      | Ot.Record r ->
-        Pb_codegen_types.gen_record_mutable module_prefix r sc;
-        F.empty_line sc;
-        Pb_codegen_default.gen_record_mutable ~gen_file_suffix ~module_prefix r
-          sc;
-        F.empty_line sc
-      | _ -> ())
-    ocaml_types
-
-let generate_yojson ocaml_types cmdline =
-  let file_suffix = "yojson" in
-
-  (* .ml file *)
-  let ml_sc = F.empty_scope () in
-  F.line ml_sc "[@@@ocaml.warning \"-27-30-39\"]";
-  F.empty_line ml_sc;
-  generate_mutable_records file_suffix ocaml_types ml_sc;
-  F.empty_line ml_sc;
-
-  generate_for_all_types ocaml_types ml_sc Pb_codegen_decode_yojson.gen_struct
-    None;
-  generate_for_all_types ocaml_types ml_sc Pb_codegen_encode_yojson.gen_struct
-    None;
-
-  (* .mli file *)
-  let mli_sc = F.empty_scope () in
-  F.linep mli_sc "(** %s YoJSON Encoding *)"
-    (Filename.basename cmdline.Cmdline.proto_file_name);
-  F.empty_line mli_sc;
-  generate_for_all_types ocaml_types mli_sc Pb_codegen_encode_yojson.gen_sig
-    (Some Pb_codegen_encode_yojson.ocamldoc_title);
-  generate_for_all_types ocaml_types mli_sc Pb_codegen_decode_yojson.gen_sig
-    (Some Pb_codegen_decode_yojson.ocamldoc_title);
-
-  open_files cmdline file_suffix @@ fun (sig_oc, struct_oc) ->
-  output_string struct_oc (F.print ml_sc);
-  output_string sig_oc (F.print mli_sc)
-
-let generate_bs ocaml_types cmdline : unit =
-  let file_suffix = "bs" in
-
-  (* .ml file *)
-  let ml_sc = F.empty_scope () in
-  F.line ml_sc "[@@@ocaml.warning \"-27-30-39\"]";
-  F.empty_line ml_sc;
-  generate_mutable_records file_suffix ocaml_types ml_sc;
-  F.empty_line ml_sc;
-
-  generate_for_all_types ocaml_types ml_sc Pb_codegen_decode_bs.gen_struct None;
-  generate_for_all_types ocaml_types ml_sc Pb_codegen_encode_bs.gen_struct None;
-
-  (* .mli file *)
-  let mli_sc = F.empty_scope () in
-  F.linep mli_sc "(** %s BuckleScript Encoding *)"
-    (Filename.basename cmdline.Cmdline.proto_file_name);
-  F.empty_line mli_sc;
-  generate_for_all_types ocaml_types mli_sc Pb_codegen_encode_bs.gen_sig
-    (Some Pb_codegen_encode_bs.ocamldoc_title);
-  generate_for_all_types ocaml_types mli_sc Pb_codegen_decode_bs.gen_sig
-    (Some Pb_codegen_decode_bs.ocamldoc_title);
-
-  open_files cmdline file_suffix @@ fun (sig_oc, struct_oc) ->
-  output_string struct_oc (F.print ml_sc);
-  output_string sig_oc (F.print mli_sc)
-
-let generate_binary ocaml_types cmdline =
-  let file_suffix = "pb" in
-
-  (* .ml file *)
-  let ml_sc = F.empty_scope () in
-  F.line ml_sc "[@@@ocaml.warning \"-27-30-39\"]";
-  F.empty_line ml_sc;
-  generate_mutable_records file_suffix ocaml_types ml_sc;
-  F.empty_line ml_sc;
-
-  generate_for_all_types ocaml_types ml_sc Pb_codegen_decode_binary.gen_struct
-    None;
-  generate_for_all_types ocaml_types ml_sc Pb_codegen_encode_binary.gen_struct
-    None;
-
-  (* .mli file *)
-  let mli_sc = F.empty_scope () in
-  F.linep mli_sc "(** %s Binary Encoding *)"
-    (Filename.basename cmdline.Cmdline.proto_file_name);
-  F.empty_line mli_sc;
-  generate_for_all_types ocaml_types mli_sc Pb_codegen_encode_binary.gen_sig
-    (Some Pb_codegen_encode_binary.ocamldoc_title);
-  generate_for_all_types ocaml_types mli_sc Pb_codegen_decode_binary.gen_sig
-    (Some Pb_codegen_decode_binary.ocamldoc_title);
-
-  open_files cmdline file_suffix @@ fun (sig_oc, struct_oc) ->
-  output_string struct_oc (F.print ml_sc);
-  output_string sig_oc (F.print mli_sc)
-
-let generate_pp ocaml_types cmdline =
-  let file_suffix = "pp" in
-
-  (* .ml file *)
-  let ml_sc = F.empty_scope () in
-  F.line ml_sc "[@@@ocaml.warning \"-27-30-39\"]";
-  F.empty_line ml_sc;
-
-  generate_for_all_types ocaml_types ml_sc Pb_codegen_pp.gen_struct None;
-
-  (* .mli file *)
-  let mli_sc = F.empty_scope () in
-  F.linep mli_sc "(** %s Pretty Printing *)"
-    (Filename.basename cmdline.Cmdline.proto_file_name);
-  F.empty_line mli_sc;
-  generate_for_all_types ocaml_types mli_sc Pb_codegen_pp.gen_sig
-    (Some Pb_codegen_pp.ocamldoc_title);
-
-  open_files cmdline file_suffix @@ fun (sig_oc, struct_oc) ->
-  output_string struct_oc (F.print ml_sc);
-  output_string sig_oc (F.print mli_sc)
+  let ocaml_mod : CG_all.ocaml_mod =
+    CG_all.codegen ocaml_types ~proto_file_options
+      ~proto_file_name:cmdline.proto_file_name plugins
+  in
 
-let generate_code ocaml_types proto_file_options cmdline : unit =
-  generate_type_and_default ocaml_types proto_file_options cmdline;
-  if !(cmdline.Cmdline.yojson) then generate_yojson ocaml_types cmdline;
-  if !(cmdline.Cmdline.binary) then generate_binary ocaml_types cmdline;
-  if !(cmdline.Cmdline.pp) then generate_pp ocaml_types cmdline;
-  if !(cmdline.Cmdline.bs) then generate_bs ocaml_types cmdline;
+  (* now write the files *)
+  open_files cmdline (fun ~ml ~mli ->
+      F.output ml ocaml_mod.ml;
+      F.output mli ocaml_mod.mli);
   ()
diff --git a/src/ocaml-protoc/ocaml_protoc_generation.mli b/src/ocaml-protoc/ocaml_protoc_generation.mli
index 3e1f7cc..1586458 100644
--- a/src/ocaml-protoc/ocaml_protoc_generation.mli
+++ b/src/ocaml-protoc/ocaml_protoc_generation.mli
@@ -3,4 +3,5 @@
 module Ot = Pb_codegen_ocaml_type
 module Cmdline = Ocaml_protoc_cmdline.Cmdline
 
-val generate_code : Ot.type_ list list -> Pb_option.set -> Cmdline.t -> unit
+val generate_code :
+  Ot.type_ list list -> proto_file_options:Pb_option.set -> Cmdline.t -> unit
diff --git a/src/runtime/dune b/src/runtime/dune
index 82b7aa0..4d317e3 100644
--- a/src/runtime/dune
+++ b/src/runtime/dune
@@ -1,6 +1,7 @@
 (library
  (name pbrt)
  (public_name pbrt)
+ (synopsis "Runtime library for ocaml-protoc")
  ; we need to increase -inline, so that the varint encoder/decoder can
  ; be remembered by the inliner.
  (ocamlopt_flags :standard -inline 100))


Test patch:
diff --git a/src/tests/benchmark/benchmark_single_ml.ml b/src/tests/benchmark/benchmark_single_ml.ml
index 70b6890..6a52e2d 100644
--- a/src/tests/benchmark/benchmark_single_ml.ml
+++ b/src/tests/benchmark/benchmark_single_ml.ml
@@ -1,6 +1,6 @@
 module type T_sig = Ocaml_test_runner.T_sig
 
-module Bench_t = Benchmark_types
+module Bench_t = Benchmark
 
 module Util = struct
   let shuffle_array a n =
diff --git a/src/tests/benchmark/dune b/src/tests/benchmark/dune
index 5e6aa0c..b52ecb9 100644
--- a/src/tests/benchmark/dune
+++ b/src/tests/benchmark/dune
@@ -1,12 +1,11 @@
 (rule
- (targets benchmark_types.ml benchmark_pb.ml benchmark_pp.ml
-   benchmark_types.mli benchmark_pb.mli benchmark_pp.mli)
+ (targets benchmark.ml benchmark.mli)
  (deps benchmark.proto)
  (action
-  (run ocaml-protoc -binary -pp -ml_out ./ %{deps})))
+  (run ocaml-protoc --binary --pp --ml_out ./ %{deps})))
 
 (executable
  (name benchmark_single_ml)
  (modules benchmark_single_ml ocaml_test_types ocaml_test_runner
-   benchmark_types benchmark_pb benchmark_pp)
+   benchmark)
  (libraries pbrt unix))
diff --git a/src/tests/benchmark/ocaml_test_runner.ml b/src/tests/benchmark/ocaml_test_runner.ml
index e7f8909..4a316ba 100644
--- a/src/tests/benchmark/ocaml_test_runner.ml
+++ b/src/tests/benchmark/ocaml_test_runner.ml
@@ -19,16 +19,16 @@ end
 (* T_sig *)
 
 module type Runner_sig = sig
-  val run : Benchmark_types.test_request -> Benchmark_types.test_response
+  val run : Benchmark.test_request -> Benchmark.test_response
   (** [run request] execute the corresponding request and returns 
       the test response 
     *)
 end
 
 module Make (T : T_sig) : Runner_sig = struct
-  let run { Benchmark_types.type_; file_name; test_id } =
+  let run { Benchmark.type_; file_name; test_id } =
     match type_ with
-    | Benchmark_types.Encode difficulty ->
+    | Benchmark.Encode difficulty ->
       let t0 = Unix.gettimeofday () in
 
       let encoder = Pbrt.Encoder.create () in
@@ -45,7 +45,7 @@ module Make (T : T_sig) : Runner_sig = struct
       let t3 = Unix.gettimeofday () in
 
       let encode_data =
-        Benchmark_types.
+        Benchmark.
           {
             creation_time = t1 -. t0;
             encode_time = t2 -. t1;
@@ -53,8 +53,8 @@ module Make (T : T_sig) : Runner_sig = struct
           }
       in
 
-      Benchmark_types.{ difficulty_size; test_id; data = Encode encode_data }
-    | Benchmark_types.Decode ->
+      Benchmark.{ difficulty_size; test_id; data = Encode encode_data }
+    | Benchmark.Decode ->
       let t0 = Unix.gettimeofday () in
       let ic = open_in file_name in
       let len = in_channel_length ic in
@@ -64,15 +64,15 @@ module Make (T : T_sig) : Runner_sig = struct
       let t1 = Unix.gettimeofday () in
       let v = T.decode (Pbrt.Decoder.of_bytes payload) in
       let t2 = Unix.gettimeofday () in
-      Benchmark_types.
+      Benchmark.
         {
           difficulty_size = T.difficulty_size v;
           test_id;
           data =
             Decode
               {
-                Benchmark_types.from_file_time = t1 -. t0;
-                Benchmark_types.decode_time = t2 -. t1;
+                Benchmark.from_file_time = t1 -. t0;
+                Benchmark.decode_time = t2 -. t1;
               };
         }
 end
diff --git a/src/tests/benchmark/ocaml_test_runner.mli b/src/tests/benchmark/ocaml_test_runner.mli
index ece44a4..e1831c1 100644
--- a/src/tests/benchmark/ocaml_test_runner.mli
+++ b/src/tests/benchmark/ocaml_test_runner.mli
@@ -31,7 +31,7 @@ end
 (* T_sig *)
 
 module type Runner_sig = sig
-  val run : Benchmark_types.test_request -> Benchmark_types.test_response
+  val run : Benchmark.test_request -> Benchmark.test_response
   (** [run request] execute the corresponding request and returns 
       the test response 
     *)
diff --git a/src/tests/benchmark/ocaml_test_types.ml b/src/tests/benchmark/ocaml_test_types.ml
index 7ef56e6..95ae0b5 100644
--- a/src/tests/benchmark/ocaml_test_types.ml
+++ b/src/tests/benchmark/ocaml_test_types.ml
@@ -9,7 +9,7 @@ let difficulty_size_of_difficulty = function
   | _ -> failwith "Invalid difficulty"
 
 module Int32List = struct
-  type t = Benchmark_types.int32_list
+  type t = Benchmark.int32_list
 
   let rec make_impl acc = function
     | 0l -> acc
@@ -17,17 +17,16 @@ module Int32List = struct
 
   let make difficulty =
     let difficulty_size = difficulty_size_of_difficulty difficulty in
-    ( Benchmark_types.
-        { int32_list = make_impl [] (Int32.of_int difficulty_size) },
+    ( Benchmark.{ int32_list = make_impl [] (Int32.of_int difficulty_size) },
       difficulty_size )
 
-  let difficulty_size { Benchmark_types.int32_list } = List.length int32_list
-  let encode = Benchmark_pb.encode_int32_list
-  let decode = Benchmark_pb.decode_int32_list
+  let difficulty_size { Benchmark.int32_list } = List.length int32_list
+  let encode = Benchmark.encode_pb_int32_list
+  let decode = Benchmark.decode_pb_int32_list
 end
 
 module IntList = struct
-  type t = Benchmark_types.int_list
+  type t = Benchmark.int_list
 
   let rec make_impl acc = function
     | 0 -> acc
@@ -35,15 +34,15 @@ module IntList = struct
 
   let make difficulty =
     let difficulty_size = difficulty_size_of_difficulty difficulty in
-    Benchmark_types.{ int_list = make_impl [] difficulty_size }, difficulty_size
+    Benchmark.{ int_list = make_impl [] difficulty_size }, difficulty_size
 
-  let difficulty_size { Benchmark_types.int_list } = List.length int_list
-  let encode = Benchmark_pb.encode_int_list
-  let decode = Benchmark_pb.decode_int_list
+  let difficulty_size { Benchmark.int_list } = List.length int_list
+  let encode = Benchmark.encode_pb_int_list
+  let decode = Benchmark.decode_pb_int_list
 end
 
 module IntRepeated = struct
-  type t = Benchmark_types.int_repeated
+  type t = Benchmark.int_repeated
 
   let make_impl n =
     let r = Pbrt.Repeated_field.make 0 in
@@ -57,18 +56,17 @@ module IntRepeated = struct
 
   let make difficulty =
     let difficulty_size = difficulty_size_of_difficulty difficulty in
-    ( Benchmark_types.{ int_repeated = make_impl difficulty_size },
-      difficulty_size )
+    Benchmark.{ int_repeated = make_impl difficulty_size }, difficulty_size
 
-  let difficulty_size { Benchmark_types.int_repeated } =
+  let difficulty_size { Benchmark.int_repeated } =
     Pbrt.Repeated_field.length int_repeated
 
-  let encode = Benchmark_pb.encode_int_repeated
-  let decode = Benchmark_pb.decode_int_repeated
+  let encode = Benchmark.encode_pb_int_repeated
+  let decode = Benchmark.decode_pb_int_repeated
 end
 
 module IntPackedRepeated = struct
-  type t = Benchmark_types.int_packed_repeated
+  type t = Benchmark.int_packed_repeated
 
   let make_impl n =
     let r = Pbrt.Repeated_field.make 0 in
@@ -82,20 +80,18 @@ module IntPackedRepeated = struct
 
   let make difficulty =
     let difficulty_size = difficulty_size_of_difficulty difficulty in
-    let v =
-      Benchmark_types.{ int_packed_repeated = make_impl difficulty_size }
-    in
+    let v = Benchmark.{ int_packed_repeated = make_impl difficulty_size } in
     v, difficulty_size
 
-  let difficulty_size { Benchmark_types.int_packed_repeated = l } =
+  let difficulty_size { Benchmark.int_packed_repeated = l } =
     Pbrt.Repeated_field.length l
 
-  let encode = Benchmark_pb.encode_int_packed_repeated
-  let decode = Benchmark_pb.decode_int_packed_repeated
+  let encode = Benchmark.encode_pb_int_packed_repeated
+  let decode = Benchmark.decode_pb_int_packed_repeated
 end
 
 let get_test = function
-  | Benchmark_types.Int_list -> (module IntList : T_sig)
-  | Benchmark_types.Int_repeated -> (module IntRepeated : T_sig)
-  | Benchmark_types.Int_packed_repeated -> (module IntPackedRepeated : T_sig)
-  | Benchmark_types.Int32_list -> (module Int32List : T_sig)
+  | Benchmark.Int_list -> (module IntList : T_sig)
+  | Benchmark.Int_repeated -> (module IntRepeated : T_sig)
+  | Benchmark.Int_packed_repeated -> (module IntPackedRepeated : T_sig)
+  | Benchmark.Int32_list -> (module Int32List : T_sig)
diff --git a/src/tests/benchmark/ocaml_test_types.mli b/src/tests/benchmark/ocaml_test_types.mli
index e1846e7..728f454 100644
--- a/src/tests/benchmark/ocaml_test_types.mli
+++ b/src/tests/benchmark/ocaml_test_types.mli
@@ -1,6 +1,6 @@
 (** Test type implementation *)
 
-val get_test : Benchmark_types.test_id -> (module Ocaml_test_runner.T_sig)
+val get_test : Benchmark.test_id -> (module Ocaml_test_runner.T_sig)
 (** [get_test test_id] returns the module for the given [test_id] which 
     can then be plugged into the Ocaml_test_runner.Make functor
   *)
diff --git a/src/tests/google_unittest/dune b/src/tests/google_unittest/dune
index 23282ff..6291319 100644
--- a/src/tests/google_unittest/dune
+++ b/src/tests/google_unittest/dune
@@ -3,12 +3,12 @@
   (libraries pbrt))
 
 (rule
-  (targets unittest_types.ml unittest_types.mli)
+  (targets unittest.ml unittest.mli)
   (deps (:file unittest.proto) unittest_import.proto)
-  (action (run ocaml-protoc -ml_out . -I . %{file})))
+  (action (run ocaml-protoc --ml_out . -I . %{file})))
 
 (rule
-  (targets unittest_import_types.ml unittest_import_types.mli)
+  (targets unittest_import.ml unittest_import.mli)
   (deps (:file unittest_import.proto))
-  (action (run ocaml-protoc -ml_out . -I . %{file})))
+  (action (run ocaml-protoc --ml_out . -I . %{file})))
 
diff --git a/src/tests/google_unittest/google_unittest.ml b/src/tests/google_unittest/google_unittest.ml
index ff83674..8e3fd2a 100644
--- a/src/tests/google_unittest/google_unittest.ml
+++ b/src/tests/google_unittest/google_unittest.ml
@@ -1,4 +1,4 @@
-module T = Unittest_types
+module T = Unittest
 
 let () =
   assert (
diff --git a/src/tests/integration-tests/dune b/src/tests/integration-tests/dune
index 2537878..1ac2574 100644
--- a/src/tests/integration-tests/dune
+++ b/src/tests/integration-tests/dune
@@ -5,103 +5,95 @@
  (libraries pbrt unix))
 
 (rule
- (targets test01_types.ml test01_pb.ml test01_pp.ml test01_types.mli
-   test01_pb.mli test01_pp.mli)
+ (targets test01.ml test01.mli)
  (deps test01.proto ../../include/ocaml-protoc/ocamloptions.proto)
  (action
   (run ocaml-protoc -I ./ -I ../../include/ocaml-protoc -I %{env:INCLUDE=/usr/include/}
-    -binary -pp -ocaml_all_types_ppx "deriving show" -ml_out ./ ./
+    --binary --pp --ocaml_all_types_ppx "deriving show" --ml_out ./ ./
     test01.proto)))
 
 (executable
  (name test01_ml)
- (modules test01_ml test01_types test01_pb test01_pp)
+ (modules test01_ml test01)
  (preprocess
   (pps ppx_deriving.show))
  (libraries test_util unix))
 
 (rule
- (targets test02_types.ml test02_pb.ml test02_pp.ml test02_types.mli
-   test02_pb.mli test02_pp.mli)
+ (targets test02.ml test02.mli)
  (deps test02.proto ../../include/ocaml-protoc/ocamloptions.proto)
  (action
-  (run ocaml-protoc -binary -pp -ml_out ./ ./ test02.proto)))
+  (run ocaml-protoc --binary --pp --ml_out ./ ./ test02.proto)))
 
 (executable
  (name test02_ml)
- (modules test02_ml test02_types test02_pb test02_pp)
+ (modules test02_ml test02)
  (libraries test_util unix))
 
 (rule
- (targets test04_types.ml test04_pb.ml test04_pp.ml test04_types.mli
-   test04_pb.mli test04_pp.mli)
+ (targets test04.ml test04.mli)
  (deps test04.proto ../../include/ocaml-protoc/ocamloptions.proto)
  (action
-  (run ocaml-protoc -binary -pp -ml_out ./ ./ test04.proto)))
+  (run ocaml-protoc --binary --pp --ml_out ./ ./ test04.proto)))
 
 (executable
  (name test04_ml)
- (modules test04_ml test04_types test04_pb test04_pp)
+ (modules test04_ml test04)
  (libraries test_util unix))
 
 (rule
- (targets test05_types.ml test05_pb.ml test05_pp.ml test05_types.mli
-   test05_pb.mli test05_pp.mli)
+ (targets test05.ml test05.mli)
  (deps test05.proto ../../include/ocaml-protoc/ocamloptions.proto)
  (action
-  (run ocaml-protoc -binary -pp -ml_out ./ ./ test05.proto)))
+  (run ocaml-protoc --binary --pp --ml_out ./ ./ test05.proto)))
 
 (executable
  (name test05_ml)
- (modules test05_ml test05_types test05_pb test05_pp)
+ (modules test05_ml test05)
  (libraries test_util unix))
 
 (rule
- (targets test06_types.ml test06_pb.ml test06_pp.ml test06_types.mli
-   test06_pb.mli test06_pp.mli)
+ (targets test06.ml test06.mli)
  (deps test06.proto ../../include/ocaml-protoc/ocamloptions.proto)
  (action
-  (run ocaml-protoc -binary -pp -ml_out ./ ./ test06.proto)))
+  (run ocaml-protoc --binary --pp --ml_out ./ ./ test06.proto)))
 
 (executable
  (name test06_ml)
- (modules test06_ml test06_types test06_pb test06_pp)
+ (modules test06_ml test06)
  (libraries test_util unix))
 
 (rule
- (targets test07_types.ml test07_pb.ml test07_pp.ml test07_types.mli
-   test07_pb.mli test07_pp.mli)
+ (targets test07.ml test07.mli)
  (deps test07.proto ../../include/ocaml-protoc/ocamloptions.proto)
  (action
-  (run ocaml-protoc -binary -pp -ml_out ./ ./ test07.proto)))
+  (run ocaml-protoc --binary --pp --ml_out ./ ./ test07.proto)))
 
 (executable
  (name test07_ml)
- (modules test07_ml test07_types test07_pb test07_pp)
+ (modules test07_ml test07)
  (libraries test_util unix))
 
 (rule
- (targets test08_types.ml test08_pb.ml test08_pp.ml test08_types.mli
-   test08_pb.mli test08_pp.mli)
+ (targets test08.ml test08.mli)
  (deps test08.proto ../../include/ocaml-protoc/ocamloptions.proto)
  (action
-  (run ocaml-protoc -binary -pp -ml_out ./ ./ test08.proto)))
+  (run ocaml-protoc --binary --pp --ml_out ./ ./ test08.proto)))
 
 (executable
  (name test08_ml)
- (modules test08_ml test08_types test08_pb test08_pp)
+ (modules test08_ml test08)
  (libraries test_util unix))
 
 (rule
- (targets test09_types.ml test09_pb.ml test09_pp.ml test09_types.mli
-   test09_pb.mli test09_pp.mli)
+ (targets test09.ml test09.mli)
  (deps test09.proto ../../include/ocaml-protoc/ocamloptions.proto)
  (action
-  (run ocaml-protoc -binary -pp -ml_out ./ ./ test09.proto)))
+  (run ocaml-protoc --binary --pp --ml_out ./ ./ test09.proto)))
 
 (library
  (name test09_gen)
- (modules test09_types test09_pb test09_pp)
+ (modules test09)
  (wrapped false)
  (libraries pbrt))
 
@@ -111,245 +103,224 @@
  (libraries test09_gen test_util unix))
 
 (rule
- (targets test10_types.ml test10_pb.ml test10_pp.ml test10_types.mli
-   test10_pb.mli test10_pp.mli)
+ (targets test10.ml test10.mli)
  (deps test10.proto ../../include/ocaml-protoc/ocamloptions.proto)
  (action
-  (run ocaml-protoc -binary -pp -ml_out ./ ./ test10.proto)))
+  (run ocaml-protoc --binary --pp --ml_out ./ ./ test10.proto)))
 
 (executable
  (name test10_ml)
- (modules test10_ml test10_types test10_pb test10_pp)
+ (modules test10_ml test10)
  (libraries test09_gen test_util unix))
 
 (rule
- (targets test11_types.ml test11_pb.ml test11_pp.ml test11_types.mli
-   test11_pb.mli test11_pp.mli)
+ (targets test11.ml test11.mli)
  (deps test11.proto ../../include/ocaml-protoc/ocamloptions.proto)
  (action
-  (run ocaml-protoc -binary -pp -ml_out ./ ./ test11.proto)))
+  (run ocaml-protoc --binary --pp --ml_out ./ ./ test11.proto)))
 
 (executable
  (name test11_ml)
- (modules test11_ml test11_types test11_pb test11_pp)
+ (modules test11_ml test11)
  (libraries test_util unix))
 
 (rule
- (targets test12_types.ml test12_pb.ml test12_pp.ml test12_types.mli
-   test12_pb.mli test12_pp.mli)
+ (targets test12.ml test12.mli)
  (deps test12.proto ../../include/ocaml-protoc/ocamloptions.proto)
  (action
-  (run ocaml-protoc -binary -pp -ml_out ./ ./ test12.proto)))
+  (run ocaml-protoc --binary --pp --ml_out ./ ./ test12.proto)))
 
 (executable
  (name test12_ml)
- (modules test12_ml test12_types test12_pb test12_pp)
+ (modules test12_ml test12)
  (libraries test_util unix))
 
 (rule
- (targets test13_types.ml test13_pb.ml test13_pp.ml test13_types.mli
-   test13_pb.mli test13_pp.mli)
+ (targets test13.ml test13.mli)
  (deps test13.proto ../../include/ocaml-protoc/ocamloptions.proto)
  (action
-  (run ocaml-protoc -binary -pp -ml_out ./ ./ test13.proto)))
+  (run ocaml-protoc --binary --pp --ml_out ./ ./ test13.proto)))
 
 (executable
  (name test13_ml)
- (modules test13_ml test13_types test13_pb test13_pp)
+ (modules test13_ml test13)
  (libraries test_util unix))
 
 (rule
- (targets test14_types.ml test14_pb.ml test14_pp.ml test14_types.mli
-   test14_pb.mli test14_pp.mli)
+ (targets test14.ml test14.mli)
  (deps test14.proto ../../include/ocaml-protoc/ocamloptions.proto)
  (action
-  (run ocaml-protoc -binary -pp -ml_out ./ ./ test14.proto)))
+  (run ocaml-protoc --binary --pp --ml_out ./ ./ test14.proto)))
 
 (executable
  (name test14_ml)
- (modules test14_ml test14_types test14_pb test14_pp)
+ (modules test14_ml test14)
  (libraries test_util unix))
 
 (rule
- (targets test15_types.ml test15_pb.ml test15_pp.ml test15_types.mli
-   test15_pb.mli test15_pp.mli)
+ (targets test15.ml test15.mli)
  (deps test15.proto ../../include/ocaml-protoc/ocamloptions.proto)
  (action
   (run ocaml-protoc -I ../../include/ocaml-protoc -I %{env:INCLUDE=/usr/include/}
-    -binary -pp -ml_out ./ ./ test15.proto)))
+    --binary --pp --ml_out ./ ./ test15.proto)))
 
 (executable
  (name test15_ml)
- (modules test15_ml test15_types test15_pb test15_pp)
+ (modules test15_ml test15)
  (libraries test_util unix))
 
 (rule
- (targets test16_types.ml test16_pb.ml test16_pp.ml test16_types.mli
-   test16_pb.mli test16_pp.mli)
+ (targets test16.ml test16.mli)
  (deps test16.proto ../../include/ocaml-protoc/ocamloptions.proto)
  (action
   (run ocaml-protoc -I ../../include/ocaml-protoc -I %{env:INCLUDE=/usr/include/}
-    -binary -pp -ml_out ./ ./ test16.proto)))
+    --binary --pp --ml_out ./ ./ test16.proto)))
 
 (executable
  (name test16_ml)
- (modules test16_ml test16_types test16_pb test16_pp)
+ (modules test16_ml test16)
  (libraries test_util unix))
 
 (rule
- (targets test17_types.ml test17_pb.ml test17_pp.ml test17_types.mli
-   test17_pb.mli test17_pp.mli)
+ (targets test17.ml test17.mli)
  (deps test17.proto ../../include/ocaml-protoc/ocamloptions.proto)
  (action
   (run ocaml-protoc -I ../../include/ocaml-protoc -I %{env:INCLUDE=/usr/include/}
-    -binary -pp -ml_out ./ ./ test17.proto)))
+    --binary --pp --ml_out ./ ./ test17.proto)))
 
 (executable
  (name test17_ml)
- (modules test17_ml test17_types test17_pb test17_pp)
+ (modules test17_ml test17)
  (libraries test_util unix))
 
 (rule
- (targets test18_types.ml test18_pb.ml test18_pp.ml test18_types.mli
-   test18_pb.mli test18_pp.mli)
+ (targets test18.ml test18.mli)
  (deps test18.proto ../../include/ocaml-protoc/ocamloptions.proto)
  (action
   (run ocaml-protoc -I ../../include/ocaml-protoc -I %{env:INCLUDE=/usr/include/}
-    -binary -pp -ml_out ./ ./ test18.proto)))
+    --binary --pp --ml_out ./ ./ test18.proto)))
 
 (executable
  (name test18_ml)
- (modules test18_ml test18_types test18_pb test18_pp)
+ (modules test18_ml test18)
  (libraries test_util unix))
 
 (rule
- (targets test19_types.ml test19_pb.ml test19_pp.ml test19_types.mli
-   test19_pb.mli test19_pp.mli)
+ (targets test19.ml test19.mli)
  (deps test19.proto ../../include/ocaml-protoc/ocamloptions.proto)
  (action
   (run ocaml-protoc -I ../../include/ocaml-protoc -I %{env:INCLUDE=/usr/include/}
-    -binary -pp -ml_out ./ ./ test19.proto)))
+    --binary --pp --ml_out ./ ./ test19.proto)))
 
 (executable
  (name test19_ml)
- (modules test19_ml test19_types test19_pb test19_pp)
+ (modules test19_ml test19)
  (libraries test_util unix))
 
 (rule
- (targets test20_types.ml test20_pb.ml test20_pp.ml test20_types.mli
-   test20_pb.mli test20_pp.mli)
+ (targets test20.ml test20.mli)
  (deps test20.proto ../../include/ocaml-protoc/ocamloptions.proto)
  (action
   (run ocaml-protoc -I ../../include/ocaml-protoc -I %{env:INCLUDE=/usr/include/}
-    -binary -pp -ml_out ./ ./ test20.proto)))
+    --binary --pp --ml_out ./ ./ test20.proto)))
 
 (executable
  (name test20_ml)
- (modules test20_ml test20_types test20_pb test20_pp)
+ (modules test20_ml test20)
  (preprocess
   (pps ppx_deriving.show))
  (libraries test_util unix))
 
 (rule
- (targets test21_types.ml test21_pb.ml test21_pp.ml test21_types.mli
-   test21_pb.mli test21_pp.mli)
+ (targets test21.ml test21.mli)
  (deps test21.proto ../../include/ocaml-protoc/ocamloptions.proto)
  (action
   (run ocaml-protoc -I ../../include/ocaml-protoc -I %{env:INCLUDE=/usr/include/}
-    -binary -pp -ml_out ./ ./ test21.proto)))
+    --binary --pp --ml_out ./ ./ test21.proto)))
 
 (executable
  (name test21_ml)
- (modules test21_ml test21_types test21_pb test21_pp)
+ (modules test21_ml test21)
  (preprocess
   (pps ppx_deriving.show))
  (libraries test_util unix))
 
 (rule
- (targets test22_types.ml test22_pb.ml test22_pp.ml test22_types.mli
-   test22_pb.mli test22_pp.mli)
+ (targets test22.ml test22.mli)
  (deps test22.proto ../../include/ocaml-protoc/ocamloptions.proto)
  (action
-  (run ocaml-protoc -binary -pp -ml_out ./ ./ test22.proto)))
+  (run ocaml-protoc --binary --pp --ml_out ./ ./ test22.proto)))
 
 (executable
  (name test22_ml)
- (modules test22_ml test22_types test22_pb test22_pp)
+ (modules test22_ml test22)
  (libraries test_util unix))
 
 (rule
- (targets test23_types.ml test23_pb.ml test23_pp.ml test23_types.mli
-   test23_pb.mli test23_pp.mli)
+ (targets test23.ml test23.mli)
  (deps test23.proto ../../include/ocaml-protoc/ocamloptions.proto)
  (action
-  (run ocaml-protoc -binary -pp -ml_out ./ ./ test23.proto)))
+  (run ocaml-protoc --binary --pp --ml_out ./ ./ test23.proto)))
 
 (executable
  (name test23_ml)
- (modules test23_ml test23_types test23_pb test23_pp)
+ (modules test23_ml test23)
  (libraries test_util unix))
 
 (rule
- (targets test24_types.ml test24_pb.ml test24_pp.ml test24_types.mli
-   test24_pb.mli test24_pp.mli)
+ (targets test24.ml test24.mli)
  (deps test24.proto ../../include/ocaml-protoc/ocamloptions.proto)
  (action
-  (run ocaml-protoc -binary -pp -ml_out ./ ./ test24.proto)))
+  (run ocaml-protoc --binary --pp --ml_out ./ ./ test24.proto)))
 
 (executable
  (name test24_ml)
- (modules test24_ml test24_types test24_pb test24_pp)
+ (modules test24_ml test24)
  (libraries test_util unix))
 
 (rule
- (targets test25_types.ml test25_pb.ml test25_pp.ml test25_types.mli
-   test25_pb.mli test25_pp.mli)
+ (targets test25.ml test25.mli)
  (deps test25.proto ../../include/ocaml-protoc/ocamloptions.proto)
  (action
-  (run ocaml-protoc -binary -pp -ml_out ./ ./ test25.proto)))
+  (run ocaml-protoc --binary --pp --ml_out ./ ./ test25.proto)))
 
 (executable
  (name test25_ml)
- (modules test25_ml test25_types test25_pb test25_pp)
+ (modules test25_ml test25)
  (libraries test_util unix))
 
 (rule
- (targets test26_types.ml test26_pb.ml test26_pp.ml test26_types.mli
-   test26_pb.mli test26_pp.mli)
+ (targets test26.ml test26.mli)
  (deps test26.proto ../../include/ocaml-protoc/ocamloptions.proto)
  (action
-  (run ocaml-protoc -I %{env:INCLUDE=/usr/include/} -binary -pp -ml_out ./ ./
+  (run ocaml-protoc -I %{env:INCLUDE=/usr/include/} --binary --pp --ml_out ./ ./
     test26.proto)))
 
 (executable
  (name test26_ml)
- (modules test26_ml test26_types test26_pb test26_pp)
+ (modules test26_ml test26)
  (libraries test_util unix))
 
 (rule
- (targets test27_types.ml test27_pb.ml test27_pp.ml test27_types.mli
-   test27_pb.mli test27_pp.mli)
+ (targets test27.ml test27.mli)
  (deps test27.proto ../../include/ocaml-protoc/ocamloptions.proto)
  (action
-  (run ocaml-protoc -binary -pp -ml_out ./ test27.proto)))
+  (run ocaml-protoc --binary --pp --ml_out ./ test27.proto)))
 
 (executable
  (name test27_ml)
- (modules test27_ml test27_types test27_pb test27_pp)
+ (modules test27_ml test27)
  (libraries test_util unix))
 
 (rule
- (targets test_proto3_optional_types.ml test_proto3_optional_pb.ml
-   test_proto3_optional_pp.ml test_proto3_optional_types.mli
-   test_proto3_optional_pb.mli test_proto3_optional_pp.mli)
+ (targets test_proto3_optional.ml test_proto3_optional.mli)
  (deps
   (:proto test_proto3_optional.proto)
   ../../include/ocaml-protoc/ocamloptions.proto)
  (action
-  (run ocaml-protoc -binary -pp -ml_out ./ %{proto})))
+  (run ocaml-protoc --binary --pp --ml_out ./ %{proto})))
 
 (executable
  (name test_proto3_optional)
  (libraries pbrt)
- (modules test_proto3_optional_types test_proto3_optional
-   test_proto3_optional_pb test_proto3_optional_pp))
+ (modules test_proto3_optional_ml test_proto3_optional))
diff --git a/src/tests/integration-tests/test01_ml.ml b/src/tests/integration-tests/test01_ml.ml
index bd285ac..4fc73a7 100644
--- a/src/tests/integration-tests/test01_ml.ml
+++ b/src/tests/integration-tests/test01_ml.ml
@@ -1,6 +1,4 @@
-module T = Test01_types
-module Pb = Test01_pb
-module Pp = Test01_pp
+module T = Test01
 
 let decode_ref_data () =
   T.
@@ -41,10 +39,10 @@ let () =
 
   match mode with
   | Test_util.Decode ->
-    Test_util.decode "test01.c2ml.data" Pb.decode_couple Pp.pp_couple
+    Test_util.decode "test01.c2ml.data" T.decode_pb_couple T.pp_couple
       (decode_ref_data ())
   | Test_util.Encode ->
-    Test_util.encode "test01.ml2c.data" Pb.encode_couple (decode_ref_data ())
+    Test_util.encode "test01.ml2c.data" T.encode_pb_couple (decode_ref_data ())
 
 let () =
   let expected_default_person =
diff --git a/src/tests/integration-tests/test02_ml.ml b/src/tests/integration-tests/test02_ml.ml
index 78206b6..6a32a9e 100644
--- a/src/tests/integration-tests/test02_ml.ml
+++ b/src/tests/integration-tests/test02_ml.ml
@@ -1,6 +1,4 @@
-module T = Test02_types
-module Pp = Test02_pp
-module Pb = Test02_pb
+module T = Test02
 
 let decode_ref_data () =
   {
@@ -26,8 +24,8 @@ let () =
 
   match mode with
   | Test_util.Decode ->
-    Test_util.decode "test02.c2ml.data" Pb.decode_all_basics_types
-      Pp.pp_all_basics_types (decode_ref_data ())
+    Test_util.decode "test02.c2ml.data" T.decode_pb_all_basics_types
+      T.pp_all_basics_types (decode_ref_data ())
   | Test_util.Encode ->
-    Test_util.encode "test02.ml2c.data" Pb.encode_all_basics_types
+    Test_util.encode "test02.ml2c.data" T.encode_pb_all_basics_types
       (decode_ref_data ())
diff --git a/src/tests/integration-tests/test04_ml.ml b/src/tests/integration-tests/test04_ml.ml
index 4bb8d29..1e2fedb 100644
--- a/src/tests/integration-tests/test04_ml.ml
+++ b/src/tests/integration-tests/test04_ml.ml
@@ -1,6 +1,4 @@
-module T = Test04_types
-module Pb = Test04_pb
-module Pp = Test04_pp
+module T = Test04
 
 let decode_ref_data () = { T.j = 456l }
 
@@ -9,7 +7,7 @@ let () =
 
   match mode with
   | Test_util.Decode ->
-    Test_util.decode "test03.c2ml.data" Pb.decode_test Pp.pp_test
+    Test_util.decode "test03.c2ml.data" T.decode_pb_test T.pp_test
       (decode_ref_data ())
   | Test_util.Encode ->
-    Test_util.encode "test04.ml2c.data" Pb.encode_test (decode_ref_data ())
+    Test_util.encode "test04.ml2c.data" T.encode_pb_test (decode_ref_data ())
diff --git a/src/tests/integration-tests/test05_ml.ml b/src/tests/integration-tests/test05_ml.ml
index 81f199b..3757e0a 100644
--- a/src/tests/integration-tests/test05_ml.ml
+++ b/src/tests/integration-tests/test05_ml.ml
@@ -1,6 +1,4 @@
-module T = Test05_types
-module Pb = Test05_pb
-module Pp = Test05_pp
+module T = Test05
 
 let max_int = 107374182l
 let min_int = -107374182l
@@ -26,7 +24,8 @@ let () =
   | Test_util.Decode ->
     let ref_data = decode_ref_data () in
     Printf.printf "List size : %i\n%!" (List.length @@ ref_data.T.l);
-    Test_util.decode ~noprint:() "test05.c2ml.data" Pb.decode_int_list
-      Pp.pp_int_list ref_data
+    Test_util.decode ~noprint:() "test05.c2ml.data" T.decode_pb_int_list
+      T.pp_int_list ref_data
   | Test_util.Encode ->
-    Test_util.encode "test05.ml2c.data" Pb.encode_int_list (decode_ref_data ())
+    Test_util.encode "test05.ml2c.data" T.encode_pb_int_list
+      (decode_ref_data ())
diff --git a/src/tests/integration-tests/test06_ml.ml b/src/tests/integration-tests/test06_ml.ml
index 2a39ef4..9ad88f3 100644
--- a/src/tests/integration-tests/test06_ml.ml
+++ b/src/tests/integration-tests/test06_ml.ml
@@ -1,6 +1,4 @@
-module T = Test06_types
-module Pb = Test06_pb
-module Pp = Test06_pp
+module T = Test06
 
 let decode_ref_data () = { T.teste_field = T.Test_e_value1 }
 
@@ -9,7 +7,7 @@ let () =
 
   match mode with
   | Test_util.Decode ->
-    Test_util.decode "test06.c2ml.data" Pb.decode_test_m Pp.pp_test_m
+    Test_util.decode "test06.c2ml.data" T.decode_pb_test_m T.pp_test_m
       (decode_ref_data ())
   | Test_util.Encode ->
-    Test_util.encode "test06.ml2c.data" Pb.encode_test_m (decode_ref_data ())
+    Test_util.encode "test06.ml2c.data" T.encode_pb_test_m (decode_ref_data ())
diff --git a/src/tests/integration-tests/test07_ml.ml b/src/tests/integration-tests/test07_ml.ml
index a2afe37..ebd845d 100644
--- a/src/tests/integration-tests/test07_ml.ml
+++ b/src/tests/integration-tests/test07_ml.ml
@@ -1,6 +1,4 @@
-module T = Test07_types
-module Pb = Test07_pb
-module Pp = Test07_pp
+module T = Test07
 
 let decode_ref_data () =
   {
@@ -14,7 +12,7 @@ let () =
 
   match mode with
   | Test_util.Decode ->
-    Test_util.decode "test07.c2ml.data" Pb.decode_node Pp.pp_node
+    Test_util.decode "test07.c2ml.data" T.decode_pb_node T.pp_node
       (decode_ref_data ())
   | Test_util.Encode ->
-    Test_util.encode "test07.ml2c.data" Pb.encode_node (decode_ref_data ())
+    Test_util.encode "test07.ml2c.data" T.encode_pb_node (decode_ref_data ())
diff --git a/src/tests/integration-tests/test08_ml.ml b/src/tests/integration-tests/test08_ml.ml
index 6e6a3dc..c4b4f19 100644
--- a/src/tests/integration-tests/test08_ml.ml
+++ b/src/tests/integration-tests/test08_ml.ml
@@ -1,6 +1,4 @@
-module T = Test08_types
-module Pb = Test08_pb
-module Pp = Test08_pp
+module T = Test08
 
 let decode_ref_data () =
   T.Node
@@ -15,7 +13,7 @@ let () =
 
   match mode with
   | Test_util.Decode ->
-    Test_util.decode "test08.c2ml.data" Pb.decode_tree Pp.pp_tree
+    Test_util.decode "test08.c2ml.data" T.decode_pb_tree T.pp_tree
       (decode_ref_data ())
   | Test_util.Encode ->
-    Test_util.encode "test08.ml2c.data" Pb.encode_tree (decode_ref_data ())
+    Test_util.encode "test08.ml2c.data" T.encode_pb_tree (decode_ref_data ())
diff --git a/src/tests/integration-tests/test09_ml.ml b/src/tests/integration-tests/test09_ml.ml
index 965fd97..f3d90bd 100644
--- a/src/tests/integration-tests/test09_ml.ml
+++ b/src/tests/integration-tests/test09_ml.ml
@@ -1,15 +1,13 @@
-module T = Test09_types
-module Pb = Test09_pb
-module Pp = Test09_pp
+module T = Test09
 
-let decode_ref_data () = { T.int32_f1 = 1l }
+let decode_pb_ref_data () = { T.int32_f1 = 1l }
 
 let () =
   let mode = Test_util.parse_args () in
 
   match mode with
   | Test_util.Decode ->
-    Test_util.decode "test09.c2ml.data" Pb.decode_m09 Pp.pp_m09
-      (decode_ref_data ())
+    Test_util.decode "test09.c2ml.data" T.decode_pb_m09 T.pp_m09
+      (decode_pb_ref_data ())
   | Test_util.Encode ->
-    Test_util.encode "test09.ml2c.data" Pb.encode_m09 (decode_ref_data ())
+    Test_util.encode "test09.ml2c.data" T.encode_pb_m09 (decode_pb_ref_data ())
diff --git a/src/tests/integration-tests/test10_ml.ml b/src/tests/integration-tests/test10_ml.ml
index 33604b9..5533685 100644
--- a/src/tests/integration-tests/test10_ml.ml
+++ b/src/tests/integration-tests/test10_ml.ml
@@ -1,15 +1,13 @@
-module T = Test10_types
-module Pb = Test10_pb
-module Pp = Test10_pp
+module T = Test10
 
-let decode_ref_data () = { T.m10_f1 = { Test09_types.int32_f1 = 1l } }
+let decode_pb_ref_data () = { T.m10_f1 = { Test09.int32_f1 = 1l } }
 
 let () =
   let mode = Test_util.parse_args () in
 
   match mode with
   | Test_util.Decode ->
-    Test_util.decode "test10.c2ml.data" Pb.decode_m10 Pp.pp_m10
-      (decode_ref_data ())
+    Test_util.decode "test10.c2ml.data" T.decode_pb_m10 T.pp_m10
+      (decode_pb_ref_data ())
   | Test_util.Encode ->
-    Test_util.encode "test10.ml2c.data" Pb.encode_m10 (decode_ref_data ())
+    Test_util.encode "test10.ml2c.data" T.encode_pb_m10 (decode_pb_ref_data ())
diff --git a/src/tests/integration-tests/test11_ml.ml b/src/tests/integration-tests/test11_ml.ml
index 1de62d5..1ea47ee 100644
--- a/src/tests/integration-tests/test11_ml.ml
+++ b/src/tests/integration-tests/test11_ml.ml
@@ -1,6 +1,4 @@
-module T = Test11_types
-module Pb = Test11_pb
-module Pp = Test11_pp
+module T = Test11
 
 let create_ref_data n =
   let rec loop l = function
@@ -43,8 +41,8 @@ let () =
 
   match mode with
   | Test_util.Decode ->
-    Test_util.decode "test11.c2ml.data" ~notest:() Pb.decode_couples
-      Pp.pp_couples (create_ref_data 100000)
+    Test_util.decode "test11.c2ml.data" ~notest:() T.decode_pb_couples
+      T.pp_couples (create_ref_data 100000)
   | Test_util.Encode ->
-    Test_util.encode "test11.ml2c.data" Pb.encode_couples
+    Test_util.encode "test11.ml2c.data" T.encode_pb_couples
       (create_ref_data 100000)
diff --git a/src/tests/integration-tests/test12_ml.ml b/src/tests/integration-tests/test12_ml.ml
index 3871003..480c2dc 100644
--- a/src/tests/integration-tests/test12_ml.ml
+++ b/src/tests/integration-tests/test12_ml.ml
@@ -1,14 +1,13 @@
-module T = Test12_types
-module Pb = Test12_pb
-module Pp = Test12_pp
+module T = Test12
 
-let decode_ref_data () = T.{ to_ = 1l; type_ = 1l }
+let decode_pb_ref_data () = T.{ to_ = 1l; type_ = 1l }
 
 let () =
   let mode = Test_util.parse_args () in
 
   match mode with
   | Test_util.Decode ->
-    Test_util.decode "test12.c2ml.data" Pb.decode_m Pp.pp_m (decode_ref_data ())
+    Test_util.decode "test12.c2ml.data" T.decode_pb_m T.pp_m
+      (decode_pb_ref_data ())
   | Test_util.Encode ->
-    Test_util.encode "test12.ml2c.data" Pb.encode_m (decode_ref_data ())
+    Test_util.encode "test12.ml2c.data" T.encode_pb_m (decode_pb_ref_data ())
diff --git a/src/tests/integration-tests/test13_ml.ml b/src/tests/integration-tests/test13_ml.ml
index 94b780a..12bb5c3 100644
--- a/src/tests/integration-tests/test13_ml.ml
+++ b/src/tests/integration-tests/test13_ml.ml
@@ -1,8 +1,6 @@
-module T = Test13_types
-module Pb = Test13_pb
-module Pp = Test13_pp
+module T = Test13
 
-let decode_ref_data () =
+let decode_pb_ref_data () =
   T.
     {
       p2 = { empty = (); sub = Sub_empty };
@@ -14,6 +12,7 @@ let mode = Test_util.parse_args ()
 let () =
   match mode with
   | Test_util.Decode ->
-    Test_util.decode "test13.c2ml.data" Pb.decode_t Pp.pp_t (decode_ref_data ())
+    Test_util.decode "test13.c2ml.data" T.decode_pb_t T.pp_t
+      (decode_pb_ref_data ())
   | Test_util.Encode ->
-    Test_util.encode "test13.ml2c.data" Pb.encode_t (decode_ref_data ())
+    Test_util.encode "test13.ml2c.data" T.encode_pb_t (decode_pb_ref_data ())
diff --git a/src/tests/integration-tests/test14_ml.ml b/src/tests/integration-tests/test14_ml.ml
index d9ec506..4e2e7b6 100644
--- a/src/tests/integration-tests/test14_ml.ml
+++ b/src/tests/integration-tests/test14_ml.ml
@@ -1,8 +1,6 @@
-module T = Test14_types
-module Pb = Test14_pb
-module Pp = Test14_pp
+module T = Test14
 
-let decode_ref_data () =
+let decode_pb_ref_data () =
   T.
     {
       da = { aa = [ 1l; 2l; 3l; 4l ] };
@@ -19,6 +17,7 @@ let mode = Test_util.parse_args ()
 let () =
   match mode with
   | Test_util.Decode ->
-    Test_util.decode "test14.c2ml.data" Pb.decode_d Pp.pp_d (decode_ref_data ())
+    Test_util.decode "test14.c2ml.data" T.decode_pb_d T.pp_d
+      (decode_pb_ref_data ())
   | Test_util.Encode ->
-    Test_util.encode "test14.ml2c.data" Pb.encode_d (decode_ref_data ())
+    Test_util.encode "test14.ml2c.data" T.encode_pb_d (decode_pb_ref_data ())
diff --git a/src/tests/integration-tests/test15_ml.ml b/src/tests/integration-tests/test15_ml.ml
index 4984b3a..73096d7 100644
--- a/src/tests/integration-tests/test15_ml.ml
+++ b/src/tests/integration-tests/test15_ml.ml
@@ -1,6 +1,4 @@
-module T = Test15_types
-module Pb = Test15_pb
-module Pp = Test15_pp
+module T = Test15
 
 (** [fill_0_to_n n] create and fill a Pbrt.Repeated_field 
     with int32 values from [0; n].
@@ -16,7 +14,7 @@ let fill_0_to_n n =
   in
   loop (-n)
 
-let decode_ref_data () =
+let decode_pb_ref_data () =
   T.
     {
       m1_l =
@@ -28,7 +26,7 @@ let mode = Test_util.parse_args ()
 let () =
   match mode with
   | Test_util.Decode ->
-    Test_util.decode "test15.c2ml.data" Pb.decode_m2 Pp.pp_m2
-      (decode_ref_data ())
+    Test_util.decode "test15.c2ml.data" T.decode_pb_m2 T.pp_m2
+      (decode_pb_ref_data ())
   | Test_util.Encode ->
-    Test_util.encode "test15.ml2c.data" Pb.encode_m2 (decode_ref_data ())
+    Test_util.encode "test15.ml2c.data" T.encode_pb_m2 (decode_pb_ref_data ())
diff --git a/src/tests/integration-tests/test16_ml.ml b/src/tests/integration-tests/test16_ml.ml
index 048cc30..90e4937 100644
--- a/src/tests/integration-tests/test16_ml.ml
+++ b/src/tests/integration-tests/test16_ml.ml
@@ -1,13 +1,11 @@
-module T = Test16_types
-module Pb = Test16_pb
-module Pp = Test16_pp
+module T = Test16
 
 let of_list e0 l =
   let f = Pbrt.Repeated_field.make e0 in
   List.iter (fun x -> Pbrt.Repeated_field.add x f) l;
   f
 
-let decode_ref_data () =
+let decode_pb_ref_data () =
   let v = T.{ f1 = 1l; f2 = 2l; f3 = 3l; f4 = of_list 0l [ 1l; 2l ] } in
   v.T.f3 <- 4l;
   v.T.f4 <- of_list 0l [ 3l; 4l ];
@@ -18,6 +16,7 @@ let mode = Test_util.parse_args ()
 let () =
   match mode with
   | Test_util.Decode ->
-    Test_util.decode "test16.c2ml.data" Pb.decode_m Pp.pp_m (decode_ref_data ())
+    Test_util.decode "test16.c2ml.data" T.decode_pb_m T.pp_m
+      (decode_pb_ref_data ())
   | Test_util.Encode ->
-    Test_util.encode "test16.ml2c.data" Pb.encode_m (decode_ref_data ())
+    Test_util.encode "test16.ml2c.data" T.encode_pb_m (decode_pb_ref_data ())
diff --git a/src/tests/integration-tests/test17_ml.ml b/src/tests/integration-tests/test17_ml.ml
index 983d584..77c1b8a 100644
--- a/src/tests/integration-tests/test17_ml.ml
+++ b/src/tests/integration-tests/test17_ml.ml
@@ -1,13 +1,14 @@
-module T = Test17_types
-module Pb = Test17_pb
-module Pp = Test17_pp
+module T = Test17
+
+let decode_pb_ref_data () =
+  T.{ m1 = { i1 = 1; i2 = 2 }; m2 = 1; m3 = 1; o = M4 4 }
 
-let decode_ref_data () = T.{ m1 = { i1 = 1; i2 = 2 }; m2 = 1; m3 = 1; o = M4 4 }
 let mode = Test_util.parse_args ()
 
 let () =
   match mode with
   | Test_util.Decode ->
-    Test_util.decode "test17.c2ml.data" Pb.decode_m Pp.pp_m (decode_ref_data ())
+    Test_util.decode "test17.c2ml.data" T.decode_pb_m T.pp_m
+      (decode_pb_ref_data ())
   | Test_util.Encode ->
-    Test_util.encode "test17.ml2c.data" Pb.encode_m (decode_ref_data ())
+    Test_util.encode "test17.ml2c.data" T.encode_pb_m (decode_pb_ref_data ())
diff --git a/src/tests/integration-tests/test18_ml.ml b/src/tests/integration-tests/test18_ml.ml
index c5a4d45..4f806ea 100644
--- a/src/tests/integration-tests/test18_ml.ml
+++ b/src/tests/integration-tests/test18_ml.ml
@@ -1,8 +1,6 @@
-module T = Test18_types
-module Pb = Test18_pb
-module Pp = Test18_pp
+module T = Test18
 
-let decode_ref_data () =
+let decode_pb_ref_data () =
   T.
     {
       string_to_string = [ "one", "two"; "three", "four" ];
@@ -31,8 +29,8 @@ let sort_by_int32_key l =
   let cmp (lhs, _) (rhs, _) = cmp lhs rhs in
   List.stable_sort cmp l
 
-let decode_maps decoder =
-  let maps = Pb.decode_maps decoder in
+let decode_pb_maps decoder =
+  let maps = T.decode_pb_maps decoder in
   let {
     T.string_to_string;
     T.string_to_int;
@@ -55,7 +53,7 @@ let decode_maps decoder =
 let () =
   match mode with
   | Test_util.Decode ->
-    Test_util.decode "test18.c2ml.data" decode_maps Pp.pp_maps
-      (decode_ref_data ())
+    Test_util.decode "test18.c2ml.data" decode_pb_maps T.pp_maps
+      (decode_pb_ref_data ())
   | Test_util.Encode ->
-    Test_util.encode "test18.ml2c.data" Pb.encode_maps (decode_ref_data ())
+    Test_util.encode "test18.ml2c.data" T.encode_pb_maps (decode_pb_ref_data ())
diff --git a/src/tests/integration-tests/test19_ml.ml b/src/tests/integration-tests/test19_ml.ml
index b99c016..b5594a9 100644
--- a/src/tests/integration-tests/test19_ml.ml
+++ b/src/tests/integration-tests/test19_ml.ml
@@ -1,8 +1,6 @@
-module T = Test19_types
-module Pb = Test19_pb
-module Pp = Test19_pp
+module T = Test19
 
-let decode_ref_data () =
+let decode_pb_ref_data () =
   let hash = Hashtbl.create 3 in
   Hashtbl.add hash "one" "1";
   Hashtbl.add hash "two" "2";
@@ -25,8 +23,8 @@ let () =
   | Test_util.Decode ->
     let buffer, _ = Test_util.get_binary_file_content "test19.c2ml.data" in
     let decoder = Pbrt.Decoder.of_bytes buffer in
-    let hash = Pb.decode_hash decoder in
+    let hash = T.decode_pb_hash decoder in
     let () = test_contain_data hash in
     print_endline "ML: -- Good --"
   | Test_util.Encode ->
-    Test_util.encode "test19.ml2c.data" Pb.encode_hash (decode_ref_data ())
+    Test_util.encode "test19.ml2c.data" T.encode_pb_hash (decode_pb_ref_data ())
diff --git a/src/tests/integration-tests/test20_ml.ml b/src/tests/integration-tests/test20_ml.ml
index 2157f19..9f3eaf5 100644
--- a/src/tests/integration-tests/test20_ml.ml
+++ b/src/tests/integration-tests/test20_ml.ml
@@ -1,15 +1,16 @@
-module T = Test20_types
-module Pb = Test20_pb
-module Pp = Test20_pp
+module T = Test20
 
-let decode_ref_data () = { T.f1 = 1l; f2 = { T.sub_f1 = 2l }; f3 = T.Eone }
-let () = Printf.printf "Show is working: %s\n" @@ T.show_m (decode_ref_data ())
+let decode_pb_ref_data () = { T.f1 = 1l; f2 = { T.sub_f1 = 2l }; f3 = T.Eone }
+
+let () =
+  Printf.printf "Show is working: %s\n" @@ T.show_m (decode_pb_ref_data ())
 
 let () =
   let mode = Test_util.parse_args () in
 
   match mode with
   | Test_util.Decode ->
-    Test_util.decode "test20.c2ml.data" Pb.decode_m Pp.pp_m (decode_ref_data ())
+    Test_util.decode "test20.c2ml.data" T.decode_pb_m T.pp_m
+      (decode_pb_ref_data ())
   | Test_util.Encode ->
-    Test_util.encode "test20.ml2c.data" Pb.encode_m (decode_ref_data ())
+    Test_util.encode "test20.ml2c.data" T.encode_pb_m (decode_pb_ref_data ())
diff --git a/src/tests/integration-tests/test21_ml.ml b/src/tests/integration-tests/test21_ml.ml
index 90e6584..cb625f4 100644
--- a/src/tests/integration-tests/test21_ml.ml
+++ b/src/tests/integration-tests/test21_ml.ml
@@ -1,15 +1,16 @@
-module T = Test21_types
-module Pb = Test21_pb
-module Pp = Test21_pp
+module T = Test21
 
-let decode_ref_data () = { T.f1 = 1l; f2 = { T.sub_f1 = 2l }; f3 = T.Eone }
-let () = Printf.printf "Show is working: %s\n" @@ T.show_m (decode_ref_data ())
+let decode_pb_ref_data () = { T.f1 = 1l; f2 = { T.sub_f1 = 2l }; f3 = T.Eone }
+
+let () =
+  Printf.printf "Show is working: %s\n" @@ T.show_m (decode_pb_ref_data ())
 
 let () =
   let mode = Test_util.parse_args () in
 
   match mode with
   | Test_util.Decode ->
-    Test_util.decode "test21.c2ml.data" Pb.decode_m Pp.pp_m (decode_ref_data ())
+    Test_util.decode "test21.c2ml.data" T.decode_pb_m T.pp_m
+      (decode_pb_ref_data ())
   | Test_util.Encode ->
-    Test_util.encode "test21.ml2c.data" Pb.encode_m (decode_ref_data ())
+    Test_util.encode "test21.ml2c.data" T.encode_pb_m (decode_pb_ref_data ())
diff --git a/src/tests/integration-tests/test22_ml.ml b/src/tests/integration-tests/test22_ml.ml
index 9a02ede..1910647 100644
--- a/src/tests/integration-tests/test22_ml.ml
+++ b/src/tests/integration-tests/test22_ml.ml
@@ -1,8 +1,6 @@
-module T = Test22_types
-module Pb = Test22_pb
-module Pp = Test22_pp
+module T = Test22
 
-let decode_ref_data () =
+let decode_pb_ref_data () =
   T.
     {
       p1 =
@@ -40,10 +38,11 @@ let () =
 
   match mode with
   | Test_util.Decode ->
-    Test_util.decode "test22.c2ml.data" Pb.decode_couple Pp.pp_couple
-      (decode_ref_data ())
+    Test_util.decode "test22.c2ml.data" T.decode_pb_couple T.pp_couple
+      (decode_pb_ref_data ())
   | Test_util.Encode ->
-    Test_util.encode "test22.ml2c.data" Pb.encode_couple (decode_ref_data ())
+    Test_util.encode "test22.ml2c.data" T.encode_pb_couple
+      (decode_pb_ref_data ())
 
 let () =
   let expected_default_person =
diff --git a/src/tests/integration-tests/test23_ml.ml b/src/tests/integration-tests/test23_ml.ml
index 68fadf1..a54f332 100644
--- a/src/tests/integration-tests/test23_ml.ml
+++ b/src/tests/integration-tests/test23_ml.ml
@@ -1,15 +1,14 @@
-module T = Test23_types
-module Pp = Test23_pp
-module Pb = Test23_pb
+module T = Test23
 
-let decode_ref_data () = { T.i = Int64.(shift_left 3L 56) }
+let decode_pb_ref_data () = { T.i = Int64.(shift_left 3L 56) }
 
 let () =
   let mode = Test_util.parse_args () in
 
   match mode with
   | Test_util.Decode ->
-    Test_util.decode "test23.c2ml.data" Pb.decode_int64_t Pp.pp_int64_t
-      (decode_ref_data ())
+    Test_util.decode "test23.c2ml.data" T.decode_pb_int64_t T.pp_int64_t
+      (decode_pb_ref_data ())
   | Test_util.Encode ->
-    Test_util.encode "test23.ml2c.data" Pb.encode_int64_t (decode_ref_data ())
+    Test_util.encode "test23.ml2c.data" T.encode_pb_int64_t
+      (decode_pb_ref_data ())
diff --git a/src/tests/integration-tests/test24_ml.ml b/src/tests/integration-tests/test24_ml.ml
index ef655f8..2e739a5 100644
--- a/src/tests/integration-tests/test24_ml.ml
+++ b/src/tests/integration-tests/test24_ml.ml
@@ -1,14 +1,13 @@
-module T = Test24_types
-module Pb = Test24_pb
-module Pp = Test24_pp
+module T = Test24
 
-let decode_ref_data () : T.a = T.Value "value"
+let decode_pb_ref_data () : T.a = T.Value "value"
 
 let () =
   let mode = Test_util.parse_args () in
 
   match mode with
   | Test_util.Decode ->
-    Test_util.decode "test24.c2ml.data" Pb.decode_a Pp.pp_a (decode_ref_data ())
+    Test_util.decode "test24.c2ml.data" T.decode_pb_a T.pp_a
+      (decode_pb_ref_data ())
   | Test_util.Encode ->
-    Test_util.encode "test24.ml2c.data" Pb.encode_a (decode_ref_data ())
+    Test_util.encode "test24.ml2c.data" T.encode_pb_a (decode_pb_ref_data ())
diff --git a/src/tests/integration-tests/test25_ml.ml b/src/tests/integration-tests/test25_ml.ml
index a496cc6..4d2ef6e 100644
--- a/src/tests/integration-tests/test25_ml.ml
+++ b/src/tests/integration-tests/test25_ml.ml
@@ -1,13 +1,12 @@
-module T = Test25_types
-module Pb = Test25_pb
-module Pp = Test25_pp
+module T = Test25
 
-let decode_ref_data () = T.{ f1 = [ 1l; 2l; 3l; 4l ] }
+let decode_pb_ref_data () = T.{ f1 = [ 1l; 2l; 3l; 4l ] }
 let mode = Test_util.parse_args ()
 
 let () =
   match mode with
   | Test_util.Decode ->
-    Test_util.decode "test25.c2ml.data" Pb.decode_m Pp.pp_m (decode_ref_data ())
+    Test_util.decode "test25.c2ml.data" T.decode_pb_m T.pp_m
+      (decode_pb_ref_data ())
   | Test_util.Encode ->
-    Test_util.encode "test25.ml2c.data" Pb.encode_m (decode_ref_data ())
+    Test_util.encode "test25.ml2c.data" T.encode_pb_m (decode_pb_ref_data ())
diff --git a/src/tests/integration-tests/test26_ml.ml b/src/tests/integration-tests/test26_ml.ml
index 28ddcf5..74d0e84 100644
--- a/src/tests/integration-tests/test26_ml.ml
+++ b/src/tests/integration-tests/test26_ml.ml
@@ -1,12 +1,10 @@
-module T = Test26_types
-module Pb = Test26_pb
-module Pp = Test26_pp
+module T = Test26
 
 let round32 v =
   let open Int32 in
   v |> bits_of_float |> float_of_bits
 
-let decode_ref_data () =
+let decode_pb_ref_data () =
   T.
     {
       double_value = Some 1.23;
@@ -25,7 +23,7 @@ let () =
 
   match mode with
   | Test_util.Decode ->
-    Test_util.decode "test26.c2ml.data" Pb.decode_test Pp.pp_test
-      (decode_ref_data ())
+    Test_util.decode "test26.c2ml.data" T.decode_pb_test T.pp_test
+      (decode_pb_ref_data ())
   | Test_util.Encode ->
-    Test_util.encode "test26.ml2c.data" Pb.encode_test (decode_ref_data ())
+    Test_util.encode "test26.ml2c.data" T.encode_pb_test (decode_pb_ref_data ())
diff --git a/src/tests/integration-tests/test27_ml.ml b/src/tests/integration-tests/test27_ml.ml
index d9d5000..6b43716 100644
--- a/src/tests/integration-tests/test27_ml.ml
+++ b/src/tests/integration-tests/test27_ml.ml
@@ -1,6 +1,4 @@
-module T = Test27_types
-module Pb = Test27_pb
-module Pp = Test27_pp
+module T = Test27
 
 let test_data () =
   T.
@@ -15,6 +13,7 @@ let () =
 
   match mode with
   | Test_util.Decode ->
-    Test_util.decode "test27.c2ml.data" Pb.decode_test Pp.pp_test (test_data ())
+    Test_util.decode "test27.c2ml.data" T.decode_pb_test T.pp_test
+      (test_data ())
   | Test_util.Encode ->
-    Test_util.encode "test27.ml2c.data" Pb.encode_test (test_data ())
+    Test_util.encode "test27.ml2c.data" T.encode_pb_test (test_data ())
diff --git a/src/tests/integration-tests/test_proto3_optional.ml b/src/tests/integration-tests/test_proto3_optional_ml.ml
similarity index 62%
rename from src/tests/integration-tests/test_proto3_optional.ml
rename to src/tests/integration-tests/test_proto3_optional_ml.ml
index bfe4686..df533b6 100644
--- a/src/tests/integration-tests/test_proto3_optional.ml
+++ b/src/tests/integration-tests/test_proto3_optional_ml.ml
@@ -1,12 +1,11 @@
-open Test_proto3_optional_types
-module Pb = Test_proto3_optional_pb
+open Test_proto3_optional
 
 let foo1 : foo = { x = 1l; y = None; name = Some "hello" }
 
 let () =
   let enc = Pbrt.Encoder.create () in
-  Pb.encode_foo foo1 enc;
+  encode_pb_foo foo1 enc;
   let s = Pbrt.Encoder.to_string enc in
   let dec = Pbrt.Decoder.of_string s in
-  let foo1' = Pb.decode_foo dec in
+  let foo1' = decode_pb_foo dec in
   assert (foo1 = foo1')
diff --git a/src/tests/unit-tests/parse_extend.ml b/src/tests/unit-tests/parse_extend.ml
index 6e8efbf..eff79f4 100644
--- a/src/tests/unit-tests/parse_extend.ml
+++ b/src/tests/unit-tests/parse_extend.ml
@@ -1,3 +1,5 @@
+module Pbpt = Pb_parsing_parse_tree
+
 let parse s =
   Pb_parsing_parser.extend_ Pb_parsing_lexer.lexer (Lexing.from_string s)
 
@@ -17,8 +19,8 @@ let () =
     :: [] ->
     assert (field_name = "i");
     assert (field_number = 1);
-    assert (field_type = "int32");
-    assert (field_options = []);
+    assert (field_type = `Int32);
+    assert (field_options = Pb_option.empty);
     assert (`Optional = field_label);
     ()
   | _ -> (assert false : unit));
@@ -36,8 +38,8 @@ let () =
   ] ->
     assert (field_name = "i");
     assert (field_number = 1);
-    assert (field_type = "int32");
-    assert (field_options = []);
+    assert (field_type = `Int32);
+    assert (field_options = Pb_option.empty);
     assert (`Optional = field_label);
     let {
       Pbpt.field_name;
@@ -50,8 +52,8 @@ let () =
     in
     assert (field_name = "j");
     assert (field_number = 2);
-    assert (field_type = "double");
-    assert (field_options = []);
+    assert (field_type = `Double);
+    assert (field_options = Pb_option.empty);
     assert (`Required = field_label);
     ()
   | _ -> (assert false : unit));
diff --git a/src/tests/yojson/dune b/src/tests/yojson/dune
new file mode 100644
index 0000000..b61fa6c
--- /dev/null
+++ b/src/tests/yojson/dune
@@ -0,0 +1,10 @@
+
+(executable
+  (name yojson_unittest)
+  (libraries pbrt pbrt_yojson))
+
+(rule
+  (targets yojson_unittest.ml yojson_unittest.mli)
+  (deps (:file yojson_unittest.proto))
+  (action (run %{project_root}/src/ocaml-protoc/ocaml_protoc.exe
+               %{file} --yojson --ml_out=.)))
diff --git a/src/tests/yojson/yojson_unittest.proto b/src/tests/yojson/yojson_unittest.proto
index 8ea2586..b0c9157 100644
--- a/src/tests/yojson/yojson_unittest.proto
+++ b/src/tests/yojson/yojson_unittest.proto
@@ -57,6 +57,10 @@ message SingleOneOfDuplicate {
   }
 }
 
+message ABasicMessage {
+  string name = 1;
+}
+
 message Test {
   AllBasicTypes all_basic_types = 1; 
   Enum test_enum0 = 2; 
@@ -68,4 +72,5 @@ message Test {
   SingleOneOf single_one_of_small_message = 8; 
   SingleOneOf single_one_of_recursive = 9; 
   repeated Enum repeated_enum = 10;
+  repeated ABasicMessage basic = 11;
 }
diff --git a/src/tests/yojson/yojson_unittest_ml.ml b/src/tests/yojson/yojson_unittest_ml.ml
index 26f181e..86ad343 100644
--- a/src/tests/yojson/yojson_unittest_ml.ml
+++ b/src/tests/yojson/yojson_unittest_ml.ml
@@ -1,5 +1,5 @@
 let test =
-  let open Yojson_unittest_types in
+  let open Yojson_unittest in
   let all_basic_types =
     {
       field01 = 1.2000001;
@@ -51,6 +51,8 @@ let test =
     Recursive_value single_one_of_small_message
   in
 
+  let basic0 : abasic_message = { name = "basic0" } in
+
   {
     all_basic_types = Some all_basic_types;
     test_enum0;
@@ -62,17 +64,18 @@ let test =
     single_one_of_small_message = Some single_one_of_small_message;
     single_one_of_recursive = Some single_one_of_recursive;
     repeated_enum = [ test_enum0; test_enum1; test_enum2 ];
+    basic = [ basic0 ];
   }
 
 let () =
   let json_str =
-    Yojson_unittest_yojson.encode_test test |> Yojson.Basic.to_string
+    Yojson_unittest.encode_json_test test |> Yojson.Basic.to_string
   in
 
   print_endline json_str;
 
   let test' =
-    json_str |> Yojson.Basic.from_string |> Yojson_unittest_yojson.decode_test
+    json_str |> Yojson.Basic.from_string |> Yojson_unittest.decode_json_test
   in
 
   assert (test = test');
@@ -96,10 +99,10 @@ let () =
 
   let test' =
     Bytes.sub_string buffer 0 buffer_len
-    |> Yojson.Basic.from_string |> Yojson_unittest_yojson.decode_test
+    |> Yojson.Basic.from_string |> Yojson_unittest.decode_json_test
   in
 
-  let open Yojson_unittest_types in
+  let open Yojson_unittest in
   assert (test'.all_basic_types = test.all_basic_types);
   assert (test'.test_enum0 = test.test_enum0);
   assert (test'.test_enum1 = test.test_enum1);