Add inspect

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

crates/tytix-core/src/lib.rs

@@ -41,6 +41,10 @@ impl InternalMessage {
             })
     }
 
+    pub fn as_ref<M: Any>(&self) -> Option<&M> {
+        self.value.downcast_ref()
+    }
+
     pub fn type_name(&self) -> &'static str {
         self.name
     }

crates/tytix/src/lib.rs

@@ -12,12 +12,17 @@ use tytix_core::anyhow;
 pub trait AddressExt<MB> {
     fn map_message<F, M, R, U>(self, f: F) -> MappedBeforeAddress<Self, F, M, R>
     where
-        MB: MessageBundle,
         F: Fn(M) -> U + 'static,
         U: Future<Output = R>,
         M: Message,
         R: Message + IsContainedInBundle<MB>,
         Self: Sized;
+
+    fn inspect<F, M, U>(self, f: F) -> Inspect<Self, F, U, M>
+    where
+        F: Fn(&M) -> U,
+        M: Message + IsContainedInBundle<MB>,
+        Self: Sized;
 }
 
 impl<MB: MessageBundle, A: Address<MB>> AddressExt<MB> for A {
@@ -42,6 +47,48 @@ impl<MB: MessageBundle, A: Address<MB>> AddressExt<MB> for A {
             _pd: PhantomData,
         }
     }
+
+    fn inspect<F, M, U>(self, f: F) -> Inspect<Self, F, U, M>
+    where
+        F: Fn(&M) -> U,
+        Self: Sized,
+        M: Message + IsContainedInBundle<MB>,
+    {
+        const {
+            let true = <M as IsContainedInBundle<MB>>::IS_CONTAINED else {
+                panic!("Message is not contained in MessageBundle",);
+            };
+        }
+
+        Inspect {
+            address: self,
+            func: f,
+            _pd: PhantomData,
+        }
+    }
+}
+
+pub struct Inspect<A, F, U, M> {
+    address: A,
+    func: F,
+    _pd: PhantomData<fn(&M) -> U>,
+}
+
+impl<A, F, U, M> InternalMessageHandler for Inspect<A, F, U, M>
+where
+    A: InternalMessageHandler,
+    M: Message,
+    F: Fn(&M) -> U,
+    U: Future<Output = ()>,
+{
+    type HandledMessages = A::HandledMessages;
+    async fn handle_message(&mut self, msg: InternalMessage) -> anyhow::Result<InternalMessage> {
+        if let Some(msg) = msg.as_ref() {
+            (self.func)(msg).await;
+        }
+
+        self.address.handle_message(msg).await
+    }
 }
 
 pub struct MappedBeforeAddress<A, F, M, R> {
@@ -117,7 +164,7 @@ mod tests {
     struct SimpleAddress;
 
     impl InternalMessageHandler for SimpleAddress {
-        type HandledMessages = (Foo,);
+        type HandledMessages = (Foo, Bar);
 
         fn handle_message(
             &mut self,
@@ -142,4 +189,22 @@ mod tests {
 
         MSG.get().expect("The message was mapped!");
     }
+
+    #[apply(test!)]
+    async fn check_inspect() {
+        static MSG: OnceLock<bool> = OnceLock::new();
+
+        let mut sa = SimpleAddress.inspect(|_b: &Bar| {
+            let _ = MSG.set(true);
+            async {}
+        });
+
+        sa.send(Foo).await.unwrap();
+
+        assert!(MSG.get().is_none());
+
+        sa.send(Bar).await.unwrap();
+
+        MSG.get().expect("The message was inspected!");
+    }
 }