From cd8d55986735b0db8150133d50b1cf51f1cf1a1b Mon Sep 17 00:00:00 2001
From: Daniel Karzel <daniel@comit.network>
Date: Fri, 22 Oct 2021 14:34:05 +1100
Subject: [PATCH] Macro attribute so we can control impl of `xtra::Message`

Co-authored-by: Mariusz Klochowicz <mariusz@klochowicz.com>
---
 xtra_productivity/src/lib.rs                  | 34 ++++++++++++++++---
 .../tests/pass/can_handle_message.rs          | 20 +++++++++++
 2 files changed, 49 insertions(+), 5 deletions(-)

diff --git a/xtra_productivity/src/lib.rs b/xtra_productivity/src/lib.rs
index ca1890e..99e1d26 100644
--- a/xtra_productivity/src/lib.rs
+++ b/xtra_productivity/src/lib.rs
@@ -1,10 +1,26 @@
 use proc_macro::TokenStream;
 use quote::quote;
-use syn::{FnArg, GenericParam, ImplItem, ItemImpl, ReturnType};
+use syn::{FnArg, GenericParam, ImplItem, ItemImpl, MetaNameValue, ReturnType};
 
 #[proc_macro_attribute]
-pub fn xtra_productivity(_attribute: TokenStream, item: TokenStream) -> TokenStream {
+pub fn xtra_productivity(attribute: TokenStream, item: TokenStream) -> TokenStream {
     let block = syn::parse::<ItemImpl>(item).unwrap();
+    let want_message_impl = if attribute.is_empty() {
+        true
+    } else {
+        let attribute = syn::parse::<MetaNameValue>(attribute).unwrap();
+        if !attribute.path.is_ident("message_impl") {
+            panic!(
+                "Unexpected attribute {:?}",
+                attribute.path.get_ident().unwrap()
+            )
+        }
+
+        matches!(
+            attribute.lit,
+            syn::Lit::Bool(syn::LitBool { value: true, .. })
+        )
+    };
 
     let actor = block.self_ty;
 
@@ -60,10 +76,18 @@ pub fn xtra_productivity(_attribute: TokenStream, item: TokenStream) -> TokenStr
                 ReturnType::Type(_, ref t) => quote! { #t }
             };
 
-            quote! {
-                impl xtra::Message for #message_type {
-                    type Result = #result_type;
+            let message_impl = if want_message_impl {
+                quote! {
+                    impl xtra::Message for #message_type {
+                        type Result = #result_type;
+                    }
                 }
+            } else {
+                quote! {}
+            };
+
+            quote! {
+                #message_impl
 
                 #[async_trait::async_trait]
                 impl<#generic_params> xtra::Handler<#message_type> for #actor
diff --git a/xtra_productivity/tests/pass/can_handle_message.rs b/xtra_productivity/tests/pass/can_handle_message.rs
index 14f54ee..91e751d 100644
--- a/xtra_productivity/tests/pass/can_handle_message.rs
+++ b/xtra_productivity/tests/pass/can_handle_message.rs
@@ -30,6 +30,21 @@ impl DummyActor {
 
 fn is_i32(_: i32) {}
 
+struct DummyMessageWithoutMessageImpl;
+
+#[xtra_productivity(message_impl = false)]
+impl DummyActor {
+    pub fn handle_dummy_message_without_message_impl(
+        &mut self,
+        _message: DummyMessageWithoutMessageImpl,
+    ) {
+    }
+}
+
+impl xtra::Message for DummyMessageWithoutMessageImpl {
+    type Result = ();
+}
+
 #[tokio::main]
 async fn main() {
     // Create dummy actor
@@ -39,4 +54,9 @@ async fn main() {
     let i32 = dummy_actor.send(DummyMessage).await.unwrap();
     is_i32(i32);
     dummy_actor.send(DummyMessageWithContext).await.unwrap();
+
+    dummy_actor
+        .send(DummyMessageWithoutMessageImpl)
+        .await
+        .unwrap();
 }