grpcio_compiler/
codegen.rs

1// Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0.
2
3// Copyright (c) 2016, Stepan Koltsov
4//
5// Permission is hereby granted, free of charge, to any person obtaining
6// a copy of this software and associated documentation files (the
7// "Software"), to deal in the Software without restriction, including
8// without limitation the rights to use, copy, modify, merge, publish,
9// distribute, sublicense, and/or sell copies of the Software, and to
10// permit persons to whom the Software is furnished to do so, subject to
11// the following conditions:
12//
13// The above copyright notice and this permission notice shall be
14// included in all copies or substantial portions of the Software.
15//
16// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
17// EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
18// MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
19// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
20// LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
21// OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
22// WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
23
24use std::collections::HashMap;
25use std::io::Write;
26
27use protobuf::compiler_plugin;
28use protobuf::descriptor::*;
29use protobuf::descriptorx::*;
30
31struct CodeWriter<'a> {
32    writer: &'a mut (dyn Write + 'a),
33    indent: String,
34}
35
36impl<'a> CodeWriter<'a> {
37    pub fn new(writer: &'a mut dyn Write) -> CodeWriter<'a> {
38        CodeWriter {
39            writer,
40            indent: "".to_string(),
41        }
42    }
43
44    pub fn write_line<S: AsRef<str>>(&mut self, line: S) {
45        (if line.as_ref().is_empty() {
46            self.writer.write_all(b"\n")
47        } else {
48            let s: String = [self.indent.as_ref(), line.as_ref(), "\n"].concat();
49            self.writer.write_all(s.as_bytes())
50        })
51        .unwrap();
52    }
53
54    pub fn write_generated(&mut self) {
55        self.write_line("// This file is generated. Do not edit");
56        self.write_generated_common();
57    }
58
59    fn write_generated_common(&mut self) {
60        // https://secure.phabricator.com/T784
61        self.write_line("// @generated");
62
63        self.write_line("");
64        self.comment("https://github.com/Manishearth/rust-clippy/issues/702");
65        self.write_line("#![allow(unknown_lints)]");
66        self.write_line("#![allow(clippy::all)]");
67        self.write_line("");
68        self.write_line("#![allow(box_pointers)]");
69        self.write_line("#![allow(dead_code)]");
70        self.write_line("#![allow(missing_docs)]");
71        self.write_line("#![allow(non_camel_case_types)]");
72        self.write_line("#![allow(non_snake_case)]");
73        self.write_line("#![allow(non_upper_case_globals)]");
74        self.write_line("#![allow(trivial_casts)]");
75        self.write_line("#![allow(unsafe_code)]");
76        self.write_line("#![allow(unused_imports)]");
77        self.write_line("#![allow(unused_results)]");
78    }
79
80    pub fn indented<F>(&mut self, cb: F)
81    where
82        F: Fn(&mut CodeWriter),
83    {
84        cb(&mut CodeWriter {
85            writer: self.writer,
86            indent: format!("{}    ", self.indent),
87        });
88    }
89
90    #[allow(dead_code)]
91    pub fn commented<F>(&mut self, cb: F)
92    where
93        F: Fn(&mut CodeWriter),
94    {
95        cb(&mut CodeWriter {
96            writer: self.writer,
97            indent: format!("// {}", self.indent),
98        });
99    }
100
101    pub fn block<F>(&mut self, first_line: &str, last_line: &str, cb: F)
102    where
103        F: Fn(&mut CodeWriter),
104    {
105        self.write_line(first_line);
106        self.indented(cb);
107        self.write_line(last_line);
108    }
109
110    pub fn expr_block<F>(&mut self, prefix: &str, cb: F)
111    where
112        F: Fn(&mut CodeWriter),
113    {
114        self.block(&format!("{prefix} {{"), "}", cb);
115    }
116
117    pub fn impl_self_block<S: AsRef<str>, F>(&mut self, name: S, cb: F)
118    where
119        F: Fn(&mut CodeWriter),
120    {
121        self.expr_block(&format!("impl {}", name.as_ref()), cb);
122    }
123
124    pub fn pub_struct<S: AsRef<str>, F>(&mut self, name: S, cb: F)
125    where
126        F: Fn(&mut CodeWriter),
127    {
128        self.expr_block(&format!("pub struct {}", name.as_ref()), cb);
129    }
130
131    pub fn pub_trait<F>(&mut self, name: &str, cb: F)
132    where
133        F: Fn(&mut CodeWriter),
134    {
135        self.expr_block(&format!("pub trait {name}"), cb);
136    }
137
138    pub fn field_entry(&mut self, name: &str, value: &str) {
139        self.write_line(&format!("{name}: {value},"));
140    }
141
142    pub fn field_decl(&mut self, name: &str, field_type: &str) {
143        self.write_line(&format!("{name}: {field_type},"));
144    }
145
146    pub fn comment(&mut self, comment: &str) {
147        if comment.is_empty() {
148            self.write_line("//");
149        } else {
150            self.write_line(&format!("// {comment}"));
151        }
152    }
153
154    pub fn fn_block<F>(&mut self, public: bool, sig: &str, cb: F)
155    where
156        F: Fn(&mut CodeWriter),
157    {
158        if public {
159            self.expr_block(&format!("pub fn {sig}"), cb);
160        } else {
161            self.expr_block(&format!("fn {sig}"), cb);
162        }
163    }
164
165    pub fn pub_fn<F>(&mut self, sig: &str, cb: F)
166    where
167        F: Fn(&mut CodeWriter),
168    {
169        self.fn_block(true, sig, cb);
170    }
171}
172
173use super::util::{self, fq_grpc, to_snake_case, MethodType};
174
175struct MethodGen<'a> {
176    proto: &'a MethodDescriptorProto,
177    service_name: String,
178    service_path: String,
179    root_scope: &'a RootScope<'a>,
180}
181
182impl<'a> MethodGen<'a> {
183    fn new(
184        proto: &'a MethodDescriptorProto,
185        service_name: String,
186        service_path: String,
187        root_scope: &'a RootScope<'a>,
188    ) -> MethodGen<'a> {
189        MethodGen {
190            proto,
191            service_name,
192            service_path,
193            root_scope,
194        }
195    }
196
197    fn input(&self) -> String {
198        format!(
199            "super::{}",
200            self.root_scope
201                .find_message(self.proto.get_input_type())
202                .rust_fq_name()
203        )
204    }
205
206    fn output(&self) -> String {
207        format!(
208            "super::{}",
209            self.root_scope
210                .find_message(self.proto.get_output_type())
211                .rust_fq_name()
212        )
213    }
214
215    fn method_type(&self) -> (MethodType, String) {
216        match (
217            self.proto.get_client_streaming(),
218            self.proto.get_server_streaming(),
219        ) {
220            (false, false) => (MethodType::Unary, fq_grpc("MethodType::Unary")),
221            (true, false) => (
222                MethodType::ClientStreaming,
223                fq_grpc("MethodType::ClientStreaming"),
224            ),
225            (false, true) => (
226                MethodType::ServerStreaming,
227                fq_grpc("MethodType::ServerStreaming"),
228            ),
229            (true, true) => (MethodType::Duplex, fq_grpc("MethodType::Duplex")),
230        }
231    }
232
233    fn service_name(&self) -> String {
234        to_snake_case(&self.service_name)
235    }
236
237    fn name(&self) -> String {
238        to_snake_case(self.proto.get_name())
239    }
240
241    fn fq_name(&self) -> String {
242        format!("\"{}/{}\"", self.service_path, &self.proto.get_name())
243    }
244
245    fn const_method_name(&self) -> String {
246        format!(
247            "METHOD_{}_{}",
248            self.service_name().to_uppercase(),
249            self.name().to_uppercase()
250        )
251    }
252
253    fn write_definition(&self, w: &mut CodeWriter) {
254        let head = format!(
255            "const {}: {}<{}, {}> = {} {{",
256            self.const_method_name(),
257            fq_grpc("Method"),
258            self.input(),
259            self.output(),
260            fq_grpc("Method")
261        );
262        let pb_mar = format!(
263            "{} {{ ser: {}, de: {} }}",
264            fq_grpc("Marshaller"),
265            fq_grpc("pb_ser"),
266            fq_grpc("pb_de")
267        );
268        w.block(&head, "};", |w| {
269            w.field_entry("ty", &self.method_type().1);
270            w.field_entry("name", &self.fq_name());
271            w.field_entry("req_mar", &pb_mar);
272            w.field_entry("resp_mar", &pb_mar);
273        });
274    }
275
276    // Method signatures
277    fn unary(&self, method_name: &str) -> String {
278        format!(
279            "{}(&self, req: &{}) -> {}<{}>",
280            method_name,
281            self.input(),
282            fq_grpc("Result"),
283            self.output()
284        )
285    }
286
287    fn unary_opt(&self, method_name: &str) -> String {
288        format!(
289            "{}_opt(&self, req: &{}, opt: {}) -> {}<{}>",
290            method_name,
291            self.input(),
292            fq_grpc("CallOption"),
293            fq_grpc("Result"),
294            self.output()
295        )
296    }
297
298    fn unary_async(&self, method_name: &str) -> String {
299        format!(
300            "{}_async(&self, req: &{}) -> {}<{}<{}>>",
301            method_name,
302            self.input(),
303            fq_grpc("Result"),
304            fq_grpc("ClientUnaryReceiver"),
305            self.output()
306        )
307    }
308
309    fn unary_async_opt(&self, method_name: &str) -> String {
310        format!(
311            "{}_async_opt(&self, req: &{}, opt: {}) -> {}<{}<{}>>",
312            method_name,
313            self.input(),
314            fq_grpc("CallOption"),
315            fq_grpc("Result"),
316            fq_grpc("ClientUnaryReceiver"),
317            self.output()
318        )
319    }
320
321    fn client_streaming(&self, method_name: &str) -> String {
322        format!(
323            "{}(&self) -> {}<({}<{}>, {}<{}>)>",
324            method_name,
325            fq_grpc("Result"),
326            fq_grpc("ClientCStreamSender"),
327            self.input(),
328            fq_grpc("ClientCStreamReceiver"),
329            self.output()
330        )
331    }
332
333    fn client_streaming_opt(&self, method_name: &str) -> String {
334        format!(
335            "{}_opt(&self, opt: {}) -> {}<({}<{}>, {}<{}>)>",
336            method_name,
337            fq_grpc("CallOption"),
338            fq_grpc("Result"),
339            fq_grpc("ClientCStreamSender"),
340            self.input(),
341            fq_grpc("ClientCStreamReceiver"),
342            self.output()
343        )
344    }
345
346    fn server_streaming(&self, method_name: &str) -> String {
347        format!(
348            "{}(&self, req: &{}) -> {}<{}<{}>>",
349            method_name,
350            self.input(),
351            fq_grpc("Result"),
352            fq_grpc("ClientSStreamReceiver"),
353            self.output()
354        )
355    }
356
357    fn server_streaming_opt(&self, method_name: &str) -> String {
358        format!(
359            "{}_opt(&self, req: &{}, opt: {}) -> {}<{}<{}>>",
360            method_name,
361            self.input(),
362            fq_grpc("CallOption"),
363            fq_grpc("Result"),
364            fq_grpc("ClientSStreamReceiver"),
365            self.output()
366        )
367    }
368
369    fn duplex_streaming(&self, method_name: &str) -> String {
370        format!(
371            "{}(&self) -> {}<({}<{}>, {}<{}>)>",
372            method_name,
373            fq_grpc("Result"),
374            fq_grpc("ClientDuplexSender"),
375            self.input(),
376            fq_grpc("ClientDuplexReceiver"),
377            self.output()
378        )
379    }
380
381    fn duplex_streaming_opt(&self, method_name: &str) -> String {
382        format!(
383            "{}_opt(&self, opt: {}) -> {}<({}<{}>, {}<{}>)>",
384            method_name,
385            fq_grpc("CallOption"),
386            fq_grpc("Result"),
387            fq_grpc("ClientDuplexSender"),
388            self.input(),
389            fq_grpc("ClientDuplexReceiver"),
390            self.output()
391        )
392    }
393
394    fn write_client(&self, w: &mut CodeWriter) {
395        let method_name = self.name();
396        match self.method_type().0 {
397            // Unary
398            MethodType::Unary => {
399                w.pub_fn(&self.unary_opt(&method_name), |w| {
400                    w.write_line(&format!(
401                        "self.client.unary_call(&{}, req, opt)",
402                        self.const_method_name()
403                    ));
404                });
405                w.write_line("");
406
407                w.pub_fn(&self.unary(&method_name), |w| {
408                    w.write_line(&format!(
409                        "self.{}_opt(req, {})",
410                        method_name,
411                        fq_grpc("CallOption::default()")
412                    ));
413                });
414                w.write_line("");
415
416                w.pub_fn(&self.unary_async_opt(&method_name), |w| {
417                    w.write_line(&format!(
418                        "self.client.unary_call_async(&{}, req, opt)",
419                        self.const_method_name()
420                    ));
421                });
422                w.write_line("");
423
424                w.pub_fn(&self.unary_async(&method_name), |w| {
425                    w.write_line(&format!(
426                        "self.{}_async_opt(req, {})",
427                        method_name,
428                        fq_grpc("CallOption::default()")
429                    ));
430                });
431            }
432
433            // Client streaming
434            MethodType::ClientStreaming => {
435                w.pub_fn(&self.client_streaming_opt(&method_name), |w| {
436                    w.write_line(&format!(
437                        "self.client.client_streaming(&{}, opt)",
438                        self.const_method_name()
439                    ));
440                });
441                w.write_line("");
442
443                w.pub_fn(&self.client_streaming(&method_name), |w| {
444                    w.write_line(&format!(
445                        "self.{}_opt({})",
446                        method_name,
447                        fq_grpc("CallOption::default()")
448                    ));
449                });
450            }
451
452            // Server streaming
453            MethodType::ServerStreaming => {
454                w.pub_fn(&self.server_streaming_opt(&method_name), |w| {
455                    w.write_line(&format!(
456                        "self.client.server_streaming(&{}, req, opt)",
457                        self.const_method_name()
458                    ));
459                });
460                w.write_line("");
461
462                w.pub_fn(&self.server_streaming(&method_name), |w| {
463                    w.write_line(&format!(
464                        "self.{}_opt(req, {})",
465                        method_name,
466                        fq_grpc("CallOption::default()")
467                    ));
468                });
469            }
470
471            // Duplex streaming
472            MethodType::Duplex => {
473                w.pub_fn(&self.duplex_streaming_opt(&method_name), |w| {
474                    w.write_line(&format!(
475                        "self.client.duplex_streaming(&{}, opt)",
476                        self.const_method_name()
477                    ));
478                });
479                w.write_line("");
480
481                w.pub_fn(&self.duplex_streaming(&method_name), |w| {
482                    w.write_line(&format!(
483                        "self.{}_opt({})",
484                        method_name,
485                        fq_grpc("CallOption::default()")
486                    ));
487                });
488            }
489        };
490    }
491
492    fn write_service(&self, w: &mut CodeWriter) {
493        let req_stream_type = format!("{}<{}>", fq_grpc("RequestStream"), self.input());
494        let (req, req_type, resp_type) = match self.method_type().0 {
495            MethodType::Unary => ("req", self.input(), "UnarySink"),
496            MethodType::ClientStreaming => ("stream", req_stream_type, "ClientStreamingSink"),
497            MethodType::ServerStreaming => ("req", self.input(), "ServerStreamingSink"),
498            MethodType::Duplex => ("stream", req_stream_type, "DuplexSink"),
499        };
500        let sig = format!(
501            "{}(&mut self, ctx: {}, _{}: {}, sink: {}<{}>)",
502            self.name(),
503            fq_grpc("RpcContext"),
504            req,
505            req_type,
506            fq_grpc(resp_type),
507            self.output()
508        );
509        w.fn_block(false, &sig, |w| {
510            w.write_line("grpcio::unimplemented_call!(ctx, sink)");
511        });
512    }
513
514    fn write_bind(&self, w: &mut CodeWriter) {
515        let add = match self.method_type().0 {
516            MethodType::Unary => "add_unary_handler",
517            MethodType::ClientStreaming => "add_client_streaming_handler",
518            MethodType::ServerStreaming => "add_server_streaming_handler",
519            MethodType::Duplex => "add_duplex_streaming_handler",
520        };
521        w.block(
522            &format!(
523                "builder = builder.{}(&{}, move |ctx, req, resp| {{",
524                add,
525                self.const_method_name()
526            ),
527            "});",
528            |w| {
529                w.write_line(&format!("instance.{}(ctx, req, resp)", self.name()));
530            },
531        );
532    }
533}
534
535struct ServiceGen<'a> {
536    proto: &'a ServiceDescriptorProto,
537    methods: Vec<MethodGen<'a>>,
538}
539
540impl<'a> ServiceGen<'a> {
541    fn new(
542        proto: &'a ServiceDescriptorProto,
543        file: &FileDescriptorProto,
544        root_scope: &'a RootScope,
545    ) -> ServiceGen<'a> {
546        let service_path = if file.get_package().is_empty() {
547            format!("/{}", proto.get_name())
548        } else {
549            format!("/{}.{}", file.get_package(), proto.get_name())
550        };
551        let methods = proto
552            .get_method()
553            .iter()
554            .map(|m| {
555                MethodGen::new(
556                    m,
557                    util::to_camel_case(proto.get_name()),
558                    service_path.clone(),
559                    root_scope,
560                )
561            })
562            .collect();
563
564        ServiceGen { proto, methods }
565    }
566
567    fn service_name(&self) -> String {
568        util::to_camel_case(self.proto.get_name())
569    }
570
571    fn client_name(&self) -> String {
572        format!("{}Client", self.service_name())
573    }
574
575    fn write_client(&self, w: &mut CodeWriter) {
576        w.write_line("#[derive(Clone)]");
577        w.pub_struct(&self.client_name(), |w| {
578            // This can also be exposed by a method. But it may introduce a name conflict
579            // between service definition and method name. Marking it public may put extra
580            // restrict on compatability, but it should not be an issue.
581            w.field_decl("pub client", "::grpcio::Client");
582        });
583
584        w.write_line("");
585
586        w.impl_self_block(&self.client_name(), |w| {
587            w.pub_fn("new(channel: ::grpcio::Channel) -> Self", |w| {
588                w.expr_block(&self.client_name(), |w| {
589                    w.field_entry("client", "::grpcio::Client::new(channel)");
590                });
591            });
592
593            for method in &self.methods {
594                w.write_line("");
595                method.write_client(w);
596            }
597            w.pub_fn(
598                "spawn<F>(&self, f: F) where F: ::std::future::Future<Output = ()> + Send + 'static",
599                |w| {
600                    w.write_line("self.client.spawn(f)");
601                },
602            )
603        });
604    }
605
606    fn write_server(&self, w: &mut CodeWriter) {
607        w.pub_trait(&self.service_name(), |w| {
608            for method in &self.methods {
609                method.write_service(w);
610            }
611        });
612
613        w.write_line("");
614
615        let s = format!(
616            "create_{}<S: {} + Send + Clone + 'static>(s: S) -> {}",
617            to_snake_case(&self.service_name()),
618            self.service_name(),
619            fq_grpc("Service")
620        );
621        w.pub_fn(&s, |w| {
622            w.write_line("let mut builder = ::grpcio::ServiceBuilder::new();");
623            for method in &self.methods[0..self.methods.len() - 1] {
624                w.write_line("let mut instance = s.clone();");
625                method.write_bind(w);
626            }
627
628            w.write_line("let mut instance = s;");
629            self.methods[self.methods.len() - 1].write_bind(w);
630
631            w.write_line("builder.build()");
632        });
633    }
634
635    fn write_method_definitions(&self, w: &mut CodeWriter) {
636        for (i, method) in self.methods.iter().enumerate() {
637            if i != 0 {
638                w.write_line("");
639            }
640
641            method.write_definition(w);
642        }
643    }
644
645    fn write(&self, w: &mut CodeWriter) {
646        self.write_method_definitions(w);
647        w.write_line("");
648        self.write_client(w);
649        w.write_line("");
650        self.write_server(w);
651    }
652}
653
654fn gen_file(
655    file: &FileDescriptorProto,
656    root_scope: &RootScope,
657) -> Option<compiler_plugin::GenResult> {
658    if file.get_service().is_empty() {
659        return None;
660    }
661
662    let base = protobuf::descriptorx::proto_path_to_rust_mod(file.get_name());
663
664    let mut v = Vec::new();
665    {
666        let mut w = CodeWriter::new(&mut v);
667        w.write_generated();
668
669        for service in file.get_service() {
670            w.write_line("");
671            ServiceGen::new(service, file, root_scope).write(&mut w);
672        }
673    }
674
675    Some(compiler_plugin::GenResult {
676        name: base + "_grpc.rs",
677        content: v,
678    })
679}
680
681pub fn gen(
682    file_descriptors: &[FileDescriptorProto],
683    files_to_generate: &[String],
684) -> Vec<compiler_plugin::GenResult> {
685    let files_map: HashMap<&str, &FileDescriptorProto> =
686        file_descriptors.iter().map(|f| (f.get_name(), f)).collect();
687
688    let root_scope = RootScope { file_descriptors };
689
690    let mut results = Vec::new();
691
692    for file_name in files_to_generate {
693        let file = files_map[&file_name[..]];
694
695        if file.get_service().is_empty() {
696            continue;
697        }
698
699        results.extend(gen_file(file, &root_scope).into_iter());
700    }
701
702    results
703}
704
705pub fn protoc_gen_grpc_rust_main() {
706    compiler_plugin::plugin_main(gen);
707}
grpcio_compiler/codegen.rs

grpcio_compiler/
codegen.rs
