prost-build: pass Config, ServiceGenerator by mut ref

Changing ServiceGenerator::{generate, finalize} to take `&mut self` is
motivated by tower-rs/tower-grpc#21. The `ServiceGenerator` change
necessitated changing Config::compile_protos to take `&mut self` as
well. I don't expect the latter change to impact many people, since it's
the more advanced API (as opposed to the standalone `compile_protos`
function), and to be useful the API required a mutable config, anyway.
It also provides flexibility in the future if we want to do something
else requiring mutable access to the config.

This is a breaking change.

Author: danburkert

benchmarks/build.rs

@@ -3,11 +3,10 @@ extern crate protobuf;
 
 fn main() {
     let benchmarks = protobuf::include().join("benchmarks");
-    let prost_build = prost_build::Config::new();
-    prost_build.compile_protos(&[benchmarks.join("benchmark_messages_proto2.proto")],
-                               &[benchmarks.clone()]).unwrap();
-    prost_build.compile_protos(&[benchmarks.join("benchmark_messages_proto3.proto")],
-                               &[benchmarks]).unwrap();
+    prost_build::compile_protos(&[benchmarks.join("benchmark_messages_proto2.proto")],
+                                &[benchmarks.clone()]).unwrap();
+    prost_build::compile_protos(&[benchmarks.join("benchmark_messages_proto3.proto")],
+                                &[benchmarks]).unwrap();
 
     println!("cargo:rustc-env=GOOGLE_MESSAGE1={}",
              protobuf::share().join("google_message1.dat").display());

prost-build/src/code_generator.rs

@@ -46,7 +46,7 @@ enum Syntax {
 }
 
 pub struct CodeGenerator<'a> {
-    config: &'a Config,
+    config: &'a mut Config,
     package: String,
     source_info: SourceCodeInfo,
     syntax: Syntax,
@@ -57,7 +57,7 @@ pub struct CodeGenerator<'a> {
 }
 
 impl <'a> CodeGenerator<'a> {
-    pub fn generate(config: &Config,
+    pub fn generate(config: &mut Config,
                     message_graph: &MessageGraph,
                     file: FileDescriptorProto,
                     buf: &mut String) {
@@ -104,14 +104,19 @@ impl <'a> CodeGenerator<'a> {
         }
         code_gen.path.pop();
 
-        if let Some(ref service_generator) = code_gen.config.service_generator {
+        if code_gen.config.service_generator.is_some() {
             code_gen.path.push(6);
             for (idx, service) in file.service.into_iter().enumerate() {
                 code_gen.path.push(idx as i32);
-                service_generator.generate(code_gen.unpack_service(service), &mut code_gen.buf);
+                code_gen.push_service(service);
                 code_gen.path.pop();
             }
-            service_generator.finalize(&mut code_gen.buf);
+
+            let buf = code_gen.buf;
+            code_gen.config.service_generator.as_mut().map(|service_generator| {
+                service_generator.finalize(buf);
+            });
+
             code_gen.path.pop();
         }
     }
@@ -226,21 +231,23 @@ impl <'a> CodeGenerator<'a> {
 
     fn append_type_attributes(&mut self, msg_name: &str) {
         assert_eq!(b'.', msg_name.as_bytes()[0]);
-        for &(ref matcher, ref attribute) in &self.config.type_attributes {
-            if match_ident(matcher, msg_name, None) {
+        // TODO: this clone is dirty, but expedious.
+        for (matcher, attribute) in self.config.type_attributes.clone() {
+            if match_ident(&matcher, msg_name, None) {
                 self.push_indent();
-                self.buf.push_str(attribute);
+                self.buf.push_str(&attribute);
                 self.buf.push('\n');
             }
         }
     }
 
     fn append_field_attributes(&mut self, msg_name: &str, field_name: &str) {
         assert_eq!(b'.', msg_name.as_bytes()[0]);
-        for &(ref matcher, ref attribute) in &self.config.field_attributes {
-            if match_ident(matcher, msg_name, Some(field_name)) {
+        // TODO: this clone is dirty, but expedious.
+        for (matcher, attribute) in self.config.field_attributes.clone() {
+            if match_ident(&matcher, msg_name, Some(field_name)) {
                 self.push_indent();
-                self.buf.push_str(attribute);
+                self.buf.push_str(&attribute);
                 self.buf.push('\n');
             }
         }
@@ -491,7 +498,7 @@ impl <'a> CodeGenerator<'a> {
         self.buf.push_str(",\n");
     }
 
-    fn unpack_service(&mut self, service: ServiceDescriptorProto) -> Service {
+    fn push_service(&mut self, service: ServiceDescriptorProto) {
         let name = service.name().to_owned();
         debug!("  service: {:?}", name);
 
@@ -531,14 +538,17 @@ impl <'a> CodeGenerator<'a> {
                              .collect();
         self.path.pop();
 
-        Service {
+        let service = Service {
             name: to_upper_camel(&name),
             proto_name: name,
             package: self.package.clone(),
             comments,
             methods,
             options: service.options.unwrap_or_default(),
-        }
+        };
+
+        let buf = &mut self.buf;
+        self.config.service_generator.as_mut().map(move |service_generator| service_generator.generate(service, buf));
     }
 
     fn push_indent(&mut self) {

prost-build/src/lib.rs

@@ -172,7 +172,7 @@ type Module = Vec<String>;
 pub trait ServiceGenerator {
     /// Generates a Rust interface or implementation for a service, writing the
     /// result to `buf`.
-    fn generate(&self, service: Service, buf: &mut String);
+    fn generate(&mut self, service: Service, buf: &mut String);
 
     /// Finalizes the generation process.
     ///
@@ -187,7 +187,7 @@ pub trait ServiceGenerator {
     /// is called once per `.proto` file.
     ///
     /// The default implementation is empty and does nothing.
-    fn finalize(&self, _buf: &mut String) {}
+    fn finalize(&mut self, _buf: &mut String) {}
 }
 
 /// Configuration options for Protobuf code generation.
@@ -371,7 +371,7 @@ impl Config {
     ///                                &["src"]).unwrap();
     /// }
     /// ```
-    pub fn compile_protos<P>(&self, protos: &[P], includes: &[P]) -> Result<()> where P: AsRef<Path> {
+    pub fn compile_protos<P>(&mut self, protos: &[P], includes: &[P]) -> Result<()> where P: AsRef<Path> {
         let target = match env::var("OUT_DIR") {
             Ok(val) => PathBuf::from(val),
             Err(env::VarError::NotPresent) => {
@@ -434,7 +434,7 @@ impl Config {
         Ok(())
     }
 
-    fn generate(&self, files: Vec<FileDescriptorProto>) -> HashMap<Module, String> {
+    fn generate(&mut self, files: Vec<FileDescriptorProto>) -> HashMap<Module, String> {
         let mut modules = HashMap::new();
 
         let message_graph = MessageGraph::new(&files);
@@ -525,7 +525,7 @@ mod tests {
     /// service methods.
     struct ServiceTraitGenerator;
     impl ServiceGenerator for ServiceTraitGenerator {
-        fn generate(&self, service: Service, buf: &mut String) {
+        fn generate(&mut self, service: Service, buf: &mut String) {
             // Generate a trait for the service.
             service.comments.append_with_indent(0, buf);
             buf.push_str(&format!("trait {} {{\n", &service.name));
@@ -542,7 +542,7 @@ mod tests {
             // Close out the trait.
             buf.push_str("}\n");
         }
-        fn finalize(&self, buf: &mut String) {
+        fn finalize(&mut self, buf: &mut String) {
             // Needs to be present only once, no matter how many services there are
             buf.push_str("pub mod utils { }\n");
         }