1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
#![doc(html_root_url = "https://docs.rs/tower-grpc-build/0.1.0")]
#![deny(rust_2018_idioms)]
#![cfg_attr(test, deny(warnings))]

mod client;
mod server;

use heck::CamelCase;
use std::io;
use std::path::Path;

/// Code generation configuration
pub struct Config {
    prost: prost_build::Config,
    build_client: bool,
    build_server: bool,
}

struct ServiceGenerator {
    client: Option<client::ServiceGenerator>,
    server: Option<server::ServiceGenerator>,
    root_scope: codegen::Scope,
}

impl Config {
    /// Returns a new `Config` with pre-configured prost.
    ///
    /// You can tweak the configuration how the proto buffers are generated and use this config.
    pub fn from_prost(prost: prost_build::Config) -> Self {
        Config {
            prost,
            // Enable client code gen by default
            build_client: true,

            // Disable server code gen by default
            build_server: false,
        }
    }

    /// Returns a new `Config` with default values.
    pub fn new() -> Self {
        Self::from_prost(prost_build::Config::new())
    }

    /// Enable gRPC client code generation
    pub fn enable_client(&mut self, enable: bool) -> &mut Self {
        self.build_client = enable;
        self
    }

    /// Enable gRPC server code generation
    pub fn enable_server(&mut self, enable: bool) -> &mut Self {
        self.build_server = enable;
        self
    }

    /// Generate code
    pub fn build<P>(&mut self, protos: &[P], includes: &[P]) -> io::Result<()>
    where
        P: AsRef<Path>,
    {
        let client = if self.build_client {
            Some(client::ServiceGenerator)
        } else {
            None
        };
        let server = if self.build_server {
            Some(server::ServiceGenerator)
        } else {
            None
        };

        // Set or reset the service generator.
        self.prost.service_generator(Box::new(ServiceGenerator {
            client,
            server,
            root_scope: codegen::Scope::new(),
        }));

        self.prost.compile_protos(protos, includes)
    }
}

impl prost_build::ServiceGenerator for ServiceGenerator {
    fn generate(&mut self, service: prost_build::Service, _buf: &mut String) {
        // Note that neither this implementation of `generate` nor the
        // implementations for `client::ServiceGenerator` and
        // `server::ServiceGenerator` will actually output any code to the
        // buffer; all code is written out in the implementation of the
        // `ServiceGenerator::finalize` function on this type.
        if let Some(ref mut client_generator) = self.client {
            client_generator.generate(&service, &mut self.root_scope);
        }
        if let Some(ref mut server_generator) = self.server {
            server_generator.generate(&service, &mut self.root_scope);
        }
    }

    fn finalize(&mut self, buf: &mut String) {
        // Rather than outputting each service to the buffer as it's generated,
        // we generate the code in our root `codegen::Scope`, which is shared
        // between the generation of each service in the proto file. Unlike a
        // string, codegen provides us with something not unlike a simplified
        // Rust AST, making it easier for us to add new items to modules
        // defined by previous service generator invocations. As we want to
        // output the client and server implementations for each service in the
        // proto file in one `client` or `server` module in the generated code,
        // we wait until all the services have been generated before actually
        // outputting to the buffer.
        let mut fmt = codegen::Formatter::new(buf);
        self.root_scope
            .fmt(&mut fmt)
            .expect("formatting root scope failed!");
        // Reset the root scope so that the service generator is ready to
        // generate another file. this prevents the code generated for *this*
        // file being present in the next file.
        self.root_scope = codegen::Scope::new();
    }
}

/// Extension trait for importing generated types into `codegen::Scope`s.
trait ImportType {
    fn import_type(&mut self, ty: &str, level: usize);
}

// ===== impl ImportType =====

impl ImportType for codegen::Scope {
    fn import_type(&mut self, ty: &str, level: usize) {
        if !is_imported_type(ty) && !is_native_type(ty) {
            let (path, ty) = super_import(ty, level);

            self.import(&path, &ty);
        }
    }
}

impl ImportType for codegen::Module {
    fn import_type(&mut self, ty: &str, level: usize) {
        self.scope().import_type(ty, level);
    }
}

// ===== utility fns =====

fn method_path(service: &prost_build::Service, method: &prost_build::Method) -> String {
    format!(
        "\"/{}.{}/{}\"",
        service.package, service.proto_name, method.proto_name
    )
}

fn lower_name(name: &str) -> String {
    let mut ret = String::new();

    for (i, ch) in name.chars().enumerate() {
        if ch.is_uppercase() {
            if i != 0 {
                ret.push('_');
            }

            ret.push(ch.to_ascii_lowercase());
        } else {
            ret.push(ch);
        }
    }

    ret
}

fn should_import(ty: &str) -> bool {
    !is_imported_type(ty) && !is_native_type(ty)
}

fn is_imported_type(ty: &str) -> bool {
    ty.split("::").map(|t| t == "super").next().unwrap()
}

fn is_native_type(ty: &str) -> bool {
    match ty {
        "()" => true,
        _ => false,
    }
}

fn super_import(ty: &str, level: usize) -> (String, String) {
    let mut v: Vec<&str> = ty.split("::").collect();

    assert!(!is_imported_type(ty));
    assert!(!is_native_type(ty));

    for _ in 0..level {
        v.insert(0, "super");
    }

    let ty = v.pop().unwrap_or(ty);

    (v.join("::"), ty.to_string())
}

fn unqualified(ty: &str, proto_ty: &str, level: usize) -> String {
    if proto_ty == ".google.protobuf.Empty" {
        return "()".to_string();
    }

    if !is_imported_type(ty) {
        return ty.to_string();
    }

    let mut v: Vec<&str> = ty.split("::").collect();

    for _ in 0..level {
        v.insert(0, "super");
    }

    v.join("::")
}

/// Converts a `snake_case` identifier to an `UpperCamel` case Rust type
/// identifier.
///
/// This is identical to the same [function] in `prost-build`, however, we
/// reproduce it here as `prost` does not publically export it.
///
/// We need this as `prost-build` will only give us the snake-case transformed
/// names for gRPC methods, but we need to use method names in types as well.
///
/// [function]: https://github.com/danburkert/prost/blob/d3b971ccd90df35d16069753d52289c0c85014e4/prost-build/src/ident.rs#L28-L38
fn to_upper_camel(s: &str) -> String {
    let mut ident = s.to_camel_case();

    // Add a trailing underscore if the identifier matches a Rust keyword
    // (https://doc.rust-lang.org/grammar.html#keywords).
    if ident == "Self" {
        ident.push('_');
    }
    ident
}

/// This function takes a `&prost_build::Comments` and formats them as a
/// `String` to be used as a RustDoc comment.
fn comments_to_rustdoc(comments: &prost_build::Comments) -> String {
    comments
        .leading
        .iter()
        .fold(String::new(), |acc, s| acc + s.trim_start() + "\n")
}