Fix Newtype vs Tuple structs and variants

jamesmunns · Oct 11, 2024 · 1f87ee8 · 1f87ee8
1 parent 0788f30
commit 1f87ee8
Show file tree

Hide file tree

Showing 2 changed files with 74 additions and 13 deletions.
diff --git a/source/postcard-derive/src/schema.rs b/source/postcard-derive/src/schema.rs
@@ -69,13 +69,21 @@ fn generate_struct(fields: &Fields) -> TokenStream {
             ]) }
         }
         syn::Fields::Unnamed(fields) => {
-            let fields = fields.unnamed.iter().map(|f| {
+            if fields.unnamed.len() == 1 {
+                let f = fields.unnamed[0].clone();
                 let ty = &f.ty;
-                quote_spanned!(f.span() => <#ty as ::postcard::experimental::schema::Schema>::SCHEMA)
-            });
-            quote! { &::postcard::experimental::schema::SdmTy::TupleStruct(&[
-                #( #fields ),*
-            ]) }
+                let qs = quote_spanned!(f.span() => <#ty as ::postcard::experimental::schema::Schema>::SCHEMA);
+
+                quote! { &::postcard::experimental::schema::SdmTy::NewtypeStruct(#qs) }
+            } else {
+                let fields = fields.unnamed.iter().map(|f| {
+                    let ty = &f.ty;
+                    quote_spanned!(f.span() => <#ty as ::postcard::experimental::schema::Schema>::SCHEMA)
+                });
+                quote! { &::postcard::experimental::schema::SdmTy::TupleStruct(&[
+                    #( #fields ),*
+                ]) }
+            }
         }
         syn::Fields::Unit => {
             quote! { &::postcard::experimental::schema::SdmTy::UnitStruct }
@@ -96,13 +104,21 @@ fn generate_variants(fields: &Fields) -> TokenStream {
             ]) }
         }
         syn::Fields::Unnamed(fields) => {
-            let fields = fields.unnamed.iter().map(|f| {
+            if fields.unnamed.len() == 1 {
+                let f = fields.unnamed[0].clone();
                 let ty = &f.ty;
-                quote_spanned!(f.span() => <#ty as ::postcard::experimental::schema::Schema>::SCHEMA)
-            });
-            quote! { &::postcard::experimental::schema::SdmTy::TupleVariant(&[
-                #( #fields ),*
-            ]) }
+                let qs = quote_spanned!(f.span() => <#ty as ::postcard::experimental::schema::Schema>::SCHEMA);
+
+                quote! { &::postcard::experimental::schema::SdmTy::NewtypeVariant(#qs) }
+            } else {
+                let fields = fields.unnamed.iter().map(|f| {
+                    let ty = &f.ty;
+                    quote_spanned!(f.span() => <#ty as ::postcard::experimental::schema::Schema>::SCHEMA)
+                });
+                quote! { &::postcard::experimental::schema::SdmTy::TupleVariant(&[
+                    #( #fields ),*
+                ]) }
+            }
         }
         syn::Fields::Unit => {
             quote! { &::postcard::experimental::schema::SdmTy::UnitVariant }

diff --git a/source/postcard/tests/schema.rs b/source/postcard/tests/schema.rs
@@ -203,7 +203,7 @@ fn owned_punning() {
 
     // TODO: This is wildly repetitive, and likely could benefit from interning of
     // repeated types, strings, etc.
-    assert_eq!(ser_borrowed_schema.len(), 282);
+    assert_eq!(ser_borrowed_schema.len(), 280);
 
     // Check that we round-trip correctly
     let deser_borrowed_schema =
@@ -213,3 +213,48 @@ fn owned_punning() {
     assert_eq!(deser_borrowed_schema, owned_schema);
     assert_eq!(deser_owned_schema, owned_schema);
 }
+
+#[allow(unused)]
+#[derive(Debug, Schema)]
+struct TestStruct3(u64);
+
+#[allow(unused)]
+#[derive(Debug, Schema)]
+struct TestStruct4(u64, bool);
+
+#[allow(unused)]
+#[derive(Debug, Schema)]
+enum TestEnum2 {
+    Nt(u64),
+    Tup(u64, bool),
+}
+
+#[test]
+fn newtype_vs_tuple() {
+    assert_eq!(
+        TestStruct3::SCHEMA,
+        &NamedType {
+            name: "TestStruct3",
+            ty: &SdmTy::NewtypeStruct(u64::SCHEMA)
+        }
+    );
+
+    assert_eq!(
+        TestStruct4::SCHEMA,
+        &NamedType {
+            name: "TestStruct4",
+            ty: &SdmTy::TupleStruct(&[u64::SCHEMA, bool::SCHEMA]),
+        }
+    );
+
+    assert_eq!(
+        TestEnum2::SCHEMA,
+        &NamedType {
+            name: "TestEnum2",
+            ty: &SdmTy::Enum(&[
+                &NamedVariant { name: "Nt", ty: &SdmTy::NewtypeVariant(u64::SCHEMA) },
+                &NamedVariant { name: "Tup", ty: &SdmTy::TupleVariant(&[u64::SCHEMA, bool::SCHEMA]) },
+            ]),
+        }
+    );
+}