nrc_protobuf_codegen/
lib.rs

1extern crate protobuf;
2extern crate heck;
3
4use std::collections::hash_map::HashMap;
5use std::fmt::Write as FmtWrite;
6use std::fs::File;
7use std::io;
8use std::io::Write;
9use std::path::Path;
10
11use protobuf::compiler_plugin;
12use protobuf::descriptor::*;
13use protobuf::descriptorx::*;
14use protobuf::Message;
15
16mod customize;
17mod enums;
18mod extensions;
19mod field;
20mod file_and_mod;
21pub mod float;
22mod ident;
23mod inside;
24mod message;
25mod oneof;
26mod rust_name;
27mod rust_types_values;
28mod serde;
29mod well_known_types;
30
31use customize::customize_from_rustproto_for_file;
32pub use customize::Customize;
33
34pub mod code_writer;
35
36use self::code_writer::CodeWriter;
37use self::enums::*;
38use self::extensions::*;
39use self::message::*;
40use inside::protobuf_crate_path;
41
42fn escape_byte(s: &mut String, b: u8) {
43    if b == b'\n' {
44        write!(s, "\\n").unwrap();
45    } else if b == b'\r' {
46        write!(s, "\\r").unwrap();
47    } else if b == b'\t' {
48        write!(s, "\\t").unwrap();
49    } else if b == b'\\' || b == b'"' {
50        write!(s, "\\{}", b as char).unwrap();
51    } else if b == b'\0' {
52        write!(s, "\\0").unwrap();
53    // ASCII printable except space
54    } else if b > 0x20 && b < 0x7f {
55        write!(s, "{}", b as char).unwrap();
56    } else {
57        write!(s, "\\x{:02x}", b).unwrap();
58    }
59}
60
61fn write_file_descriptor_data(
62    file: &FileDescriptorProto,
63    customize: &Customize,
64    w: &mut CodeWriter,
65) {
66    let fdp_bytes = file.write_to_bytes().unwrap();
67    w.write_line("static file_descriptor_proto_data: &'static [u8] = b\"\\");
68    w.indented(|w| {
69        const MAX_LINE_LEN: usize = 72;
70
71        let mut s = String::new();
72        for &b in &fdp_bytes {
73            let prev_len = s.len();
74            escape_byte(&mut s, b);
75            let truncate = s.len() > MAX_LINE_LEN;
76            if truncate {
77                s.truncate(prev_len);
78            }
79            if truncate || s.len() == MAX_LINE_LEN {
80                write!(s, "\\").unwrap();
81                w.write_line(&s);
82                s.clear();
83            }
84            if truncate {
85                escape_byte(&mut s, b);
86            }
87        }
88        if !s.is_empty() {
89            write!(s, "\\").unwrap();
90            w.write_line(&s);
91            s.clear();
92        }
93    });
94    w.write_line("\";");
95    w.write_line("");
96    w.lazy_static_protobuf_path(
97        "file_descriptor_proto_lazy",
98        &format!(
99            "{}::descriptor::FileDescriptorProto",
100            protobuf_crate_path(customize)
101        ),
102        protobuf_crate_path(customize),
103    );
104    w.write_line("");
105    w.def_fn(
106        "parse_descriptor_proto() -> ::protobuf::descriptor::FileDescriptorProto",
107        |w| {
108            w.write_line("::protobuf::parse_from_bytes(file_descriptor_proto_data).unwrap()");
109        },
110    );
111    w.write_line("");
112    w.pub_fn(
113        "file_descriptor_proto() -> &'static ::protobuf::descriptor::FileDescriptorProto",
114        |w| {
115            w.unsafe_expr(|w| {
116                w.block("file_descriptor_proto_lazy.get(|| {", "})", |w| {
117                    w.write_line("parse_descriptor_proto()");
118                });
119            });
120        },
121    );
122}
123
124fn gen_file(
125    file: &FileDescriptorProto,
126    _files_map: &HashMap<&str, &FileDescriptorProto>,
127    root_scope: &RootScope,
128    customize: &Customize,
129) -> Option<compiler_plugin::GenResult> {
130    // TODO: use it
131    let mut customize = customize.clone();
132    // options specified in invocation have precedence over options specified in file
133    customize.update_with(&customize_from_rustproto_for_file(file.get_options()));
134
135    let scope = FileScope {
136        file_descriptor: file,
137    }
138    .to_scope();
139    let lite_runtime = customize.lite_runtime.unwrap_or_else(|| {
140        file.get_options().get_optimize_for() == FileOptions_OptimizeMode::LITE_RUNTIME
141    });
142
143    let mut v = Vec::new();
144
145    {
146        let mut w = CodeWriter::new(&mut v);
147
148        w.write_generated_by("rust-protobuf", env!("CARGO_PKG_VERSION"));
149        w.write_line(&format!("//! Generated file from `{}`", file.get_name()));
150
151        w.write_line("");
152        w.write_line("use protobuf::Message as Message_imported_for_functions;");
153        w.write_line("use protobuf::ProtobufEnum as ProtobufEnum_imported_for_functions;");
154        if customize.inside_protobuf != Some(true) {
155            w.write_line("");
156            w.write_line("/// Generated files are compatible only with the same version");
157            w.write_line("/// of protobuf runtime.");
158            w.write_line(&format!(
159                "const _PROTOBUF_VERSION_CHECK: () = {}::{};",
160                protobuf_crate_path(&customize),
161                protobuf::VERSION_IDENT
162            ));
163        }
164
165        for message in &scope.get_messages() {
166            // ignore map entries, because they are not used in map fields
167            if message.map_entry().is_none() {
168                w.write_line("");
169                MessageGen::new(message, &root_scope, &customize).write(&mut w);
170            }
171        }
172        for enum_type in &scope.get_enums() {
173            w.write_line("");
174            EnumGen::new(enum_type, file, &customize, root_scope).write(&mut w);
175        }
176
177        write_extensions(file, &root_scope, &mut w, &customize);
178
179        if !lite_runtime {
180            w.write_line("");
181            write_file_descriptor_data(file, &customize, &mut w);
182        }
183    }
184
185    Some(compiler_plugin::GenResult {
186        name: format!("{}.rs", proto_path_to_rust_mod(file.get_name())),
187        content: v,
188    })
189}
190
191// This function is also used externally by cargo plugin
192// https://github.com/plietar/rust-protobuf-build
193// So be careful changing its signature.
194pub fn gen(
195    file_descriptors: &[FileDescriptorProto],
196    files_to_generate: &[String],
197    customize: &Customize,
198) -> Vec<compiler_plugin::GenResult> {
199    let root_scope = RootScope {
200        file_descriptors: file_descriptors,
201    };
202
203    let mut results: Vec<compiler_plugin::GenResult> = Vec::new();
204    let files_map: HashMap<&str, &FileDescriptorProto> =
205        file_descriptors.iter().map(|f| (f.get_name(), f)).collect();
206
207    let all_file_names: Vec<&str> = file_descriptors.iter().map(|f| f.get_name()).collect();
208
209    for file_name in files_to_generate {
210        let file = files_map.get(&file_name[..]).expect(&format!(
211            "file not found in file descriptors: {:?}, files: {:?}",
212            file_name, all_file_names
213        ));
214        results.extend(gen_file(file, &files_map, &root_scope, customize));
215    }
216    results
217}
218
219pub fn gen_and_write(
220    file_descriptors: &[FileDescriptorProto],
221    files_to_generate: &[String],
222    out_dir: &Path,
223    customize: &Customize,
224) -> io::Result<()> {
225    let results = gen(file_descriptors, files_to_generate, customize);
226
227    for r in &results {
228        let mut file_path = out_dir.to_owned();
229        file_path.push(&r.name);
230        let mut file_writer = File::create(&file_path)?;
231        file_writer.write_all(&r.content)?;
232        file_writer.flush()?;
233    }
234
235    Ok(())
236}
237
238pub fn protoc_gen_rust_main() {
239    compiler_plugin::plugin_main_2(|r| {
240        let customize = Customize::parse_from_parameter(r.parameter).expect("parse options");
241        gen(r.file_descriptors, r.files_to_generate, &customize)
242    });
243}