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