1use std::{collections::HashMap, path::PathBuf, sync::Arc};
2
3use faststr::FastStr;
4use fxhash::{FxHashMap, FxHashSet};
5use itertools::Itertools;
6use normpath::PathExt;
7use protobuf::descriptor::{
8 DescriptorProto,
9 EnumDescriptorProto, field_descriptor_proto::{Label, Type}, ServiceDescriptorProto,
10};
11pub use protobuf::descriptor::FileDescriptorProto;
12
13use crate::{
14 index::Idx,
15 ir::{self, FieldKind, Item, Path, TyKind},
16 symbol::{EnumRepr, FileId, Ident, IdentName},
17 tags::{
18 PilotaName,
19 protobuf::{ClientStreaming, OneOf, ProstType, Repeated, ServerStreaming}, Tags,
20 },
21};
22
23use super::Parser;
24
25#[derive(Default)]
26pub struct ProtobufParser {
27 inner: protobuf_parse::Parser,
28 include_dirs: Vec<PathBuf>,
29 input_files: FxHashSet<PathBuf>,
30}
31
32#[derive(PartialEq, Eq)]
33pub enum Syntax {
34 Proto2,
35 Proto3,
36}
37
38struct Lower {
39 next_file_id: FileId,
40 files: FxHashMap<String, FileId>,
41 cur_package: Option<String>,
42 cur_syntax: Syntax,
43}
44
45impl Default for Lower {
46 fn default() -> Self {
47 Self {
48 next_file_id: FileId::from_u32(0),
49 files: Default::default(),
50 cur_package: None,
51 cur_syntax: Syntax::Proto3,
52 }
53 }
54}
55
56impl Lower {
57 fn str2path(&self, s: &str) -> ir::Path {
58 ir::Path {
59 segments: Arc::from_iter(s.split('.').map(FastStr::new).map(Ident::from)),
60 }
61 }
62
63 fn lower_ty(
64 &self,
65 type_: Option<protobuf::EnumOrUnknown<protobuf::descriptor::field_descriptor_proto::Type>>,
66 type_name: Option<&str>,
67 nested_messages: &FxHashMap<String, &DescriptorProto>,
68 message_name: Option<&str>,
69 ) -> ir::Ty {
70 if let Some(name) = type_name {
71 if let Some(msg) = nested_messages.get(name) {
72 if msg.options.has_map_entry() {
73 let key = &msg.field[0];
74 let value = &msg.field[1];
75 assert_eq!("key", key.name());
76 assert_eq!("value", value.name());
77 return ir::Ty {
78 kind: ir::TyKind::Map(
79 Arc::from(self.lower_ty(
80 key.type_,
81 key.type_name.as_deref(),
82 nested_messages,
83 message_name,
84 )),
85 Arc::from(self.lower_ty(
86 value.type_,
87 value.type_name.as_deref(),
88 nested_messages,
89 message_name,
90 )),
91 ),
92 tags: Default::default(),
93 };
94 }
95 }
96
97 assert_eq!(".", &name[..1]);
98
99 return ir::Ty {
100 kind: ir::TyKind::Path(self.str2path(&name[1..])),
101 tags: Default::default(),
102 };
103 }
104 let Some(ty) = type_ else {
105 panic!()
106 };
107
108 let mut tags = Tags::default();
109 let kind = match ty.enum_value().unwrap() {
110 protobuf::descriptor::field_descriptor_proto::Type::TYPE_DOUBLE => ir::TyKind::F64,
111 protobuf::descriptor::field_descriptor_proto::Type::TYPE_FLOAT => ir::TyKind::F32,
112 protobuf::descriptor::field_descriptor_proto::Type::TYPE_INT64 => ir::TyKind::I64,
113 protobuf::descriptor::field_descriptor_proto::Type::TYPE_UINT64 => ir::TyKind::UInt64,
114 protobuf::descriptor::field_descriptor_proto::Type::TYPE_INT32 => ir::TyKind::I32,
115 protobuf::descriptor::field_descriptor_proto::Type::TYPE_FIXED64 => {
116 tags.insert(ProstType::Fixed64);
117 ir::TyKind::UInt64
118 }
119 protobuf::descriptor::field_descriptor_proto::Type::TYPE_FIXED32 => {
120 tags.insert(ProstType::Fixed32);
121 ir::TyKind::UInt32
122 }
123 protobuf::descriptor::field_descriptor_proto::Type::TYPE_BOOL => ir::TyKind::Bool,
124 protobuf::descriptor::field_descriptor_proto::Type::TYPE_STRING => ir::TyKind::String,
125 protobuf::descriptor::field_descriptor_proto::Type::TYPE_GROUP => todo!(),
126 protobuf::descriptor::field_descriptor_proto::Type::TYPE_BYTES => ir::TyKind::Bytes,
127 protobuf::descriptor::field_descriptor_proto::Type::TYPE_UINT32 => ir::TyKind::UInt32,
128 protobuf::descriptor::field_descriptor_proto::Type::TYPE_SFIXED32 => {
129 tags.insert(ProstType::SFixed32);
130 ir::TyKind::I32
131 }
132 protobuf::descriptor::field_descriptor_proto::Type::TYPE_SFIXED64 => {
133 tags.insert(ProstType::SFixed64);
134 ir::TyKind::I64
135 }
136 protobuf::descriptor::field_descriptor_proto::Type::TYPE_SINT32 => {
137 tags.insert(ProstType::SInt32);
138 ir::TyKind::I32
139 }
140 protobuf::descriptor::field_descriptor_proto::Type::TYPE_SINT64 => {
141 tags.insert(ProstType::SInt64);
142 ir::TyKind::I64
143 }
144
145 protobuf::descriptor::field_descriptor_proto::Type::TYPE_MESSAGE
146 | protobuf::descriptor::field_descriptor_proto::Type::TYPE_ENUM => unreachable!(),
147 };
148
149 ir::Ty {
150 kind,
151 tags: Arc::new(tags),
152 }
153 }
154
155 fn lower_enum(&self, e: &EnumDescriptorProto) -> ir::Item {
156 ir::Item {
157 related_items: Default::default(),
158 tags: Default::default(),
159 kind: ir::ItemKind::Enum(ir::Enum {
160 name: FastStr::new(e.name()).into(),
161 variants: e
162 .value
163 .iter()
164 .map(|v| ir::EnumVariant {
165 id: v.number,
166 name: FastStr::new(v.name()).into(),
167 discr: v.number.map(|v| v as i64),
168 tags: Default::default(),
169 fields: Default::default(),
170 })
171 .collect_vec(),
172 repr: Some(EnumRepr::I32),
173 }),
174 }
175 }
176
177 fn lower_message(
178 &self,
179 message: &DescriptorProto,
180 parent_messages: &mut Vec<String>,
181 ) -> Vec<ir::Item> {
182 let fq_message_name = format!(
183 "{}{}.{}{}",
184 if self.cur_package.is_none() { "" } else { "." },
185 self.cur_package.as_deref().unwrap_or(""),
186 {
187 let mut s = String::new();
188 parent_messages.iter().for_each(|m| {
189 s.push_str(m);
190 s.push('.')
191 });
192 s
193 },
194 message.name()
195 );
196
197 let nested_messages = message
198 .nested_type
199 .iter()
200 .map(|m| (format!("{}.{}", fq_message_name, m.name()), m))
201 .collect::<FxHashMap<_, _>>();
202
203 let mut fields = Vec::default();
204 let mut oneof_fields = FxHashMap::default();
205
206 message.field.iter().enumerate().for_each(|(idx, field)| {
207 if field.proto3_optional.unwrap_or(false) {
208 fields.push((idx, field))
209 } else if let Some(oneof_index) = field.oneof_index {
210 oneof_fields
211 .entry(oneof_index)
212 .or_insert_with(Vec::default)
213 .push((idx, field))
214 } else {
215 fields.push((idx, field))
216 }
217 });
218
219 let mut nested_items: Vec<_> = Default::default();
220
221 let mut extra_fields = Vec::default();
222
223 message.oneof_decl.iter().enumerate().for_each(|(idx, d)| {
224 if let Some(fields) = oneof_fields.remove(&(idx as i32)) {
225 nested_items.push(Arc::new(ir::Item {
226 related_items: Default::default(),
227 tags: Arc::new(crate::tags!(OneOf)),
228 kind: ir::ItemKind::Enum(ir::Enum {
229 name: FastStr::new(d.name()).into(),
230 repr: None,
231 variants: fields
232 .iter()
233 .map(|(_, f)| ir::EnumVariant {
234 discr: None,
235 id: f.number,
236 name: FastStr::new(f.name()).into(),
237 fields: vec![self.lower_ty(
238 f.type_,
239 f.type_name.as_deref(),
240 &nested_messages,
241 Some(message.name()),
242 )],
243 tags: Default::default(),
244 })
245 .collect_vec(),
246 }),
247 }));
248
249 extra_fields.push((
250 fields[0].0,
251 ir::Field {
252 name: FastStr::new(d.name()).into(),
253 id: -1,
254 ty: ir::Ty {
255 kind: ir::TyKind::Path(Path {
256 segments: Arc::from([
257 FastStr::new(message.name()).into(),
258 FastStr::new(d.name()).into(),
259 ]),
260 }),
261 tags: Default::default(),
262 },
263 tags: Arc::new(crate::tags!(OneOf)),
264 kind: ir::FieldKind::Optional,
265 default: None,
266 },
267 ));
268 }
269 });
270
271 parent_messages.push(message.name().into());
272
273 nested_messages
274 .iter()
275 .filter(|(_, m)| !m.options.has_map_entry())
276 .for_each(|(_, m)| {
277 self.lower_message(m, parent_messages)
278 .into_iter()
279 .for_each(|item| nested_items.push(Arc::new(item)))
280 });
281
282 parent_messages.pop();
283
284 message
285 .enum_type
286 .iter()
287 .for_each(|e| nested_items.push(Arc::new(self.lower_enum(e))));
288
289 let item = ir::Item {
290 related_items: Default::default(),
291 tags: Default::default(),
292 kind: ir::ItemKind::Message(ir::Message {
293 fields: fields
294 .iter()
295 .map(|(idx, f)| {
296 let mut ty = self.lower_ty(
297 f.type_,
298 f.type_name.as_deref(),
299 &nested_messages,
300 Some(message.name()),
301 );
302
303 let is_map = matches!(ty.kind, TyKind::Map(_, _));
304 let repeated = !is_map && matches!(f.label(), Label::LABEL_REPEATED);
305
306 if repeated {
307 ty = ir::Ty {
308 kind: ir::TyKind::Vec(Arc::from(ty)),
309 tags: Default::default(),
310 }
311 }
312
313 let optional = (|| {
314 if is_map {
315 return false;
316 }
317
318 match self.cur_syntax {
319 Syntax::Proto3 => {
320 f.proto3_optional()
321 || (!repeated && matches!(f.type_(), Type::TYPE_MESSAGE))
322 }
323 Syntax::Proto2 => f.label() == Label::LABEL_OPTIONAL,
324 }
325 })();
326
327 let mut tags = Tags::default();
328 if repeated {
329 tags.insert(Repeated);
330 }
331
332 (
333 *idx,
334 ir::Field {
335 default: None,
336 id: f.number(),
337 name: FastStr::new(f.name()).into(),
338 ty,
339 tags: Arc::new(tags),
340 kind: if optional {
341 FieldKind::Optional
342 } else {
343 FieldKind::Required
344 },
345 },
346 )
347 })
348 .chain(extra_fields)
349 .sorted_unstable_by_key(|(idx, _)| *idx)
350 .map(|(_, f)| f)
351 .collect(),
352 name: FastStr::new(message.name()).into(),
353 }),
354 };
355
356 if nested_items.is_empty() {
357 vec![item]
358 } else {
359 let name = item.name();
360 let mut tags = Tags::default();
361 tags.insert(PilotaName(name.0.mod_ident()));
362 vec![
363 item,
364 Item {
365 related_items: Default::default(),
366 tags: Arc::from(tags),
367 kind: ir::ItemKind::Mod(ir::Mod {
368 name: Ident { sym: name },
369 items: nested_items,
370 }),
371 },
372 ]
373 }
374 }
375
376 pub fn lower_service(&self, service: &ServiceDescriptorProto) -> ir::Item {
377 ir::Item {
378 tags: Default::default(),
379 related_items: Default::default(),
380 kind: ir::ItemKind::Service(ir::Service {
381 name: FastStr::new(service.name()).into(),
382 methods: service
383 .method
384 .iter()
385 .map(|m| {
386 let mut tags = Tags::default();
387 if m.client_streaming() {
388 tags.insert(ClientStreaming);
389 }
390 if m.server_streaming() {
391 tags.insert(ServerStreaming);
392 }
393 ir::Method {
394 name: FastStr::new(m.name()).into(),
395 tags: Arc::new(tags),
396 args: vec![ir::Arg {
397 name: "req".into(),
398 id: -1,
399 ty: self.lower_ty(
400 None,
401 m.input_type.as_deref(),
402 &Default::default(),
403 Some(service.name()),
404 ),
405 tags: Arc::new(Tags::default()),
406 }],
407 oneway: false,
408 ret: self.lower_ty(
409 None,
410 m.output_type.as_deref(),
411 &Default::default(),
412 Some(service.name()),
413 ),
414 exceptions: None,
415 }
416 })
417 .collect_vec(),
418 extend: vec![],
419 }),
420 }
421 }
422
423 pub fn lower(
424 &mut self,
425 files: &[protobuf::descriptor::FileDescriptorProto],
426 ) -> Vec<Arc<ir::File>> {
427 let mut file_map = HashMap::with_capacity(files.len());
428 files.iter().for_each(|f| {
429 self.files
430 .insert(f.name().to_string(), self.next_file_id.inc_one());
431 file_map.insert(f.name(), f);
432 });
433
434 files
435 .iter()
436 .map(|f| {
437 self.cur_package = f.package.clone();
438 self.cur_syntax = match f.syntax() {
439 "proto3" => Syntax::Proto3,
440 _ => Syntax::Proto2,
441 };
442
443 let file_id = *self.files.get(f.name()).unwrap();
444
445 let package = self.str2path(f.package());
446
447 let enums = f.enum_type.iter().map(|e| self.lower_enum(e));
448 let messages = f
449 .message_type
450 .iter()
451 .map(|m| self.lower_message(m, &mut Vec::new()));
452 let services = f.service.iter().map(|s| self.lower_service(s));
453
454 let f = Arc::from(ir::File {
455 package,
456 uses: f
457 .dependency
458 .iter()
459 .map(|d| {
460 (
461 self.str2path(file_map.get(&**d).unwrap().package()),
462 *self.files.get(d).unwrap(),
463 )
464 })
465 .collect(),
466 id: file_id,
467 items: messages
468 .flatten()
469 .chain(enums)
470 .chain(services)
471 .map(Arc::from)
472 .collect::<Vec<_>>(),
473 });
474
475 self.cur_package = None;
476
477 f
478 })
479 .collect::<Vec<_>>()
480 }
481}
482
483impl ProtobufParser {
484 pub fn parse_and_typecheck(&self) -> (Vec<protobuf::descriptor::FileDescriptorProto>, super::ParseResult) {
485 let descriptors = self.inner.parse_and_typecheck().unwrap().file_descriptors;
486 let mut input_file_ids = vec![];
487
488 let mut lower = Lower::default();
489
490 let files = lower.lower(&descriptors);
491
492 let mut file_ids = FxHashMap::default();
493
494 descriptors.iter().for_each(|f| {
495 self.include_dirs.iter().for_each(|p| {
496 let path = p.join(f.name());
497 if path.exists() {
498 println!("cargo:rerun-if-changed={}", path.display());
499 file_ids.insert(
500 Arc::from(path.normalize().unwrap().into_path_buf()),
501 *lower.files.get(f.name()).unwrap(),
502 );
503 if self
504 .input_files
505 .contains(path.normalize().unwrap().as_path())
506 {
507 input_file_ids.push(*lower.files.get(f.name()).unwrap());
508 }
509 }
510 });
511 });
512
513 (descriptors, super::ParseResult {
514 files,
515 input_files: input_file_ids,
516 file_ids_map: file_ids,
517 })
518 }
519}
520
521impl Parser for ProtobufParser {
522 fn input<P: AsRef<std::path::Path>>(&mut self, path: P) {
523 let p = path.as_ref();
524 self.input_files.insert(
525 p.normalize()
526 .unwrap_or_else(|_| panic!("normalize path failed: {}", p.display()))
527 .into_path_buf(),
528 );
529 self.inner.input(path);
530 }
531
532 fn include_dirs(&mut self, dirs: Vec<std::path::PathBuf>) {
533 self.include_dirs = dirs.clone();
534 self.inner.includes(dirs);
535 }
536
537 fn parse(self) -> super::ParseResult {
538 let (_, ret) = self.parse_and_typecheck();
539 ret
540 }
541}