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