Skip to content

Commit

Permalink
add support for proto2 groups
Browse files Browse the repository at this point in the history
  • Loading branch information
kristopherbullinger committed Oct 2, 2024
1 parent 3ba69de commit 6927801
Show file tree
Hide file tree
Showing 12 changed files with 252 additions and 44 deletions.
8 changes: 5 additions & 3 deletions protobuf-codegen/src/gen/field/accessor.rs
Original file line number Diff line number Diff line change
Expand Up @@ -155,9 +155,11 @@ impl FieldGen<'_> {
type_params: vec!["_".to_owned()],
callback_params: self.make_accessor_fns_lambda(),
},
FieldElem::Group => {
unreachable!("no accessor for group field");
}
FieldElem::Group(m) => AccessorFn {
name: "make_message_field_accessor".to_owned(),
type_params: vec![format!("{}", m.rust_name_relative(&self.file_and_mod()))],
callback_params: self.make_accessor_fns_lambda(),
},
}
}

Expand Down
30 changes: 22 additions & 8 deletions protobuf-codegen/src/gen/field/elem.rs
Original file line number Diff line number Diff line change
Expand Up @@ -74,7 +74,7 @@ pub(crate) enum FieldElem<'a> {
Primitive(Type, PrimitiveTypeVariant),
Message(FieldElemMessage<'a>),
Enum(FieldElemEnum<'a>),
Group,
Group(FieldElemMessage<'a>),
}

pub(crate) enum HowToGetMessageSize {
Expand All @@ -86,7 +86,7 @@ impl<'a> FieldElem<'a> {
pub(crate) fn proto_type(&self) -> Type {
match *self {
FieldElem::Primitive(t, ..) => t,
FieldElem::Group => Type::TYPE_GROUP,
FieldElem::Group(_) => Type::TYPE_GROUP,
FieldElem::Message(..) => Type::TYPE_MESSAGE,
FieldElem::Enum(..) => Type::TYPE_ENUM,
}
Expand All @@ -106,7 +106,7 @@ impl<'a> FieldElem<'a> {
RustType::Bytes
}
FieldElem::Primitive(.., PrimitiveTypeVariant::TokioBytes) => unreachable!(),
FieldElem::Group => RustType::Group,
FieldElem::Group(ref m) => m.rust_type(reference),
FieldElem::Message(ref m) => m.rust_type(reference),
FieldElem::Enum(ref en) => en.enum_or_unknown_rust_type(reference),
}
Expand Down Expand Up @@ -213,6 +213,18 @@ impl<'a> FieldElem<'a> {
protobuf_crate_path(customize),
));
}
Type::TYPE_GROUP => {
match how_to_get_message_size {
HowToGetMessageSize::Compute => {
w.write_line(&format!("let len = {}.compute_size();", item_var.value))
}
HowToGetMessageSize::GetCached => w.write_line(&format!(
"let len = {}.cached_size() as u64;",
item_var.value
)),
}
w.write_line(&format!("{sum_var} += {tag_size} * 2 + len;",));
}
_ => {
w.write_line(&format!(
"{sum_var} += {};",
Expand All @@ -232,7 +244,7 @@ impl<'a> FieldElem<'a> {
w: &mut CodeWriter,
) {
match self.proto_type() {
Type::TYPE_MESSAGE => {
Type::TYPE_MESSAGE | Type::TYPE_GROUP => {
let param_type = RustType::Ref(Box::new(self.rust_storage_elem_type(file_and_mod)));

w.write_line(&format!(
Expand Down Expand Up @@ -273,10 +285,7 @@ pub(crate) fn field_elem<'a>(
if let RuntimeFieldType::Map(..) = field.field.runtime_field_type() {
unreachable!();
}

if field.field.proto().type_() == Type::TYPE_GROUP {
FieldElem::Group
} else if field.field.proto().has_type_name() {
if field.field.proto().has_type_name() {
let message_or_enum = root_scope
.find_message_or_enum(&ProtobufAbsPath::from(field.field.proto().type_name()));
match (field.field.proto().type_(), message_or_enum) {
Expand All @@ -285,6 +294,11 @@ pub(crate) fn field_elem<'a>(
message: message.clone(),
})
}
(Type::TYPE_GROUP, MessageOrEnumWithScope::Message(message)) => {
FieldElem::Group(FieldElemMessage {
message: message.clone(),
})
}
(Type::TYPE_ENUM, MessageOrEnumWithScope::Enum(enum_with_scope)) => {
let default_value = if field.field.proto().has_default_value() {
enum_with_scope.value_by_name(field.field.proto().default_value())
Expand Down
83 changes: 64 additions & 19 deletions protobuf-codegen/src/gen/field/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -603,26 +603,22 @@ impl<'a> FieldGen<'a> {
}

pub(crate) fn write_struct_field(&self, w: &mut CodeWriter) {
if self.proto_type == field_descriptor_proto::Type::TYPE_GROUP {
w.comment(&format!("{}: <group>", &self.rust_name));
} else {
w.all_documentation(self.info, &self.path);
w.all_documentation(self.info, &self.path);

write_protoc_insertion_point_for_field(w, &self.customize, &self.proto_field.field);
w.field_decl_vis(
Visibility::Public,
&self.rust_name.to_string(),
&self
.full_storage_type(
&self
.proto_field
.message
.scope
.file_and_mod(self.customize.clone()),
)
.to_code(&self.customize),
);
}
write_protoc_insertion_point_for_field(w, &self.customize, &self.proto_field.field);
w.field_decl_vis(
Visibility::Public,
&self.rust_name.to_string(),
&self
.full_storage_type(
&self
.proto_field
.message
.scope
.file_and_mod(self.customize.clone()),
)
.to_code(&self.customize),
);
}

fn write_if_let_self_field_is_some<F>(&self, s: &SingularField, w: &mut CodeWriter, cb: F)
Expand Down Expand Up @@ -1035,13 +1031,58 @@ impl<'a> FieldGen<'a> {
self.rust_name,
));
}
FieldElem::Group(..) => self.write_merge_delimited_group_case_block(w),
_ => {
let read_proc = s.elem.read_one_liner();
self.write_self_field_assign_some(w, s, &read_proc);
}
})
}

fn write_merge_delimited_repeated_group_case_block(&self, w: &mut CodeWriter<'_>) {
let name = &self.rust_name;
let proto_path = protobuf_crate_path(&self.customize);
w.write_line(&format!(
"let end_tag = {proto_path}::rt::set_wire_type_to_end_group(tag);"
));
let value_type = &self
.elem()
.rust_storage_elem_type(
&self
.proto_field
.message
.scope
.file_and_mod(self.customize.clone()),
)
.to_code(&self.customize);

w.write_line(&format!("let mut {name} = {value_type}::default();"));
w.write_line(&format!("{name}.merge_delimited(is, end_tag)?;"));
w.write_line(&format!("self.{name}.push({name});",));
}

fn write_merge_delimited_group_case_block(&self, w: &mut CodeWriter<'_>) {
let name = &self.rust_name;
let proto_path = protobuf_crate_path(&self.customize);
w.write_line(&format!(
"let end_tag = {proto_path}::rt::set_wire_type_to_end_group(tag);"
));
let value_type = &self
.elem()
.rust_storage_elem_type(
&self
.proto_field
.message
.scope
.file_and_mod(self.customize.clone()),
)
.to_code(&self.customize);

w.write_line(&format!("let mut {name} = {value_type}::default();"));
w.write_line(&format!("{name}.merge_delimited(is, end_tag)?;"));
w.write_line(&format!("self.{name} = Some({name});",));
}

// Write `merge_from` part for this repeated field
fn write_merge_from_repeated_case_block(&self, w: &mut CodeWriter) {
let field = match self.kind {
Expand All @@ -1057,6 +1098,9 @@ impl<'a> FieldGen<'a> {
self.write_merge_from_field_message_string_bytes_repeated(field, w);
})
}
FieldElem::Group(_) => w.case_block(&format!("{}", self.tag()), |w| {
self.write_merge_delimited_repeated_group_case_block(w);
}),
FieldElem::Enum(..) => {
w.case_block(
&format!("{}", self.tag_with_wire_type(WireType::Varint)),
Expand Down Expand Up @@ -1338,6 +1382,7 @@ impl<'a> FieldGen<'a> {

match singular.elem {
FieldElem::Message(..) => self.write_message_field_get_singular_message(singular, w),
FieldElem::Group(..) => self.write_message_field_get_singular_message(singular, w),
FieldElem::Enum(ref en) => {
self.write_message_field_get_singular_enum(singular.flag, en, w)
}
Expand Down
2 changes: 1 addition & 1 deletion protobuf-codegen/src/gen/field/type_ext.rs
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@ impl TypeExt for Type {

fn is_copy(&self) -> bool {
match self {
Type::TYPE_MESSAGE | Type::TYPE_STRING | Type::TYPE_BYTES => false,
Type::TYPE_GROUP | Type::TYPE_MESSAGE | Type::TYPE_STRING | Type::TYPE_BYTES => false,
_ => true,
}
}
Expand Down
38 changes: 33 additions & 5 deletions protobuf-codegen/src/gen/message.rs
Original file line number Diff line number Diff line change
Expand Up @@ -271,7 +271,7 @@ impl<'a> MessageGen<'a> {
self.rust_name()
),
|w| {
for f in &self.fields_except_oneof_and_group() {
for f in self.fields_except_oneof() {
w.field_entry(
&f.rust_name.to_string(),
&f.kind
Expand Down Expand Up @@ -310,6 +310,31 @@ impl<'a> MessageGen<'a> {
);
}

fn write_merge_delimited(&self, w: &mut CodeWriter) {
let sig = format!(
"merge_delimited(&mut self, is: &mut {}::CodedInputStream<'_>, end_tag: u32) -> {}::Result<()>",
protobuf_crate_path(&self.customize.for_elem),
protobuf_crate_path(&self.customize.for_elem),
);
w.comment("read and merge values from the stream, stopping when `end_tag` tag is reached");
w.pub_fn(&sig, |w| {
w.while_block("let Some(tag) = is.read_raw_tag_or_eof()?", |w| {
w.if_stmt("tag == end_tag", |w| {
w.write_line("return ::std::result::Result::Ok(());");
});
w.match_block("tag", |w| {
for f in &self.fields {
f.write_merge_from_field_case_block(w);
}
w.case_block("tag", |w| {
w.write_line(&format!("{}::rt::read_unknown(tag, is, self.special_fields.mut_unknown_fields())?;", protobuf_crate_path(&self.customize.for_elem)));
});
});
});
w.write_line("::std::result::Result::Ok(())");
});
}

fn write_compute_size(&self, w: &mut CodeWriter) {
// Append sizes of messages in the tree to the specified vector.
// First appended element is size of self, and then nested message sizes.
Expand All @@ -320,7 +345,7 @@ impl<'a> MessageGen<'a> {
w.def_fn("compute_size(&self) -> u64", |w| {
// To have access to its methods but not polute the name space.
w.write_line("let mut my_size = 0;");
for field in self.fields_except_oneof_and_group() {
for field in &self.fields_except_oneof() {
field.write_message_compute_field_size("my_size", w);
}
self.write_match_each_oneof_variant(w, |w, variant, v| {
Expand All @@ -338,7 +363,7 @@ impl<'a> MessageGen<'a> {
}

fn write_field_accessors(&self, w: &mut CodeWriter) {
for f in self.fields_except_group() {
for f in &self.fields {
f.write_message_single_field_accessors(w);
}
}
Expand All @@ -351,6 +376,9 @@ impl<'a> MessageGen<'a> {

self.write_field_accessors(w);

w.write_line("");
self.write_merge_delimited(w);

if !self.lite_runtime {
w.write_line("");
self.write_generated_message_descriptor_data(w);
Expand Down Expand Up @@ -385,11 +413,11 @@ impl<'a> MessageGen<'a> {
w.def_fn(&sig, |w| {
w.while_block("let Some(tag) = is.read_raw_tag_or_eof()?", |w| {
w.match_block("tag", |w| {
for f in &self.fields_except_group() {
for f in &self.fields {
f.write_merge_from_field_case_block(w);
}
w.case_block("tag", |w| {
w.write_line(&format!("{}::rt::read_unknown_or_skip_group(tag, is, self.special_fields.mut_unknown_fields())?;", protobuf_crate_path(&self.customize.for_elem)));
w.write_line(&format!("{}::rt::read_unknown(tag, is, self.special_fields.mut_unknown_fields())?;", protobuf_crate_path(&self.customize.for_elem)));
});
});
});
Expand Down
3 changes: 0 additions & 3 deletions protobuf-codegen/src/gen/rust_types_values.rs
Original file line number Diff line number Diff line change
Expand Up @@ -58,8 +58,6 @@ pub(crate) enum RustType {
Bytes,
// chars::Chars
Chars,
// group
Group,
}

impl RustType {
Expand Down Expand Up @@ -101,7 +99,6 @@ impl RustType {
protobuf_crate_path(customize),
name
),
RustType::Group => format!("<group>"),
RustType::Bytes => format!("::bytes::Bytes"),
RustType::Chars => format!("{}::Chars", protobuf_crate_path(customize)),
}
Expand Down
Loading

0 comments on commit 6927801

Please sign in to comment.