From 2f3f9b3a33edbf3f943f20334d46d3238d44f186 Mon Sep 17 00:00:00 2001
From: Thomas Eizinger <thomas@eizinger.io>
Date: Fri, 22 Oct 2021 17:21:34 +1100
Subject: [PATCH] Support generic params

---
 xtra_productivity/src/lib.rs                  | 32 +++++++++++++++++--
 .../tests/pass/actor_with_generics.rs         | 28 ++++++++++++++++
 2 files changed, 58 insertions(+), 2 deletions(-)
 create mode 100644 xtra_productivity/tests/pass/actor_with_generics.rs

diff --git a/xtra_productivity/src/lib.rs b/xtra_productivity/src/lib.rs
index 1a608cb..19a5411 100644
--- a/xtra_productivity/src/lib.rs
+++ b/xtra_productivity/src/lib.rs
@@ -1,6 +1,6 @@
 use proc_macro::TokenStream;
 use quote::quote;
-use syn::{FnArg, ImplItem, ItemImpl, ReturnType};
+use syn::{FnArg, GenericParam, ImplItem, ItemImpl, ReturnType};
 
 #[proc_macro_attribute]
 pub fn xtra_productivity(_attribute: TokenStream, item: TokenStream) -> TokenStream {
@@ -8,6 +8,30 @@ pub fn xtra_productivity(_attribute: TokenStream, item: TokenStream) -> TokenStr
 
     let actor = block.self_ty;
 
+    let generic_params = &block.generics.params;
+
+    let generic_types = block
+        .generics
+        .params
+        .iter()
+        .filter_map(|param| match param {
+            GenericParam::Type(ty) => Some(ty.ident.clone()),
+            _ => None,
+        })
+        .collect::<Vec<_>>();
+
+    let additional_bounds = block
+        .generics
+        .where_clause
+        .map(|bounds| {
+            let predicates = bounds.predicates;
+
+            quote! {
+                #predicates
+            }
+        })
+        .unwrap_or_default();
+
     let code = block
         .items
         .into_iter()
@@ -42,7 +66,11 @@ pub fn xtra_productivity(_attribute: TokenStream, item: TokenStream) -> TokenStr
                 }
 
                 #[async_trait]
-                impl xtra::Handler<#message_type> for #actor {
+                impl<#generic_params> xtra::Handler<#message_type> for #actor
+                    where
+                        #additional_bounds
+                        #(#generic_types: Send + 'static),*
+                {
                     async fn handle(&mut self, #message_arg, #context_arg) #method_return #method_block
                 }
             }
diff --git a/xtra_productivity/tests/pass/actor_with_generics.rs b/xtra_productivity/tests/pass/actor_with_generics.rs
new file mode 100644
index 0000000..143cb6e
--- /dev/null
+++ b/xtra_productivity/tests/pass/actor_with_generics.rs
@@ -0,0 +1,28 @@
+use async_trait::async_trait;
+use std::marker::PhantomData;
+use xtra_productivity::xtra_productivity;
+
+struct ActorWithParam<C> {
+    ty: PhantomData<C>,
+}
+
+struct DummyMessage;
+
+trait Foo {}
+
+impl<C: 'static + Send> xtra::Actor for ActorWithParam<C> {}
+
+// Dummy actor, xtra::Handler and xtra::Message impls generated by xtra_productivity
+#[xtra_productivity]
+impl<C> ActorWithParam<C>
+where
+    C: Foo,
+{
+    pub fn handle_dummy_message(&mut self, _message: DummyMessage) {
+        assert_impls_foo::<C>();
+    }
+}
+
+fn assert_impls_foo<T: Foo>() {}
+
+fn main() {}