From 9f3823ba2a862b58e6810149989a38fb1dc13aa6 Mon Sep 17 00:00:00 2001 From: LIU JIE Date: Mon, 25 Sep 2023 20:30:40 +0800 Subject: [PATCH] fix union empty encode (#195) --- Cargo.lock | 2 +- pilota-build/Cargo.toml | 2 +- pilota-build/src/codegen/thrift/mod.rs | 18 ++-- pilota-build/test_data/thrift/union.rs | 95 ++++++++++++++++++++++ pilota-build/test_data/thrift/union.thrift | 4 + 5 files changed, 110 insertions(+), 11 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 02849676..6c469cd2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -738,7 +738,7 @@ dependencies = [ [[package]] name = "pilota-build" -version = "0.8.4" +version = "0.8.5" dependencies = [ "anyhow", "async-trait", diff --git a/pilota-build/Cargo.toml b/pilota-build/Cargo.toml index 6402e498..cdde7e2a 100644 --- a/pilota-build/Cargo.toml +++ b/pilota-build/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "pilota-build" -version = "0.8.4" +version = "0.8.5" edition = "2021" description = "Compile thrift and protobuf idl into rust code at compile-time." documentation = "https://docs.rs/pilota-build" diff --git a/pilota-build/src/codegen/thrift/mod.rs b/pilota-build/src/codegen/thrift/mod.rs index 1e3f6948..b1e3b05e 100644 --- a/pilota-build/src/codegen/thrift/mod.rs +++ b/pilota-build/src/codegen/thrift/mod.rs @@ -536,6 +536,10 @@ impl CodegenBackend for ThriftBackend { }); } + if e.variants.is_empty() { + encode_variants.push_str("_ => {},"); + } + let mut variants_size = e .variants .iter() @@ -560,6 +564,10 @@ impl CodegenBackend for ThriftBackend { }) } + if e.variants.is_empty() { + variants_size.push_str("_ => 0,"); + } + let variant_is_void = |v: &EnumVariant| { &*v.name.sym == "Ok" && v.fields.len() == 1 && v.fields[0].kind == TyKind::Void }; @@ -577,21 +585,13 @@ impl CodegenBackend for ThriftBackend { protocol.write_struct_end()?; Ok(())"# }, - if !e.variants.is_empty() { format! { r#"protocol.struct_begin_len(&::pilota::thrift::TStructIdentifier {{ name: "{name_str}", }}) + match self {{ {variants_size} }} + protocol.field_stop_len() + protocol.struct_end_len()"# - } - } else { - format! { - r#"protocol.struct_begin_len(&::pilota::thrift::TStructIdentifier {{ - name: "{name_str}", - }}) + protocol.field_stop_len() + protocol.struct_end_len()"# - } - }, + }, |helper| { let record_ptr = if keep && !helper.is_async { r#"let mut __pilota_offset = 0; diff --git a/pilota-build/test_data/thrift/union.rs b/pilota-build/test_data/thrift/union.rs index 8682efdc..2d7cf323 100644 --- a/pilota-build/test_data/thrift/union.rs +++ b/pilota-build/test_data/thrift/union.rs @@ -158,6 +158,101 @@ pub mod union { + protocol.struct_end_len() } } + #[derive(PartialOrd, Hash, Eq, Ord, Debug, Clone, PartialEq)] + + pub enum Empty {} + + #[::async_trait::async_trait] + impl ::pilota::thrift::Message for Empty { + fn encode( + &self, + protocol: &mut T, + ) -> ::std::result::Result<(), ::pilota::thrift::EncodeError> { + #[allow(unused_imports)] + use ::pilota::thrift::TOutputProtocolExt; + protocol + .write_struct_begin(&::pilota::thrift::TStructIdentifier { name: "Empty" })?; + match self { + _ => {} + } + protocol.write_field_stop()?; + protocol.write_struct_end()?; + Ok(()) + } + + fn decode( + protocol: &mut T, + ) -> ::std::result::Result { + #[allow(unused_imports)] + use ::pilota::{thrift::TLengthProtocolExt, Buf}; + let mut ret = None; + protocol.read_struct_begin()?; + loop { + let field_ident = protocol.read_field_begin()?; + if field_ident.field_type == ::pilota::thrift::TType::Stop { + protocol.field_stop_len(); + break; + } else { + protocol.field_begin_len(field_ident.field_type, field_ident.id); + } + match field_ident.id { + _ => { + protocol.skip(field_ident.field_type)?; + } + } + } + protocol.read_field_end()?; + protocol.read_struct_end()?; + if let Some(ret) = ret { + Ok(ret) + } else { + Err(::pilota::thrift::DecodeError::new( + ::pilota::thrift::DecodeErrorKind::InvalidData, + "received empty union from remote Message", + )) + } + } + + async fn decode_async( + protocol: &mut T, + ) -> ::std::result::Result { + let mut ret = None; + protocol.read_struct_begin().await?; + loop { + let field_ident = protocol.read_field_begin().await?; + if field_ident.field_type == ::pilota::thrift::TType::Stop { + break; + } else { + } + match field_ident.id { + _ => { + protocol.skip(field_ident.field_type).await?; + } + } + } + protocol.read_field_end().await?; + protocol.read_struct_end().await?; + if let Some(ret) = ret { + Ok(ret) + } else { + Err(::pilota::thrift::DecodeError::new( + ::pilota::thrift::DecodeErrorKind::InvalidData, + "received empty union from remote Message", + )) + } + } + + fn size(&self, protocol: &mut T) -> usize { + #[allow(unused_imports)] + use ::pilota::thrift::TLengthProtocolExt; + protocol.struct_begin_len(&::pilota::thrift::TStructIdentifier { name: "Empty" }) + + match self { + _ => 0, + } + + protocol.field_stop_len() + + protocol.struct_end_len() + } + } #[derive(PartialOrd, Hash, Eq, Ord, Debug, Default, Clone, PartialEq)] pub struct A { pub u: ::std::option::Option, diff --git a/pilota-build/test_data/thrift/union.thrift b/pilota-build/test_data/thrift/union.thrift index 8f2a2a5a..a1bbd6b8 100644 --- a/pilota-build/test_data/thrift/union.thrift +++ b/pilota-build/test_data/thrift/union.thrift @@ -5,4 +5,8 @@ union Union { struct A { 1: Union u, +} + +union Empty { + } \ No newline at end of file