add tests,docs,fix the feature

book/src/config/field.md

@@ -68,3 +68,45 @@ may be applied in a single attribute, e.g. `#[superstruct(partial_getter(copy, n
 The error type for partial getters can currently only be configured on a per-struct basis
 via the [`partial_getter_error`](./struct.md#partial-getter-error) attribute, although this may
 change in a future release.
+
+## Flatten
+
+```
+#[superstruct(flatten)]
+```
+
+This attribute can only be applied to enum fields that whose variants match each variant of the
+superstruct. This is useful for nesting superstructs whose variant types should be linked.
+
+This will automatically create a partial getter for each variant. The following two examples are equivalent.
+
+Using `flatten`:
+```
+#[superstruct(variants(A, B))]
+struct InnerMessage {
+    pub x: u64,
+    pub y: u64,
+}
+
+#[superstruct(variants(A, B))]
+struct Message {
+    #[superstruct(flatten)]
+    pub inner: InnerMessage,
+}
+```
+Equivalent without `flatten`:
+```
+#[superstruct(variants(A, B))]
+struct InnerMessage {
+    pub x: u64,
+    pub y: u64,
+}
+
+#[superstruct(variants(A, B))]
+struct Message {
+    #[superstruct(only(A), partial_getter(rename = "inner_a"))]
+    pub inner: InnerMessageA,
+    #[superstruct(only(B), partial_getter(rename = "inner_b"))]
+    pub inner: InnerMessageB,
+}
+```

src/lib.rs

@@ -12,8 +12,8 @@ use quote::{format_ident, quote, ToTokens};
 use std::collections::HashMap;
 use std::iter::{self, FromIterator};
 use syn::{
-    parse2, parse_macro_input, Attribute, AttributeArgs, Expr, Field, GenericParam, Ident,
-    ItemEnum, ItemStruct, Lifetime, LifetimeDef, Type, TypeGenerics, TypeParamBound,
+    parse_macro_input, Attribute, AttributeArgs, Expr, Field, GenericParam, Ident, ItemStruct,
+    Lifetime, LifetimeDef, Type, TypeGenerics, TypeParamBound,
 };
 
 mod attributes;
@@ -76,7 +76,7 @@ struct FieldOpts {
 }
 
 /// Getter configuration for a specific field
-#[derive(Debug, Default, FromMeta)]
+#[derive(Clone, Debug, Default, FromMeta)]
 struct GetterOpts {
     #[darling(default)]
     copy: bool,
@@ -208,39 +208,45 @@ pub fn superstruct(args: TokenStream, input: TokenStream) -> TokenStream {
         let partial_getter_opts = field_opts.partial_getter.unwrap_or_default();
 
         if field_opts.flatten {
-            // Parse the inner type, making sure it's an enum.
-            let ty = &output_field.ty;
-
-            let inner_enum: ItemEnum = parse2(output_field.ty.to_token_stream())
-                .expect(format!("inner type must be an enum {ty:?}").as_str());
-
-            // Extract the names of the variants for the inner enum.
-            let variant_names_inner: Vec<_> = inner_enum
-                .variants
-                .iter()
-                .map(|v| v.ident.clone())
-                .collect();
+            for variant in variant_names {
+                // Update the struct name for this variant.
+                let mut next_variant_field = output_field.clone();
+                match &mut next_variant_field.ty {
+                    Type::Path(ref mut p) => {
+                        let first_segment = &mut p
+                            .path
+                            .segments
+                            .first_mut()
+                            .expect("path should have at least one segment");
+                        let inner_ty_name = first_segment.ident.clone();
+                        let next_variant_ty_name = format_ident!("{}{}", inner_ty_name, variant);
+                        first_segment.ident = next_variant_ty_name;
+                    }
+                    _ => panic!("field must be a path"),
+                };
 
-            // Compare the sets of variant names.
-            assert_eq!(
-                variant_names, &variant_names_inner,
-                "variant names must match"
-            );
+                // Create a partial getter for the field.
+                let partial_getter_rename =
+                    format_ident!("{}_{}", name, variant.to_string().to_lowercase());
+                let partial_getter_opts = GetterOpts {
+                    rename: Some(partial_getter_rename),
+                    ..<_>::default()
+                };
 
-            for variant in inner_enum.variants {
-                assert!(
-                    variant.fields.len() == 1,
-                    "only one field allowed in flattened enum variants"
-                );
-                let inner_field = variant.fields.into_iter().next().unwrap();
-                // Add the variant name as a suffix to the getter name.
                 fields.push(FieldData {
-                    name: format_ident!("{}_{}", name, variant.ident.to_string().to_lowercase()),
-                    field: inner_field,
-                    only: None,
+                    name: name.clone(),
+                    field: next_variant_field.clone(),
+                    // Make sure the field is only accessible from this variant.
+                    only: Some(vec![variant.clone()]),
                     getter_opts: <_>::default(),
-                    partial_getter_opts: <_>::default(),
+                    partial_getter_opts,
                 });
+
+                // Update the variant field map
+                let fields = variant_fields
+                    .get_mut(variant)
+                    .expect("invalid variant name");
+                *fields.get_mut(0).expect("no fields for variant") = next_variant_field;
             }
         } else {
             fields.push(FieldData {

tests/flatten.rs

@@ -0,0 +1,62 @@
+use superstruct::superstruct;
+
+#[test]
+fn flatten() {
+    #[superstruct(variants(A, B), variant_attributes(derive(Debug, PartialEq, Eq)))]
+    #[derive(Debug, PartialEq, Eq)]
+    struct InnerMessage {
+        pub x: u64,
+        #[superstruct(only(B))]
+        pub y: u64,
+    }
+
+    #[superstruct(variants(A, B), variant_attributes(derive(Debug, PartialEq, Eq)))]
+    #[derive(Debug, PartialEq, Eq)]
+    struct Message {
+        #[superstruct(flatten)]
+        pub inner: InnerMessage,
+    }
+
+    let message_a = Message::A(MessageA {
+        inner: InnerMessageA { x: 1 },
+    });
+    let message_b = Message::B(MessageB {
+        inner: InnerMessageB { x: 3, y: 4 },
+    });
+    assert_eq!(message_a.inner_a().unwrap().x, 1);
+    assert!(message_a.inner_b().is_err());
+    assert_eq!(message_b.inner_b().unwrap().x, 3);
+    assert_eq!(message_b.inner_b().unwrap().y, 4);
+    assert!(message_b.inner_a().is_err());
+
+    let message_a_ref = MessageRef::A(&MessageA {
+        inner: InnerMessageA { x: 1 },
+    });
+    let message_b_ref = MessageRef::B(&MessageB {
+        inner: InnerMessageB { x: 3, y: 4 },
+    });
+    assert_eq!(message_a_ref.inner_a().unwrap().x, 1);
+    assert!(message_a_ref.inner_b().is_err());
+    assert_eq!(message_b_ref.inner_b().unwrap().x, 3);
+    assert_eq!(message_b_ref.inner_b().unwrap().y, 4);
+    assert!(message_b_ref.inner_a().is_err());
+
+    let mut inner_a = MessageA {
+        inner: InnerMessageA { x: 1 },
+    };
+    let mut inner_b = MessageB {
+        inner: InnerMessageB { x: 3, y: 4 },
+    };
+
+    // Re-initialize the struct to avoid borrow checker errors.
+    let mut message_a_ref_mut = MessageRefMut::A(&mut inner_a);
+    assert_eq!(message_a_ref_mut.inner_a_mut().map(|inner| inner.x), Ok(1));
+    let mut message_a_ref_mut = MessageRefMut::A(&mut inner_a);
+    assert!(message_a_ref_mut.inner_b_mut().is_err());
+    let mut message_b_ref_mut = MessageRefMut::B(&mut inner_b);
+    assert_eq!(message_b_ref_mut.inner_b_mut().unwrap().x, 3);
+    let mut message_b_ref_mut = MessageRefMut::B(&mut inner_b);
+    assert_eq!(message_b_ref_mut.inner_b_mut().unwrap().y, 4);
+    let mut message_b_ref_mut = MessageRefMut::B(&mut inner_b);
+    assert!(message_b_ref_mut.inner_a_mut().is_err());
+}