Skip to content

Commit ce49069

Browse files
FoyonaCZYGgiggle
FoyonaCZY
authored and
Ggiggle
committed
feat: implement methods to get descriptor for pb nested items
1 parent 0a018a3 commit ce49069

File tree

20 files changed

+731
-223
lines changed

20 files changed

+731
-223
lines changed

‎pilota-build/src/codegen/mod.rs‎

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -372,11 +372,6 @@ where
372372
"#
373373
});
374374

375-
if self.config.with_descriptor {
376-
self.backend
377-
.codegen_enum_descrptor_getter_impl(def_id, stream, e);
378-
}
379-
380375
self.backend.codegen_enum_impl(def_id, stream, e);
381376
}
382377

‎pilota-build/src/codegen/pb/mod.rs‎

Lines changed: 94 additions & 76 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ use crate::{
1111
middle::{
1212
context::Mode,
1313
ext::{
14-
FileExts, ModExts,
14+
FileExts, ItemExts, ModExts,
1515
pb::{Extendee, ExtendeeKind, Extendees, FieldType},
1616
},
1717
ty::{self},
@@ -422,7 +422,7 @@ impl CodegenBackend for ProtobufBackend {
422422
const PROTOCOL: &'static str = "protobuf";
423423

424424
fn codegen_struct_impl(&self, def_id: DefId, stream: &mut String, s: &rir::Message) {
425-
let idl_name = s.name.sym.0.clone();
425+
let idl_name = s.name.raw_str();
426426
let name = self.cx.rust_name(def_id);
427427
let mut encoded_len = s
428428
.fields
@@ -546,25 +546,45 @@ impl CodegenBackend for ProtobufBackend {
546546

547547
if self.cx.config.with_descriptor {
548548
let file_id = self.cx.node(def_id).unwrap().file_id;
549-
if self.cx.file_paths().get(&file_id).is_some() {
550-
let filename = self.file_name(file_id).unwrap().replace(".", "_");
551-
let filename_lower = filename.to_lowercase();
549+
let mut getter_impl = String::new();
550+
if let ItemExts::Pb(pb) = &s.item_exts {
551+
if let Some(p) = &pb.parent {
552+
if self.cx.file_paths().get(&file_id).is_some() {
553+
let path = self.cx.related_item_path(def_id, p.did);
554+
let name = self.cx.rust_name(def_id);
555+
let idl_name = s.name.raw_str();
556+
getter_impl = format!(
557+
r#"
558+
impl {name} {{
559+
pub fn get_descriptor_proto() -> Option<&'static ::pilota::pb::descriptor::DescriptorProto> {{
560+
let message_descriptor = {path}::get_descriptor_proto()?;
561+
message_descriptor.get_message_descriptor_proto("{idl_name}")
562+
}}
563+
}}
564+
"#
565+
);
566+
};
567+
} else if self.cx.file_paths().get(&file_id).is_some() {
568+
let filename = self.file_name(file_id).unwrap().replace(".", "_");
569+
let filename_lower = filename.to_lowercase();
552570

553-
let file = &self.cx.files().get(&file_id).unwrap().package;
554-
let path = self.cx.item_path(def_id);
555-
let super_mods = "super::".repeat(path.len() - file.len() - 1);
571+
let file = &self.cx.files().get(&file_id).unwrap().package;
572+
let path = self.cx.item_path(def_id);
573+
let super_mods = "super::".repeat(path.len() - file.len() - 1);
556574

557-
stream.push_str(&format!(
558-
r#"
559-
impl {name} {{
560-
fn get_descriptor_proto() -> Option<&'static ::pilota::pb::descriptor::DescriptorProto> {{
561-
let file_descriptor = {super_mods}file_descriptor_proto_{filename_lower}();
562-
file_descriptor.get_message_descriptor_proto("{idl_name}")
563-
}}
564-
}}
565-
"#
566-
));
575+
getter_impl = format!(
576+
r#"
577+
impl {name} {{
578+
pub fn get_descriptor_proto() -> Option<&'static ::pilota::pb::descriptor::DescriptorProto> {{
579+
let file_descriptor = {super_mods}file_descriptor_proto_{filename_lower}();
580+
file_descriptor.get_message_descriptor_proto("{idl_name}")
581+
}}
582+
}}
583+
"#
584+
);
585+
}
567586
}
587+
stream.push_str(&getter_impl);
568588
}
569589

570590
stream.push_str(&format!(
@@ -612,6 +632,51 @@ impl CodegenBackend for ProtobufBackend {
612632
fn codegen_enum_impl(&self, def_id: DefId, stream: &mut String, e: &rir::Enum) {
613633
let node = self.cx.node(def_id).unwrap();
614634
if !self.cx.contains_tag::<OneOf>(node.tags) {
635+
if self.cx.config.with_descriptor {
636+
let file_id = self.cx.node(def_id).unwrap().file_id;
637+
let mut getter_impl = String::new();
638+
639+
if let ItemExts::Pb(pb) = &e.item_exts {
640+
if let Some(p) = &pb.parent {
641+
if self.cx.file_paths().get(&file_id).is_some() {
642+
let path = self.cx.related_item_path(def_id, p.did);
643+
let name = self.cx.rust_name(def_id);
644+
let idl_name = e.name.raw_str();
645+
getter_impl = format!(
646+
r#"
647+
impl {name} {{
648+
pub fn get_descriptor_proto() -> Option<&'static ::pilota::pb::descriptor::EnumDescriptorProto> {{
649+
let message_descriptor = {path}::get_descriptor_proto()?;
650+
message_descriptor.get_enum_descriptor_proto("{idl_name}")
651+
}}
652+
}}
653+
"#
654+
);
655+
};
656+
} else if self.cx.file_paths().get(&file_id).is_some() {
657+
let name = self.cx.rust_name(def_id);
658+
let idl_name = e.name.raw_str();
659+
let filename = self.file_name(file_id).unwrap().replace(".", "_");
660+
let filename_lower = filename.to_lowercase();
661+
662+
let file = &self.cx.files().get(&file_id).unwrap().package;
663+
let path = self.cx.item_path(def_id);
664+
let super_mods = "super::".repeat(path.len() - file.len() - 1);
665+
666+
getter_impl = format!(
667+
r#"
668+
impl {name} {{
669+
pub fn get_descriptor_proto() -> Option<&'static ::pilota::pb::descriptor::EnumDescriptorProto> {{
670+
let file_descriptor = {super_mods}file_descriptor_proto_{filename_lower}();
671+
file_descriptor.get_enum_descriptor_proto("{idl_name}")
672+
}}
673+
}}
674+
"#
675+
);
676+
}
677+
}
678+
stream.push_str(&getter_impl);
679+
}
615680
return;
616681
}
617682
let name = self.cx.rust_name(def_id);
@@ -672,31 +737,25 @@ impl CodegenBackend for ProtobufBackend {
672737
}).join("");
673738

674739
let file_id = self.cx.node(def_id).unwrap().file_id;
675-
let mut getter_impl = match &e.oneof_parent {
676-
Some(ty) => match &self.cx.file_paths().get(&file_id) {
677-
Some(_) => {
678-
let path = self.codegen_item_ty(ty.kind.clone());
679-
let path = format!("{path}");
740+
let mut getter_impl = String::new();
741+
if let ItemExts::Pb(pb) = &e.item_exts {
742+
if let Some(p) = &pb.parent {
743+
if self.cx.file_paths().get(&file_id).is_some() {
744+
let path = self.cx.related_item_path(def_id, p.did);
680745
let name = self.cx.rust_name(def_id);
681-
let idl_name = e.name.sym.0.clone();
682-
format!(
746+
let idl_name = e.name.raw_str();
747+
getter_impl = format!(
683748
r#"
684749
impl {name} {{
685-
fn get_descriptor_proto() -> Option<&'static ::pilota::pb::descriptor::OneofDescriptorProto> {{
750+
pub fn get_descriptor_proto() -> Option<&'static ::pilota::pb::descriptor::OneofDescriptorProto> {{
686751
let message_descriptor = {path}::get_descriptor_proto()?;
687752
message_descriptor.get_oneof_descriptor_proto("{idl_name}")
688753
}}
689754
}}
690755
"#
691-
)
692-
}
693-
_ => "".to_string(),
694-
},
695-
_ => "".to_string(),
696-
};
697-
698-
if !self.config.with_descriptor {
699-
getter_impl.clear();
756+
);
757+
};
758+
}
700759
}
701760

702761
stream.push_str(&format! {
@@ -736,47 +795,6 @@ impl CodegenBackend for ProtobufBackend {
736795
});
737796
}
738797

739-
fn codegen_enum_descrptor_getter_impl(
740-
&self,
741-
def_id: DefId,
742-
stream: &mut String,
743-
e: &rir::Enum,
744-
) {
745-
if !self.cx.config.with_descriptor {
746-
return;
747-
}
748-
749-
let file_id = self.cx.node(def_id).unwrap().file_id;
750-
751-
let getter_impl = match self.cx.file_paths().get(&file_id) {
752-
Some(_) => {
753-
let name = self.cx.rust_name(def_id);
754-
let idl_name = e.name.sym.0.clone();
755-
let filename = self.file_name(file_id).unwrap().replace(".", "_");
756-
let filename_lower = filename.to_lowercase();
757-
758-
let file = &self.cx.files().get(&file_id).unwrap().package;
759-
let path = self.cx.item_path(def_id);
760-
let super_mods = "super::".repeat(path.len() - file.len() - 1);
761-
762-
format!(
763-
r#"
764-
765-
impl {name} {{
766-
fn get_descriptor_proto() -> Option<&'static ::pilota::pb::descriptor::EnumDescriptorProto> {{
767-
let file_descriptor = {super_mods}file_descriptor_proto_{filename_lower}();
768-
file_descriptor.get_enum_descriptor_proto("{idl_name}")
769-
}}
770-
}}
771-
772-
"#
773-
)
774-
}
775-
_ => "".to_string(),
776-
};
777-
stream.push_str(&getter_impl);
778-
}
779-
780798
fn cx(&self) -> &Context {
781799
&self.cx
782800
}

‎pilota-build/src/codegen/thrift/mod.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -622,7 +622,7 @@ impl CodegenBackend for ThriftBackend {
622622
));
623623

624624
if !s.is_wrapper && self.config.with_descriptor {
625-
let idl_name = s.name.sym.0.clone();
625+
let idl_name = s.name.raw_str();
626626
stream.push_str(&format! {
627627
r#"impl {name} {{
628628
pub fn get_descriptor() -> Option<&'static ::pilota_thrift_reflect::thrift_reflection::StructDescriptor> {{
@@ -720,7 +720,7 @@ impl CodegenBackend for ThriftBackend {
720720
})
721721
.join("");
722722

723-
let idl_name = s.name.sym.0.clone();
723+
let idl_name = s.name.raw_str();
724724
stream.push_str(&format! {
725725
r#"impl {name} {{
726726
pub fn get_descriptor() -> Option<&'static ::pilota_thrift_reflect::thrift_reflection::StructDescriptor> {{

‎pilota-build/src/codegen/traits.rs‎

Lines changed: 0 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -28,13 +28,6 @@ pub trait CodegenBackend: Clone {
2828
Default::default()
2929
}
3030
fn codegen_enum_impl(&self, _def_id: DefId, _stream: &mut String, _e: &rir::Enum) {}
31-
fn codegen_enum_descrptor_getter_impl(
32-
&self,
33-
_def_id: DefId,
34-
_stream: &mut String,
35-
_e: &rir::Enum,
36-
) {
37-
}
3831
fn codegen_newtype_impl(&self, _def_id: DefId, _stream: &mut String, _t: &rir::NewType) {}
3932
fn codegen_file_descriptor(&self, _stream: &mut String, _f: &rir::File, _has_direct: bool) {}
4033
fn codegen_register_mod_file_descriptor(

‎pilota-build/src/ir/ext/mod.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@ pub enum ItemExts {
5252
impl ItemExts {
5353
pub fn has_used_options(&self) -> bool {
5454
match self {
55-
ItemExts::Pb(pb::ItemExts { used_options }) => !used_options.is_empty(),
55+
ItemExts::Pb(pb::ItemExts { used_options, .. }) => !used_options.is_empty(),
5656
ItemExts::Thrift => false,
5757
}
5858
}

‎pilota-build/src/ir/ext/pb.rs‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,10 @@
11
use std::sync::Arc;
22

3-
use crate::{ir::Ty, parser::protobuf::WellKnownFileName, symbol::Ident};
3+
use crate::{
4+
ir::{Path, Ty},
5+
parser::protobuf::WellKnownFileName,
6+
symbol::Ident,
7+
};
48

59
/// Extension for protobuf custom options
610
/// As defined in [google/protobuf/descriptor.proto](https://github.com/protocolbuffers/protobuf/blob/main/src/google/protobuf/descriptor.proto#L124), the extendee is following the syntax of [FieldDescriptorProto](https://github.com/protocolbuffers/protobuf/blob/main/src/google/protobuf/descriptor.proto#L243)
@@ -139,9 +143,11 @@ impl ModExts {
139143

140144
/// The extension for item
141145
/// - used_options, the used options
146+
/// - parent, the parent path of nested message, enum and oneof
142147
#[derive(Clone, Debug)]
143148
pub struct ItemExts {
144149
pub used_options: UsedOptions,
150+
pub parent: Option<Path>,
145151
}
146152

147153
impl ItemExts {

‎pilota-build/src/middle/ext/mod.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,7 @@ pub enum ItemExts {
6666
impl ItemExts {
6767
pub fn has_used_options(&self) -> bool {
6868
match self {
69-
ItemExts::Pb(pb::ItemExts { used_options }) => !used_options.is_empty(),
69+
ItemExts::Pb(pb::ItemExts { used_options, .. }) => !used_options.is_empty(),
7070
ItemExts::Thrift => false,
7171
}
7272
}

‎pilota-build/src/middle/ext/pb.rs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
use std::sync::Arc;
22

3-
use crate::{ir, parser::protobuf::WellKnownFileName, symbol::Ident, ty::Ty};
3+
use crate::{ir, parser::protobuf::WellKnownFileName, rir::Path, symbol::Ident, ty::Ty};
44

55
#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
66
pub struct ExtendeeIndex {
@@ -123,4 +123,5 @@ pub struct ModExts {
123123
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
124124
pub struct ItemExts {
125125
pub used_options: UsedOptions,
126+
pub parent: Option<Path>,
126127
}

‎pilota-build/src/middle/rir.rs‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -135,7 +135,6 @@ pub struct Enum {
135135
pub leading_comments: FastStr,
136136
pub trailing_comments: FastStr,
137137
pub item_exts: ItemExts,
138-
pub oneof_parent: Option<Ty>,
139138
}
140139

141140
#[derive(Clone, Debug, PartialEq, Eq, Hash)]

0 commit comments

Comments
 (0)