1 use std::collections::HashSet; 2 3 use protobuf::descriptor::*; 4 5 use crate::customize::ctx::CustomizeElemCtx; 6 use crate::customize::rustproto_proto::customize_from_rustproto_for_enum; 7 use crate::gen::code_writer::CodeWriter; 8 use crate::gen::code_writer::Visibility; 9 use crate::gen::descriptor::write_fn_descriptor; 10 use crate::gen::inside::protobuf_crate_path; 11 use crate::gen::protoc_insertion_point::write_protoc_insertion_point_for_enum; 12 use crate::gen::protoc_insertion_point::write_protoc_insertion_point_for_enum_value; 13 use crate::gen::rust::ident::RustIdent; 14 use crate::gen::rust::ident_with_path::RustIdentWithPath; 15 use crate::gen::rust::snippets::EXPR_NONE; 16 use crate::gen::scope::EnumValueWithContext; 17 use crate::gen::scope::EnumWithScope; 18 use crate::gen::scope::RootScope; 19 use crate::gen::scope::WithScope; 20 21 #[derive(Clone)] 22 pub(crate) struct EnumValueGen<'a> { 23 value: EnumValueWithContext<'a>, 24 enum_rust_name: RustIdentWithPath, 25 } 26 27 impl<'a> EnumValueGen<'a> { parse( value: EnumValueWithContext<'a>, enum_rust_name: &RustIdentWithPath, ) -> EnumValueGen<'a>28 fn parse( 29 value: EnumValueWithContext<'a>, 30 enum_rust_name: &RustIdentWithPath, 31 ) -> EnumValueGen<'a> { 32 EnumValueGen { 33 value: value.clone(), 34 enum_rust_name: enum_rust_name.clone(), 35 } 36 } 37 38 // enum value number(&self) -> i3239 fn number(&self) -> i32 { 40 self.value.proto.proto().number() 41 } 42 43 // name of enum variant in generated rust code rust_name_inner(&self) -> RustIdent44 pub fn rust_name_inner(&self) -> RustIdent { 45 self.value.rust_name() 46 } 47 rust_name_outer(&self) -> RustIdentWithPath48 pub fn rust_name_outer(&self) -> RustIdentWithPath { 49 self.enum_rust_name 50 .to_path() 51 .with_ident(self.rust_name_inner()) 52 } 53 } 54 55 // Codegen for enum definition 56 pub(crate) struct EnumGen<'a> { 57 enum_with_scope: &'a EnumWithScope<'a>, 58 type_name: RustIdentWithPath, 59 lite_runtime: bool, 60 customize: CustomizeElemCtx<'a>, 61 path: &'a [i32], 62 info: Option<&'a SourceCodeInfo>, 63 } 64 65 impl<'a> EnumGen<'a> { new( enum_with_scope: &'a EnumWithScope<'a>, customize: &CustomizeElemCtx<'a>, _root_scope: &RootScope, path: &'a [i32], info: Option<&'a SourceCodeInfo>, ) -> EnumGen<'a>66 pub fn new( 67 enum_with_scope: &'a EnumWithScope<'a>, 68 customize: &CustomizeElemCtx<'a>, 69 _root_scope: &RootScope, 70 path: &'a [i32], 71 info: Option<&'a SourceCodeInfo>, 72 ) -> EnumGen<'a> { 73 let customize = customize.child( 74 &customize_from_rustproto_for_enum(enum_with_scope.en.proto().options.get_or_default()), 75 &enum_with_scope.en, 76 ); 77 let lite_runtime = customize.for_elem.lite_runtime.unwrap_or_else(|| { 78 enum_with_scope 79 .file_descriptor() 80 .proto() 81 .options 82 .optimize_for() 83 == file_options::OptimizeMode::LITE_RUNTIME 84 }); 85 86 EnumGen { 87 enum_with_scope, 88 type_name: enum_with_scope.rust_name().to_path(), 89 lite_runtime, 90 customize, 91 path, 92 info, 93 } 94 } 95 allow_alias(&self) -> bool96 fn allow_alias(&self) -> bool { 97 self.enum_with_scope 98 .en 99 .proto() 100 .options 101 .get_or_default() 102 .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(p, &self.type_name)); 109 } 110 r 111 } 112 values_unique(&self) -> Vec<EnumValueGen>113 fn values_unique(&self) -> Vec<EnumValueGen> { 114 let mut used = HashSet::new(); 115 let mut r = Vec::new(); 116 for p in self.enum_with_scope.values() { 117 if !used.insert(p.proto.proto().number()) { 118 continue; 119 } 120 r.push(EnumValueGen::parse(p, &self.type_name)); 121 } 122 r 123 } 124 write(&self, w: &mut CodeWriter)125 pub fn write(&self, w: &mut CodeWriter) { 126 self.write_enum(w); 127 if self.allow_alias() { 128 w.write_line(""); 129 self.write_impl_eq(w); 130 w.write_line(""); 131 self.write_impl_hash(w); 132 } 133 w.write_line(""); 134 self.write_impl_enum(w); 135 if !self.lite_runtime { 136 w.write_line(""); 137 self.write_impl_enum_full(w); 138 } 139 w.write_line(""); 140 self.write_impl_default(w); 141 w.write_line(""); 142 self.write_impl_self(w); 143 } 144 write_impl_self(&self, w: &mut CodeWriter)145 fn write_impl_self(&self, w: &mut CodeWriter) { 146 if !self.lite_runtime { 147 w.impl_self_block(&format!("{}", self.type_name), |w| { 148 self.write_generated_enum_descriptor_data(w); 149 }); 150 } 151 } 152 write_enum(&self, w: &mut CodeWriter)153 fn write_enum(&self, w: &mut CodeWriter) { 154 w.all_documentation(self.info, self.path); 155 156 let mut derive = Vec::new(); 157 derive.push("Clone"); 158 derive.push("Copy"); 159 if !self.allow_alias() { 160 derive.push("PartialEq"); 161 } 162 derive.push("Eq"); 163 derive.push("Debug"); 164 if !self.allow_alias() { 165 derive.push("Hash"); 166 } else { 167 w.comment("Note: you cannot use pattern matching for enums with allow_alias option"); 168 } 169 w.derive(&derive); 170 let ref type_name = self.type_name; 171 write_protoc_insertion_point_for_enum( 172 w, 173 &self.customize.for_elem, 174 &self.enum_with_scope.en, 175 ); 176 w.expr_block(&format!("pub enum {}", type_name), |w| { 177 for value in self.values_all() { 178 write_protoc_insertion_point_for_enum_value( 179 w, 180 &self.customize.for_children, 181 &value.value.proto, 182 ); 183 if self.allow_alias() { 184 w.write_line(&format!( 185 "{}, // {}", 186 value.rust_name_inner(), 187 value.number() 188 )); 189 } else { 190 w.write_line(&format!( 191 "{} = {},", 192 value.rust_name_inner(), 193 value.number() 194 )); 195 } 196 } 197 }); 198 } 199 write_impl_enum_fn_value(&self, w: &mut CodeWriter)200 fn write_impl_enum_fn_value(&self, w: &mut CodeWriter) { 201 w.def_fn("value(&self) -> i32", |w| { 202 if self.allow_alias() { 203 w.match_expr("*self", |w| { 204 for value in self.values_all() { 205 w.case_expr( 206 &format!("{}", value.rust_name_outer()), 207 &format!("{}", value.number()), 208 ); 209 } 210 }); 211 } else { 212 w.write_line("*self as i32") 213 } 214 }); 215 } 216 write_impl_enum_const_name(&self, w: &mut CodeWriter)217 fn write_impl_enum_const_name(&self, w: &mut CodeWriter) { 218 w.write_line(&format!( 219 "const NAME: &'static str = \"{}\";", 220 self.enum_with_scope.en.name() 221 )); 222 } 223 write_impl_enum_fn_from_i32(&self, w: &mut CodeWriter)224 fn write_impl_enum_fn_from_i32(&self, w: &mut CodeWriter) { 225 w.def_fn( 226 &format!( 227 "from_i32(value: i32) -> ::std::option::Option<{}>", 228 self.type_name 229 ), 230 |w| { 231 w.match_expr("value", |w| { 232 let values = self.values_unique(); 233 for value in values { 234 w.write_line(&format!( 235 "{} => ::std::option::Option::Some({}),", 236 value.number(), 237 value.rust_name_outer() 238 )); 239 } 240 w.write_line(&format!("_ => {}", EXPR_NONE)); 241 }); 242 }, 243 ); 244 } 245 write_impl_enum_const_values(&self, w: &mut CodeWriter)246 fn write_impl_enum_const_values(&self, w: &mut CodeWriter) { 247 w.write_line(&format!("const VALUES: &'static [{}] = &[", self.type_name)); 248 w.indented(|w| { 249 for value in self.values_all() { 250 w.write_line(&format!("{},", value.rust_name_outer())); 251 } 252 }); 253 w.write_line("];"); 254 } 255 write_impl_enum(&self, w: &mut CodeWriter)256 fn write_impl_enum(&self, w: &mut CodeWriter) { 257 w.impl_for_block( 258 &format!("{}::Enum", protobuf_crate_path(&self.customize.for_elem)), 259 &format!("{}", self.type_name), 260 |w| { 261 self.write_impl_enum_const_name(w); 262 w.write_line(""); 263 self.write_impl_enum_fn_value(w); 264 w.write_line(""); 265 self.write_impl_enum_fn_from_i32(w); 266 w.write_line(""); 267 self.write_impl_enum_const_values(w); 268 }, 269 ); 270 } 271 write_impl_enum_full(&self, w: &mut CodeWriter)272 fn write_impl_enum_full(&self, w: &mut CodeWriter) { 273 let ref type_name = self.type_name; 274 w.impl_for_block( 275 &format!( 276 "{}::EnumFull", 277 protobuf_crate_path(&self.customize.for_elem) 278 ), 279 &format!("{}", type_name), 280 |w| { 281 self.write_impl_enum_full_fn_enum_descriptor(w); 282 w.write_line(""); 283 self.write_impl_enum_full_fn_descriptor(w); 284 }, 285 ); 286 } 287 write_impl_enum_full_fn_enum_descriptor(&self, w: &mut CodeWriter)288 fn write_impl_enum_full_fn_enum_descriptor(&self, w: &mut CodeWriter) { 289 write_fn_descriptor( 290 &self.enum_with_scope.en, 291 self.enum_with_scope.scope(), 292 &self.customize.for_elem, 293 w, 294 ); 295 } 296 rust_enum_descriptor_is_enum_index(&self) -> bool297 fn rust_enum_descriptor_is_enum_index(&self) -> bool { 298 if self.allow_alias() { 299 false 300 } else { 301 self.values_all() 302 .into_iter() 303 .enumerate() 304 .all(|(i, value)| (i as i32) == value.number()) 305 } 306 } 307 write_impl_enum_full_fn_descriptor(&self, w: &mut CodeWriter)308 fn write_impl_enum_full_fn_descriptor(&self, w: &mut CodeWriter) { 309 let sig = format!( 310 "descriptor(&self) -> {}::reflect::EnumValueDescriptor", 311 protobuf_crate_path(&self.customize.for_elem) 312 ); 313 w.def_fn(&sig, |w| { 314 if self.rust_enum_descriptor_is_enum_index() { 315 w.write_line("let index = *self as usize;"); 316 } else { 317 w.write_line("let index = match self {"); 318 w.indented(|w| { 319 for (i, value) in self.values_all().into_iter().enumerate() { 320 w.write_line(&format!( 321 "{}::{} => {},", 322 self.type_name, 323 value.rust_name_inner(), 324 i 325 )); 326 } 327 }); 328 w.write_line("};"); 329 } 330 w.write_line(&format!("Self::enum_descriptor().value_by_index(index)")); 331 }); 332 } 333 write_generated_enum_descriptor_data(&self, w: &mut CodeWriter)334 fn write_generated_enum_descriptor_data(&self, w: &mut CodeWriter) { 335 let sig = format!( 336 "generated_enum_descriptor_data() -> {}::reflect::GeneratedEnumDescriptorData", 337 protobuf_crate_path(&self.customize.for_elem) 338 ); 339 w.fn_block( 340 Visibility::Path( 341 self.enum_with_scope 342 .scope() 343 .rust_path_to_file() 344 .to_reverse(), 345 ), 346 &sig, 347 |w| { 348 w.write_line(&format!( 349 "{}::reflect::GeneratedEnumDescriptorData::new::<{}>(\"{}\")", 350 protobuf_crate_path(&self.customize.for_elem), 351 self.type_name, 352 self.enum_with_scope.name_to_package(), 353 )); 354 }, 355 ); 356 } 357 write_impl_eq(&self, w: &mut CodeWriter)358 fn write_impl_eq(&self, w: &mut CodeWriter) { 359 assert!(self.allow_alias()); 360 w.impl_for_block( 361 "::std::cmp::PartialEq", 362 &format!("{}", self.type_name), 363 |w| { 364 w.def_fn("eq(&self, other: &Self) -> bool", |w| { 365 w.write_line(&format!( 366 "{}::Enum::value(self) == {}::Enum::value(other)", 367 protobuf_crate_path(&self.customize.for_elem), 368 protobuf_crate_path(&self.customize.for_elem) 369 )); 370 }); 371 }, 372 ); 373 } 374 write_impl_hash(&self, w: &mut CodeWriter)375 fn write_impl_hash(&self, w: &mut CodeWriter) { 376 assert!(self.allow_alias()); 377 w.impl_for_block("::std::hash::Hash", &format!("{}", self.type_name), |w| { 378 w.def_fn("hash<H : ::std::hash::Hasher>(&self, state: &mut H)", |w| { 379 w.write_line(&format!( 380 "state.write_i32({}::Enum::value(self))", 381 protobuf_crate_path(&self.customize.for_elem) 382 )); 383 }); 384 }); 385 } 386 write_impl_default(&self, w: &mut CodeWriter)387 fn write_impl_default(&self, w: &mut CodeWriter) { 388 let first_value = &self.enum_with_scope.values()[0]; 389 if first_value.proto.proto().number() != 0 { 390 // This warning is emitted only for proto2 391 // (because in proto3 first enum variant number is always 0). 392 // `Default` implemented unconditionally to simplify certain 393 // generic operations, e. g. reading a map. 394 // Also, note that even in proto2 some operations fallback to 395 // first enum value, e. g. `get_xxx` for unset field, 396 // so this implementation is not completely unreasonable. 397 w.comment("Note, `Default` is implemented although default value is not 0"); 398 } 399 w.impl_for_block( 400 "::std::default::Default", 401 &format!("{}", self.type_name), 402 |w| { 403 w.def_fn("default() -> Self", |w| { 404 w.write_line(&format!( 405 "{}::{}", 406 &self.type_name, 407 &first_value.rust_name() 408 )) 409 }); 410 }, 411 ); 412 } 413 } 414