add flatten (#34)

* add flatten

* add tests,docs,fix the feature

* remove clone from getter opt

* fix bug in generated struct field placement

* flatten subset

* add comment

* Fix bug in field_index usage

* Use *last* path segment not first

---------

Co-authored-by: Michael Sproul <michael@sigmaprime.io>

book/src/config/field.md

@@ -68,3 +68,61 @@ 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 with variants that 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`:
+```rust
+#[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`:
+```rust
+#[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,
+}
+```
+
+If you wish to only flatten into only a subset of variants, you can define them like so: 
+
+```rust
+#[superstruct(variants(A, B))]
+struct InnerMessage {
+    pub x: u64,
+    pub y: u64,
+}
+
+#[superstruct(variants(A, B, C))]
+struct Message {
+    #[superstruct(flatten(A,B))]
+    pub inner: InnerMessage,
+}
+```

src/lib.rs

@@ -1,4 +1,5 @@
 use attributes::{IdentList, NestedMetaList};
+use darling::util::Override;
 use darling::FromMeta;
 use from::{
     generate_from_enum_trait_impl_for_ref, generate_from_variant_trait_impl,
@@ -65,6 +66,11 @@ struct StructOpts {
 /// Field-level configuration.
 #[derive(Debug, Default, FromMeta)]
 struct FieldOpts {
+    // TODO: When we update darling, we can replace `Override`
+    // with a custom enum and use `#[darling(word)]` on the variant
+    // we want to use for `#[superstruct(flatten)]` (no variants specified).
+    #[darling(default)]
+    flatten: Option<Override<HashMap<Ident, ()>>>,
     #[darling(default)]
     only: Option<HashMap<Ident, ()>>,
     #[darling(default)]
@@ -73,6 +79,13 @@ struct FieldOpts {
     partial_getter: Option<GetterOpts>,
 }
 
+fn should_skip(flatten: &Override<HashMap<Ident, ()>>, key: &Ident) -> bool {
+    match flatten {
+        Override::Inherit => false,
+        Override::Explicit(map) => !map.is_empty() && !map.contains_key(key),
+    }
+}
+
 /// Getter configuration for a specific field
 #[derive(Debug, Default, FromMeta)]
 struct GetterOpts {
@@ -159,7 +172,7 @@ pub fn superstruct(args: TokenStream, input: TokenStream) -> TokenStream {
     let mut variant_fields =
         HashMap::<_, _>::from_iter(variant_names.iter().zip(iter::repeat(vec![])));
 
-    for field in item.fields.iter() {
+    for field in &item.fields {
         let name = field.ident.clone().expect("named fields only");
         let field_opts = field
             .attrs
@@ -193,20 +206,86 @@ pub fn superstruct(args: TokenStream, input: TokenStream) -> TokenStream {
             panic!("can't configure `only` and `getter` on the same field");
         } else if field_opts.only.is_none() && field_opts.partial_getter.is_some() {
             panic!("can't set `partial_getter` options on common field");
+        } else if field_opts.flatten.is_some() && field_opts.only.is_some() {
+            panic!("can't set `flatten` and `only` on the same field");
+        } else if field_opts.flatten.is_some() && field_opts.getter.is_some() {
+            panic!("can't set `flatten` and `getter` on the same field");
+        } else if field_opts.flatten.is_some() && field_opts.partial_getter.is_some() {
+            panic!("can't set `flatten` and `partial_getter` on the same field");
         }
 
         let only = field_opts.only.map(|only| only.keys().cloned().collect());
         let getter_opts = field_opts.getter.unwrap_or_default();
         let partial_getter_opts = field_opts.partial_getter.unwrap_or_default();
 
-        // Add to list of all fields
-        fields.push(FieldData {
-            name,
-            field: output_field,
-            only,
-            getter_opts,
-            partial_getter_opts,
-        });
+        if let Some(flatten_opts) = field_opts.flatten {
+            for variant in variant_names {
+                let variant_field_index = variant_fields
+                    .get(variant)
+                    .expect("invalid variant name")
+                    .iter()
+                    .position(|f| f.ident.as_ref() == Some(&name))
+                    .expect("flattened fields are present on all variants");
+
+                if should_skip(&flatten_opts, variant) {
+                    // Remove the field from the field map
+                    let fields = variant_fields
+                        .get_mut(variant)
+                        .expect("invalid variant name");
+                    fields.remove(variant_field_index);
+                    continue;
+                }
+
+                // 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 last_segment = &mut p
+                            .path
+                            .segments
+                            .last_mut()
+                            .expect("path should have at least one segment");
+                        let inner_ty_name = last_segment.ident.clone();
+                        let next_variant_ty_name = format_ident!("{}{}", inner_ty_name, variant);
+                        last_segment.ident = next_variant_ty_name;
+                    }
+                    _ => panic!("field must be a path"),
+                };
+
+                // 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()
+                };
+
+                fields.push(FieldData {
+                    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,
+                });
+
+                // Update the variant field map
+                let fields = variant_fields
+                    .get_mut(variant)
+                    .expect("invalid variant name");
+                *fields
+                    .get_mut(variant_field_index)
+                    .expect("invalid field index") = next_variant_field;
+            }
+        } else {
+            fields.push(FieldData {
+                name,
+                field: output_field,
+                only,
+                getter_opts,
+                partial_getter_opts,
+            });
+        }
     }
 
     // Generate structs for all of the variants.

tests/flatten.rs

@@ -0,0 +1,170 @@
+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());
+}
+
+#[test]
+fn flatten_subset() {
+    #[superstruct(variants(A, B), variant_attributes(derive(Debug, PartialEq, Eq)))]
+    #[derive(Debug, PartialEq, Eq)]
+    struct InnerMessageSubset {
+        pub x: u64,
+        #[superstruct(only(B))]
+        pub y: u64,
+    }
+
+    #[superstruct(variants(A, B, C), variant_attributes(derive(Debug, PartialEq, Eq)))]
+    #[derive(Debug, PartialEq, Eq)]
+    struct MessageSubset {
+        #[superstruct(flatten(A, B))]
+        pub inner: InnerMessageSubset,
+    }
+
+    let message_a = MessageSubset::A(MessageSubsetA {
+        inner: InnerMessageSubsetA { x: 1 },
+    });
+    let message_b = MessageSubset::B(MessageSubsetB {
+        inner: InnerMessageSubsetB { x: 3, y: 4 },
+    });
+    let message_c = MessageSubset::C(MessageSubsetC {});
+    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());
+    assert!(message_c.inner_a().is_err());
+    assert!(message_c.inner_b().is_err());
+}
+
+#[test]
+fn flatten_not_first_field() {
+    use test_mod::*;
+
+    // Put this type in a submodule to test path parsing in `flatten`.
+    pub mod test_mod {
+        use superstruct::superstruct;
+
+        #[superstruct(variants(A, B), variant_attributes(derive(Debug, PartialEq, Eq)))]
+        #[derive(Debug, PartialEq, Eq)]
+        pub struct InnerMessageTwo {
+            pub x: u64,
+            #[superstruct(only(B))]
+            pub y: u64,
+        }
+    }
+
+    #[superstruct(variants(A, B), variant_attributes(derive(Debug, PartialEq, Eq)))]
+    #[derive(Debug, PartialEq, Eq)]
+    struct MessageTwo {
+        #[superstruct(only(A), partial_getter(copy))]
+        pub other: u64,
+        #[superstruct(flatten)]
+        pub inner: test_mod::InnerMessageTwo,
+    }
+
+    let message_a = MessageTwo::A(MessageTwoA {
+        other: 21,
+        inner: InnerMessageTwoA { x: 1 },
+    });
+    let message_b = MessageTwo::B(MessageTwoB {
+        inner: InnerMessageTwoB { x: 3, y: 4 },
+    });
+    assert_eq!(message_a.other().unwrap(), 21);
+    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 = MessageTwoRef::A(&MessageTwoA {
+        other: 21,
+        inner: InnerMessageTwoA { x: 1 },
+    });
+    let message_b_ref = MessageTwoRef::B(&MessageTwoB {
+        inner: InnerMessageTwoB { x: 3, y: 4 },
+    });
+    assert_eq!(message_a.other().unwrap(), 21);
+    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 = MessageTwoA {
+        other: 21,
+        inner: InnerMessageTwoA { x: 1 },
+    };
+    let mut inner_b = MessageTwoB {
+        inner: InnerMessageTwoB { x: 3, y: 4 },
+    };
+
+    // Re-initialize the struct to avoid borrow checker errors.
+    let mut message_a_ref_mut = MessageTwoRefMut::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 = MessageTwoRefMut::A(&mut inner_a);
+    assert!(message_a_ref_mut.inner_b_mut().is_err());
+    let mut message_b_ref_mut = MessageTwoRefMut::B(&mut inner_b);
+    assert_eq!(message_b_ref_mut.inner_b_mut().unwrap().x, 3);
+    let mut message_b_ref_mut = MessageTwoRefMut::B(&mut inner_b);
+    assert_eq!(message_b_ref_mut.inner_b_mut().unwrap().y, 4);
+    let mut message_b_ref_mut = MessageTwoRefMut::B(&mut inner_b);
+    assert!(message_b_ref_mut.inner_a_mut().is_err());
+}