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