From 67f59a4a9721d286c4b84604590242a5520b9abc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jakub=20Ber=C3=A1nek?= <berykubik@gmail.com>
Date: Thu, 14 Dec 2023 17:36:50 +0100
Subject: [PATCH] Allow passing arbitrary attributes to segments

---
 CHANGELOG.md                | 11 +++++++++++
 src/attributes.rs           | 32 ++++++++++++++++++++++++--------
 tests/segment_attributes.rs | 25 +++++++++++++++++++++++++
 3 files changed, 60 insertions(+), 8 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 570a1e8..8ec9351 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,3 +1,14 @@
+# Dev
+- Allow passing arbitrary attributes to delegation segments:
+```rust
+impl Foo {
+  delegate! {
+    #[inline(always)]
+    to self.0 { ... }
+  }
+}
+```
+
 # 0.11.0 (4. 12. 2023)
 - Allow delegating an associated function (not just a method).
 ```rust
diff --git a/src/attributes.rs b/src/attributes.rs
index 5392887..69bd4d1 100644
--- a/src/attributes.rs
+++ b/src/attributes.rs
@@ -225,6 +225,7 @@ pub struct SegmentAttributes {
     pub expressions: Vec<ReturnExpression>,
     pub generate_await: Option<bool>,
     pub target_trait: Option<TypePath>,
+    pub other_attrs: Vec<Attribute>,
 }
 
 pub fn parse_segment_attributes(attrs: &[Attribute]) -> SegmentAttributes {
@@ -232,10 +233,7 @@ pub fn parse_segment_attributes(attrs: &[Attribute]) -> SegmentAttributes {
     let mut generate_await: Option<bool> = None;
     let mut target_trait: Option<TraitTarget> = None;
 
-    let (parsed, mut other) = parse_attributes(attrs);
-    if other.next().is_some() {
-        panic!("Only return expression attributes can be used on a segment (into, try_into, unwrap or await).");
-    }
+    let (parsed, other) = parse_attributes(attrs);
 
     for attribute in parsed {
         match attribute {
@@ -261,23 +259,31 @@ pub fn parse_segment_attributes(attrs: &[Attribute]) -> SegmentAttributes {
         expressions,
         generate_await,
         target_trait: target_trait.map(|t| t.type_path),
+        other_attrs: other.cloned().collect::<Vec<_>>(),
     }
 }
 
 /// Applies default values from the segment and adds them to the method attributes.
 pub fn combine_attributes<'a>(
     mut method_attrs: MethodAttributes<'a>,
-    segment_attrs: &SegmentAttributes,
+    segment_attrs: &'a SegmentAttributes,
 ) -> MethodAttributes<'a> {
+    let SegmentAttributes {
+        expressions,
+        generate_await,
+        target_trait,
+        other_attrs,
+    } = segment_attrs;
+
     if method_attrs.generate_await.is_none() {
-        method_attrs.generate_await = segment_attrs.generate_await;
+        method_attrs.generate_await = *generate_await;
     }
 
     if method_attrs.target_trait.is_none() {
-        method_attrs.target_trait = segment_attrs.target_trait.clone();
+        method_attrs.target_trait = target_trait.clone();
     }
 
-    for expr in &segment_attrs.expressions {
+    for expr in expressions {
         match expr {
             ReturnExpression::Into(path) => {
                 if !method_attrs
@@ -294,5 +300,15 @@ pub fn combine_attributes<'a>(
         }
     }
 
+    for other_attr in other_attrs {
+        if !method_attrs
+            .attributes
+            .iter()
+            .any(|attr| attr.path().get_ident() != other_attr.path().get_ident())
+        {
+            method_attrs.attributes.push(other_attr);
+        }
+    }
+
     method_attrs
 }
diff --git a/tests/segment_attributes.rs b/tests/segment_attributes.rs
index 319c5dc..95b3772 100644
--- a/tests/segment_attributes.rs
+++ b/tests/segment_attributes.rs
@@ -186,3 +186,28 @@ fn test_segment_through_trait() {
     assert_eq!(bar.f(), 0);
     assert_eq!(bar.f2(), 1);
 }
+
+#[test]
+fn test_segment_inline() {
+    struct Foo;
+
+    impl Foo {
+        fn f(&self) -> u32 {
+            0
+        }
+    }
+
+    struct Bar(Foo);
+
+    impl Bar {
+        delegate! {
+            #[inline(always)]
+            to self.0 {
+                fn f(&self) -> u32;
+            }
+        }
+    }
+
+    let bar = Bar(Foo);
+    assert_eq!(bar.f(), 0);
+}