fix union empty encode (#195)

cloudwego · Sep 25, 2023 · 9f3823b · 9f3823b
1 parent 541ffb0
commit 9f3823b
Show file tree

Hide file tree

Showing 5 changed files with 110 additions and 11 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git 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
@@ -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
@@ -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<T: ::pilota::thrift::TOutputProtocol>(
+                &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<T: ::pilota::thrift::TInputProtocol>(
+                protocol: &mut T,
+            ) -> ::std::result::Result<Self, ::pilota::thrift::DecodeError> {
+                #[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<T: ::pilota::thrift::TAsyncInputProtocol>(
+                protocol: &mut T,
+            ) -> ::std::result::Result<Self, ::pilota::thrift::DecodeError> {
+                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<T: ::pilota::thrift::TLengthProtocol>(&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<Union>,

diff --git 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 {
+
 }