Also map the reply

Signed-off-by: Marcel Müller <neikos@neikos.email>

crates/tytix-core/src/lib.rs

@@ -62,14 +62,13 @@ pub trait InternalMessageHandler {
     ) -> impl Future<Output = anyhow::Result<InternalMessage>>;
 }
 
-impl<MB, IMH> Address<MB> for IMH
+impl<IMH> Address<IMH::HandledMessages> for IMH
 where
-    MB: MessageBundle,
+    IMH: InternalMessageHandler,
-    IMH: InternalMessageHandler<HandledMessages = MB>,
 {
     fn send<M: Message>(&mut self, message: M) -> impl Future<Output = anyhow::Result<M::Reply>> {
         const {
-            let true = <M as IsContainedInBundle<MB>>::IS_CONTAINED else {
+            let true = <M as IsContainedInBundle<IMH::HandledMessages>>::IS_CONTAINED else {
                 panic!("Message is not contained in MessageBundle",);
             };
         }

crates/tytix/src/lib.rs

@@ -1,7 +1,6 @@
 use std::marker::PhantomData;
 
 use futures::FutureExt;
-use tytix_core::Address;
 use tytix_core::InternalMessage;
 use tytix_core::InternalMessageHandler;
 use tytix_core::IsContainedInBundle;
@@ -9,13 +8,16 @@ use tytix_core::Message;
 use tytix_core::MessageBundle;
 use tytix_core::anyhow;
 
+pub type ReplyOf<M> = <M as Message>::Reply;
+
 pub trait AddressExt<MB> {
-    fn map_message<F, M, R, U>(self, f: F) -> MappedMessage<Self, F, M, R>
+    fn map<M, R, F, U, RF, RU>(self, f: F, r: RF) -> MappedMessage<Self, M, F, U, RF, RU>
     where
-        F: Fn(M) -> U + 'static,
-        U: Future<Output = R>,
         M: Message,
         R: Message + IsContainedInBundle<MB>,
+        F: FnMut(M) -> U + 'static,
+        U: Future<Output = R>,
+        RF: FnMut(<R as Message>::Reply) -> RU + 'static,
         Self: Sized;
 
     fn inspect<F, M, U>(self, f: F) -> Inspect<Self, F, U, M>
@@ -25,14 +27,13 @@ pub trait AddressExt<MB> {
         Self: Sized;
 }
 
-impl<MB: MessageBundle, A: Address<MB>> AddressExt<MB> for A {
+impl<MB: MessageBundle, A: InternalMessageHandler<HandledMessages = MB>> AddressExt<MB> for A {
-    fn map_message<F, M, R, U>(self, f: F) -> MappedMessage<Self, F, M, R>
+    fn map<M, R, F, U, RF, RU>(self, f: F, r: RF) -> MappedMessage<Self, M, F, U, RF, RU>
     where
-        MB: MessageBundle,
-        F: Fn(M) -> U + 'static,
-        U: Future<Output = R>,
         M: Message,
         R: Message + IsContainedInBundle<MB>,
+        F: FnMut(M) -> U + 'static,
+        RF: FnMut(<R as Message>::Reply) -> RU + 'static,
         Self: Sized,
     {
         const {
@@ -44,6 +45,7 @@ impl<MB: MessageBundle, A: Address<MB>> AddressExt<MB> for A {
         MappedMessage {
             address: self,
             func: f,
+            reply: r,
             _pd: PhantomData,
         }
     }
@@ -91,10 +93,11 @@ where
     }
 }
 
-pub struct MappedMessage<A, F, M, R> {
+pub struct MappedMessage<A, M, F, U, RF, RU> {
     address: A,
     func: F,
-    _pd: PhantomData<fn(M) -> R>,
+    reply: RF,
+    _pd: PhantomData<fn(M) -> (U, RU)>,
 }
 
 pub struct MappedHandledMessages<MB, M, R> {
@@ -110,13 +113,15 @@ where
     const IDS: tytix_core::BundleChain = MB::IDS.without::<R>().with::<M>();
 }
 
-impl<A, F, M, R, U> InternalMessageHandler for MappedMessage<A, F, M, R>
+impl<A, M, R, F, U, RF, RU> InternalMessageHandler for MappedMessage<A, M, F, U, RF, RU>
 where
     A: InternalMessageHandler,
     M: Message,
     R: Message + IsContainedInBundle<A::HandledMessages>,
     F: FnMut(M) -> U,
     U: Future<Output = R>,
+    RF: FnMut(<R as Message>::Reply) -> RU,
+    RU: Future<Output = <M as Message>::Reply>,
 {
     type HandledMessages = MappedHandledMessages<A::HandledMessages, M, R>;
 
@@ -135,7 +140,15 @@ where
         async {
             let msg = msg.await;
 
-            self.address.handle_message(msg).await
+            let msg_reply = self.address.handle_message(msg).await?;
+
+            let Ok(msg_reply) = msg_reply.into_inner() else {
+                return Err(anyhow::anyhow!("Internal Error: Received the wrong type"));
+            };
+
+            let msg_reply = (self.reply)(msg_reply).await;
+
+            Ok(InternalMessage::new(msg_reply))
         }
     }
 }
@@ -146,6 +159,7 @@ mod tests {
 
     use macro_rules_attribute::apply;
     use smol_macros::test;
+    use tytix_core::Address;
 
     use super::*;
 
@@ -180,10 +194,13 @@ mod tests {
     async fn check_mapping() {
         static MSG: OnceLock<bool> = OnceLock::new();
 
-        let mut sa = SimpleAddress.map_message(|_b: Bar| {
+        let mut sa = SimpleAddress.map(
-            let _ = MSG.set(true);
+            |_b: Bar| {
-            async { Foo }
+                let _ = MSG.set(true);
-        });
+                async { Foo }
+            },
+            |_a| async {},
+        );
 
         sa.send(Bar).await.unwrap();