Chih-Hung Hsieh | 92ff605 | 2020-06-10 20:18:39 -0700 | [diff] [blame] | 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; |
Haibo Huang | 72fec01 | 2020-07-10 20:24:04 -0700 | [diff] [blame] | 8 | use inside::protobuf_crate_path; |
Chih-Hung Hsieh | 92ff605 | 2020-06-10 20:18:39 -0700 | [diff] [blame] | 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 { |
| 24 | 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 |
| 37 | fn number(&self) -> i32 { |
| 38 | self.proto.get_number() |
| 39 | } |
| 40 | |
| 41 | // name of enum variant in generated rust code |
| 42 | fn rust_name_inner(&self) -> String { |
| 43 | self.variant_rust_name.clone() |
| 44 | } |
| 45 | |
| 46 | 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> { |
| 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, |
Haibo Huang | ba676d3 | 2020-08-12 13:52:13 -0700 | [diff] [blame] | 80 | customize, |
Chih-Hung Hsieh | 92ff605 | 2020-06-10 20:18:39 -0700 | [diff] [blame] | 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 | |
| 101 | fn allow_alias(&self) -> bool { |
| 102 | self.enum_with_scope.en.get_options().get_allow_alias() |
| 103 | } |
| 104 | |
| 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 | |
| 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 |
| 136 | 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 | |
| 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 | |
| 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(w, &self.customize, "derive(Serialize, Deserialize)"); |
| 174 | let ref type_name = self.type_name; |
| 175 | w.expr_block(&format!("pub enum {}", type_name), |w| { |
| 176 | for value in self.values_all() { |
| 177 | if self.allow_alias() { |
| 178 | w.write_line(&format!( |
| 179 | "{}, // {}", |
| 180 | value.rust_name_inner(), |
| 181 | value.number() |
| 182 | )); |
| 183 | } else { |
| 184 | w.write_line(&format!( |
| 185 | "{} = {},", |
| 186 | value.rust_name_inner(), |
| 187 | value.number() |
| 188 | )); |
| 189 | } |
| 190 | } |
| 191 | }); |
| 192 | } |
| 193 | |
| 194 | fn write_fn_value(&self, w: &mut CodeWriter) { |
| 195 | w.def_fn("value(&self) -> i32", |w| { |
| 196 | if self.allow_alias() { |
| 197 | w.match_expr("*self", |w| { |
| 198 | for value in self.values_all() { |
| 199 | w.case_expr(value.rust_name_outer(), format!("{}", value.number())); |
| 200 | } |
| 201 | }); |
| 202 | } else { |
| 203 | w.write_line("*self as i32") |
| 204 | } |
| 205 | }); |
| 206 | } |
| 207 | |
| 208 | fn write_impl_enum(&self, w: &mut CodeWriter) { |
| 209 | let ref type_name = self.type_name; |
Haibo Huang | ba676d3 | 2020-08-12 13:52:13 -0700 | [diff] [blame] | 210 | w.impl_for_block( |
| 211 | &format!("{}::ProtobufEnum", protobuf_crate_path(&self.customize)), |
| 212 | &format!("{}", type_name), |
| 213 | |w| { |
| 214 | self.write_fn_value(w); |
Chih-Hung Hsieh | 92ff605 | 2020-06-10 20:18:39 -0700 | [diff] [blame] | 215 | |
Chih-Hung Hsieh | 92ff605 | 2020-06-10 20:18:39 -0700 | [diff] [blame] | 216 | w.write_line(""); |
Haibo Huang | ba676d3 | 2020-08-12 13:52:13 -0700 | [diff] [blame] | 217 | let ref type_name = self.type_name; |
Chih-Hung Hsieh | 92ff605 | 2020-06-10 20:18:39 -0700 | [diff] [blame] | 218 | w.def_fn( |
| 219 | &format!( |
Haibo Huang | ba676d3 | 2020-08-12 13:52:13 -0700 | [diff] [blame] | 220 | "from_i32(value: i32) -> ::std::option::Option<{}>", |
| 221 | type_name |
Chih-Hung Hsieh | 92ff605 | 2020-06-10 20:18:39 -0700 | [diff] [blame] | 222 | ), |
| 223 | |w| { |
Haibo Huang | ba676d3 | 2020-08-12 13:52:13 -0700 | [diff] [blame] | 224 | w.match_expr("value", |w| { |
| 225 | let values = self.values_unique(); |
| 226 | for value in values { |
Chih-Hung Hsieh | 92ff605 | 2020-06-10 20:18:39 -0700 | [diff] [blame] | 227 | w.write_line(&format!( |
Haibo Huang | ba676d3 | 2020-08-12 13:52:13 -0700 | [diff] [blame] | 228 | "{} => ::std::option::Option::Some({}),", |
| 229 | value.number(), |
| 230 | value.rust_name_outer() |
| 231 | )); |
| 232 | } |
| 233 | w.write_line(&format!("_ => ::std::option::Option::None")); |
| 234 | }); |
| 235 | }, |
| 236 | ); |
| 237 | |
| 238 | w.write_line(""); |
| 239 | w.def_fn(&format!("values() -> &'static [Self]"), |w| { |
| 240 | w.write_line(&format!("static values: &'static [{}] = &[", type_name)); |
| 241 | w.indented(|w| { |
| 242 | for value in self.values_all() { |
| 243 | w.write_line(&format!("{},", value.rust_name_outer())); |
| 244 | } |
| 245 | }); |
| 246 | w.write_line("];"); |
| 247 | w.write_line("values"); |
| 248 | }); |
| 249 | |
| 250 | if !self.lite_runtime { |
| 251 | w.write_line(""); |
| 252 | w.def_fn( |
| 253 | &format!( |
| 254 | "enum_descriptor_static() -> &'static {}::reflect::EnumDescriptor", |
| 255 | protobuf_crate_path(&self.customize) |
| 256 | ), |
| 257 | |w| { |
| 258 | w.lazy_static_decl_get( |
| 259 | "descriptor", |
| 260 | &format!( |
| 261 | "{}::reflect::EnumDescriptor", |
| 262 | protobuf_crate_path(&self.customize) |
| 263 | ), |
| 264 | &self.customize, |
| 265 | |w| { |
| 266 | let ref type_name = self.type_name; |
| 267 | w.write_line(&format!( |
| 268 | "{}::reflect::EnumDescriptor::new_pb_name::<{}>(\"{}\", {})", |
| 269 | protobuf_crate_path(&self.customize), |
Chih-Hung Hsieh | 92ff605 | 2020-06-10 20:18:39 -0700 | [diff] [blame] | 270 | type_name, |
| 271 | self.enum_with_scope.name_to_package(), |
| 272 | file_descriptor_proto_expr(&self.enum_with_scope.scope) |
| 273 | )); |
Haibo Huang | ba676d3 | 2020-08-12 13:52:13 -0700 | [diff] [blame] | 274 | }, |
| 275 | ); |
| 276 | }, |
| 277 | ); |
| 278 | } |
| 279 | }, |
| 280 | ); |
Chih-Hung Hsieh | 92ff605 | 2020-06-10 20:18:39 -0700 | [diff] [blame] | 281 | } |
| 282 | |
| 283 | fn write_impl_value(&self, w: &mut CodeWriter) { |
Haibo Huang | ba676d3 | 2020-08-12 13:52:13 -0700 | [diff] [blame] | 284 | w.impl_for_block( |
| 285 | &format!( |
| 286 | "{}::reflect::ProtobufValue", |
| 287 | protobuf_crate_path(&self.customize) |
| 288 | ), |
| 289 | &format!("{}", self.type_name), |
| 290 | |w| { |
| 291 | w.def_fn( |
| 292 | &format!( |
| 293 | "as_ref(&self) -> {}::reflect::ReflectValueRef", |
| 294 | protobuf_crate_path(&self.customize) |
| 295 | ), |
| 296 | |w| { |
| 297 | w.write_line(&format!( |
| 298 | "{}::reflect::ReflectValueRef::Enum({}::ProtobufEnum::descriptor(self))", |
| 299 | protobuf_crate_path(&self.customize), |
| 300 | protobuf_crate_path(&self.customize) |
| 301 | )) |
| 302 | }, |
| 303 | ) |
| 304 | }, |
| 305 | ) |
Chih-Hung Hsieh | 92ff605 | 2020-06-10 20:18:39 -0700 | [diff] [blame] | 306 | } |
| 307 | |
| 308 | fn write_impl_copy(&self, w: &mut CodeWriter) { |
| 309 | w.impl_for_block("::std::marker::Copy", &self.type_name, |_w| {}); |
| 310 | } |
| 311 | |
| 312 | fn write_impl_eq(&self, w: &mut CodeWriter) { |
| 313 | assert!(self.allow_alias()); |
Haibo Huang | 72fec01 | 2020-07-10 20:24:04 -0700 | [diff] [blame] | 314 | w.impl_for_block( |
| 315 | "::std::cmp::PartialEq", |
| 316 | &format!("{}", self.type_name), |
| 317 | |w| { |
| 318 | w.def_fn("eq(&self, other: &Self) -> bool", |w| { |
| 319 | w.write_line(&format!( |
| 320 | "{}::ProtobufEnum::value(self) == {}::ProtobufEnum::value(other)", |
| 321 | protobuf_crate_path(&self.customize), |
| 322 | protobuf_crate_path(&self.customize) |
| 323 | )); |
| 324 | }); |
| 325 | }, |
| 326 | ); |
Chih-Hung Hsieh | 92ff605 | 2020-06-10 20:18:39 -0700 | [diff] [blame] | 327 | } |
| 328 | |
| 329 | fn write_impl_hash(&self, w: &mut CodeWriter) { |
| 330 | assert!(self.allow_alias()); |
Haibo Huang | 72fec01 | 2020-07-10 20:24:04 -0700 | [diff] [blame] | 331 | w.impl_for_block("::std::hash::Hash", &format!("{}", self.type_name), |w| { |
Chih-Hung Hsieh | 92ff605 | 2020-06-10 20:18:39 -0700 | [diff] [blame] | 332 | w.def_fn("hash<H : ::std::hash::Hasher>(&self, state: &mut H)", |w| { |
Haibo Huang | 72fec01 | 2020-07-10 20:24:04 -0700 | [diff] [blame] | 333 | w.write_line(&format!( |
| 334 | "state.write_i32({}::ProtobufEnum::value(self))", |
| 335 | protobuf_crate_path(&self.customize) |
| 336 | )); |
Chih-Hung Hsieh | 92ff605 | 2020-06-10 20:18:39 -0700 | [diff] [blame] | 337 | }); |
| 338 | }); |
| 339 | } |
| 340 | |
| 341 | fn write_impl_default(&self, w: &mut CodeWriter) { |
| 342 | let first_value = &self.enum_with_scope.values()[0]; |
| 343 | if first_value.proto.get_number() != 0 { |
| 344 | // This warning is emitted only for proto2 |
| 345 | // (because in proto3 first enum variant number is always 0). |
| 346 | // `Default` implemented unconditionally to simplify certain |
| 347 | // generic operations, e. g. reading a map. |
| 348 | // Also, note that even in proto2 some operations fallback to |
| 349 | // first enum value, e. g. `get_xxx` for unset field, |
| 350 | // so this implementation is not completely unreasonable. |
| 351 | w.comment("Note, `Default` is implemented although default value is not 0"); |
| 352 | } |
| 353 | w.impl_for_block("::std::default::Default", &self.type_name, |w| { |
| 354 | w.def_fn("default() -> Self", |w| { |
| 355 | w.write_line(&format!( |
| 356 | "{}::{}", |
| 357 | &self.type_name, |
| 358 | &first_value.rust_name() |
| 359 | )) |
| 360 | }); |
| 361 | }); |
| 362 | } |
| 363 | } |