1 use std::collections::HashSet;
2 
3 use protobuf::descriptor::*;
4 
5 use super::code_writer::*;
6 use super::customize::Customize;
7 use file_descriptor::file_descriptor_proto_expr;
8 use inside::protobuf_crate_path;
9 use protobuf_name::ProtobufAbsolutePath;
10 use rust_types_values::type_name_to_rust_relative;
11 use scope::EnumWithScope;
12 use scope::RootScope;
13 use scope::WithScope;
14 use serde;
15 
16 #[derive(Clone)]
17 pub struct EnumValueGen {
18     proto: EnumValueDescriptorProto,
19     enum_rust_name: String,
20     variant_rust_name: String,
21 }
22 
23 impl EnumValueGen {
parse( proto: &EnumValueDescriptorProto, enum_rust_name: &str, variant_rust_name: &str, ) -> EnumValueGen24     fn parse(
25         proto: &EnumValueDescriptorProto,
26         enum_rust_name: &str,
27         variant_rust_name: &str,
28     ) -> EnumValueGen {
29         EnumValueGen {
30             proto: proto.clone(),
31             enum_rust_name: enum_rust_name.to_string(),
32             variant_rust_name: variant_rust_name.to_string(),
33         }
34     }
35 
36     // enum value
number(&self) -> i3237     fn number(&self) -> i32 {
38         self.proto.get_number()
39     }
40 
41     // name of enum variant in generated rust code
rust_name_inner(&self) -> String42     fn rust_name_inner(&self) -> String {
43         self.variant_rust_name.clone()
44     }
45 
rust_name_outer(&self) -> String46     pub fn rust_name_outer(&self) -> String {
47         let mut r = String::new();
48         r.push_str(&self.enum_rust_name);
49         r.push_str("::");
50         r.push_str(&self.rust_name_inner());
51         r
52     }
53 }
54 
55 pub(crate) struct EnumGen<'a> {
56     enum_with_scope: &'a EnumWithScope<'a>,
57     type_name: String,
58     lite_runtime: bool,
59     customize: Customize,
60 }
61 
62 impl<'a> EnumGen<'a> {
new( enum_with_scope: &'a EnumWithScope<'a>, current_file: &FileDescriptorProto, customize: &Customize, root_scope: &RootScope, ) -> EnumGen<'a>63     pub fn new(
64         enum_with_scope: &'a EnumWithScope<'a>,
65         current_file: &FileDescriptorProto,
66         customize: &Customize,
67         root_scope: &RootScope,
68     ) -> EnumGen<'a> {
69         let rust_name = if enum_with_scope.get_scope().get_file_descriptor().get_name()
70             == current_file.get_name()
71         {
72             // field type is a message or enum declared in the same file
73             enum_with_scope.rust_name().to_string()
74         } else {
75             type_name_to_rust_relative(
76                 &ProtobufAbsolutePath::from(enum_with_scope.name_absolute()),
77                 current_file,
78                 false,
79                 root_scope,
80                 customize,
81             )
82             .to_string()
83         };
84         let lite_runtime = customize.lite_runtime.unwrap_or_else(|| {
85             enum_with_scope
86                 .get_scope()
87                 .get_file_descriptor()
88                 .get_options()
89                 .get_optimize_for()
90                 == FileOptions_OptimizeMode::LITE_RUNTIME
91         });
92 
93         EnumGen {
94             enum_with_scope,
95             type_name: rust_name,
96             lite_runtime,
97             customize: customize.clone(),
98         }
99     }
100 
allow_alias(&self) -> bool101     fn allow_alias(&self) -> bool {
102         self.enum_with_scope.en.get_options().get_allow_alias()
103     }
104 
values_all(&self) -> Vec<EnumValueGen>105     fn values_all(&self) -> Vec<EnumValueGen> {
106         let mut r = Vec::new();
107         for p in self.enum_with_scope.values() {
108             r.push(EnumValueGen::parse(
109                 &p.proto,
110                 &self.type_name,
111                 p.rust_name().get(),
112             ));
113         }
114         r
115     }
116 
values_unique(&self) -> Vec<EnumValueGen>117     pub fn values_unique(&self) -> Vec<EnumValueGen> {
118         let mut used = HashSet::new();
119         let mut r = Vec::new();
120         for p in self.enum_with_scope.values() {
121             // skipping non-unique enums
122             // TODO: should support it
123             if !used.insert(p.proto.get_number()) {
124                 continue;
125             }
126             r.push(EnumValueGen::parse(
127                 p.proto,
128                 &self.type_name,
129                 p.rust_name().get(),
130             ));
131         }
132         r
133     }
134 
135     // find enum value by name
value_by_name(&'a self, name: &str) -> EnumValueGen136     pub fn value_by_name(&'a self, name: &str) -> EnumValueGen {
137         let v = self.enum_with_scope.value_by_name(name);
138         EnumValueGen::parse(v.proto, &self.type_name, v.rust_name().get())
139     }
140 
write(&self, w: &mut CodeWriter)141     pub fn write(&self, w: &mut CodeWriter) {
142         self.write_struct(w);
143         if self.allow_alias() {
144             w.write_line("");
145             self.write_impl_eq(w);
146             w.write_line("");
147             self.write_impl_hash(w);
148         }
149         w.write_line("");
150         self.write_impl_enum(w);
151         w.write_line("");
152         self.write_impl_copy(w);
153         w.write_line("");
154         self.write_impl_default(w);
155         w.write_line("");
156         self.write_impl_value(w);
157     }
158 
write_struct(&self, w: &mut CodeWriter)159     fn write_struct(&self, w: &mut CodeWriter) {
160         let mut derive = Vec::new();
161         derive.push("Clone");
162         if !self.allow_alias() {
163             derive.push("PartialEq");
164         }
165         derive.push("Eq");
166         derive.push("Debug");
167         if !self.allow_alias() {
168             derive.push("Hash");
169         } else {
170             w.comment("Note: you cannot use pattern matching for enums with allow_alias option");
171         }
172         w.derive(&derive);
173         serde::write_serde_attr(
174             w,
175             &self.customize,
176             "derive(::serde::Serialize, ::serde::Deserialize)",
177         );
178         let ref type_name = self.type_name;
179         w.expr_block(&format!("pub enum {}", type_name), |w| {
180             for value in self.values_all() {
181                 if self.allow_alias() {
182                     w.write_line(&format!(
183                         "{}, // {}",
184                         value.rust_name_inner(),
185                         value.number()
186                     ));
187                 } else {
188                     w.write_line(&format!(
189                         "{} = {},",
190                         value.rust_name_inner(),
191                         value.number()
192                     ));
193                 }
194             }
195         });
196     }
197 
write_fn_value(&self, w: &mut CodeWriter)198     fn write_fn_value(&self, w: &mut CodeWriter) {
199         w.def_fn("value(&self) -> i32", |w| {
200             if self.allow_alias() {
201                 w.match_expr("*self", |w| {
202                     for value in self.values_all() {
203                         w.case_expr(value.rust_name_outer(), format!("{}", value.number()));
204                     }
205                 });
206             } else {
207                 w.write_line("*self as i32")
208             }
209         });
210     }
211 
write_impl_enum(&self, w: &mut CodeWriter)212     fn write_impl_enum(&self, w: &mut CodeWriter) {
213         let ref type_name = self.type_name;
214         w.impl_for_block(
215             &format!("{}::ProtobufEnum", protobuf_crate_path(&self.customize)),
216             &format!("{}", type_name),
217             |w| {
218                 self.write_fn_value(w);
219 
220                 w.write_line("");
221                 let ref type_name = self.type_name;
222                 w.def_fn(
223                     &format!(
224                         "from_i32(value: i32) -> ::std::option::Option<{}>",
225                         type_name
226                     ),
227                     |w| {
228                         w.match_expr("value", |w| {
229                             let values = self.values_unique();
230                             for value in values {
231                                 w.write_line(&format!(
232                                     "{} => ::std::option::Option::Some({}),",
233                                     value.number(),
234                                     value.rust_name_outer()
235                                 ));
236                             }
237                             w.write_line(&format!("_ => ::std::option::Option::None"));
238                         });
239                     },
240                 );
241 
242                 w.write_line("");
243                 w.def_fn(&format!("values() -> &'static [Self]"), |w| {
244                     w.write_line(&format!("static values: &'static [{}] = &[", type_name));
245                     w.indented(|w| {
246                         for value in self.values_all() {
247                             w.write_line(&format!("{},", value.rust_name_outer()));
248                         }
249                     });
250                     w.write_line("];");
251                     w.write_line("values");
252                 });
253 
254                 if !self.lite_runtime {
255                     w.write_line("");
256                     w.def_fn(
257                         &format!(
258                             "enum_descriptor_static() -> &'static {}::reflect::EnumDescriptor",
259                             protobuf_crate_path(&self.customize)
260                         ),
261                         |w| {
262                             w.lazy_static_decl_get(
263                                 "descriptor",
264                                 &format!(
265                                     "{}::reflect::EnumDescriptor",
266                                     protobuf_crate_path(&self.customize)
267                                 ),
268                                 &self.customize,
269                                 |w| {
270                                     let ref type_name = self.type_name;
271                                     w.write_line(&format!(
272                                     "{}::reflect::EnumDescriptor::new_pb_name::<{}>(\"{}\", {})",
273                                     protobuf_crate_path(&self.customize),
274                                     type_name,
275                                     self.enum_with_scope.name_to_package(),
276                                     file_descriptor_proto_expr(&self.enum_with_scope.scope)
277                                 ));
278                                 },
279                             );
280                         },
281                     );
282                 }
283             },
284         );
285     }
286 
write_impl_value(&self, w: &mut CodeWriter)287     fn write_impl_value(&self, w: &mut CodeWriter) {
288         w.impl_for_block(
289             &format!(
290                 "{}::reflect::ProtobufValue",
291                 protobuf_crate_path(&self.customize)
292             ),
293             &format!("{}", self.type_name),
294             |w| {
295                 w.def_fn(
296                     &format!(
297                         "as_ref(&self) -> {}::reflect::ReflectValueRef",
298                         protobuf_crate_path(&self.customize)
299                     ),
300                     |w| {
301                         w.write_line(&format!(
302                         "{}::reflect::ReflectValueRef::Enum({}::ProtobufEnum::descriptor(self))",
303                         protobuf_crate_path(&self.customize),
304                         protobuf_crate_path(&self.customize)
305                     ))
306                     },
307                 )
308             },
309         )
310     }
311 
write_impl_copy(&self, w: &mut CodeWriter)312     fn write_impl_copy(&self, w: &mut CodeWriter) {
313         w.impl_for_block("::std::marker::Copy", &self.type_name, |_w| {});
314     }
315 
write_impl_eq(&self, w: &mut CodeWriter)316     fn write_impl_eq(&self, w: &mut CodeWriter) {
317         assert!(self.allow_alias());
318         w.impl_for_block(
319             "::std::cmp::PartialEq",
320             &format!("{}", self.type_name),
321             |w| {
322                 w.def_fn("eq(&self, other: &Self) -> bool", |w| {
323                     w.write_line(&format!(
324                         "{}::ProtobufEnum::value(self) == {}::ProtobufEnum::value(other)",
325                         protobuf_crate_path(&self.customize),
326                         protobuf_crate_path(&self.customize)
327                     ));
328                 });
329             },
330         );
331     }
332 
write_impl_hash(&self, w: &mut CodeWriter)333     fn write_impl_hash(&self, w: &mut CodeWriter) {
334         assert!(self.allow_alias());
335         w.impl_for_block("::std::hash::Hash", &format!("{}", self.type_name), |w| {
336             w.def_fn("hash<H : ::std::hash::Hasher>(&self, state: &mut H)", |w| {
337                 w.write_line(&format!(
338                     "state.write_i32({}::ProtobufEnum::value(self))",
339                     protobuf_crate_path(&self.customize)
340                 ));
341             });
342         });
343     }
344 
write_impl_default(&self, w: &mut CodeWriter)345     fn write_impl_default(&self, w: &mut CodeWriter) {
346         let first_value = &self.enum_with_scope.values()[0];
347         if first_value.proto.get_number() != 0 {
348             // This warning is emitted only for proto2
349             // (because in proto3 first enum variant number is always 0).
350             // `Default` implemented unconditionally to simplify certain
351             // generic operations, e. g. reading a map.
352             // Also, note that even in proto2 some operations fallback to
353             // first enum value, e. g. `get_xxx` for unset field,
354             // so this implementation is not completely unreasonable.
355             w.comment("Note, `Default` is implemented although default value is not 0");
356         }
357         w.impl_for_block("::std::default::Default", &self.type_name, |w| {
358             w.def_fn("default() -> Self", |w| {
359                 w.write_line(&format!(
360                     "{}::{}",
361                     &self.type_name,
362                     &first_value.rust_name()
363                 ))
364             });
365         });
366     }
367 }
368