core/muxing: Introduce StreamMuxerEvent::map_inbound_stream (#2691)

Co-authored-by: Max Inden <[email protected]>
libp2p · Jun 15, 2022 · 3c120ef · 3c120ef
1 parent 04f31cd
commit 3c120ef
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 17 deletions.
diff --git a/core/CHANGELOG.md b/core/CHANGELOG.md
@@ -1,3 +1,9 @@
+# 0.33.1 - unreleased
+
+- Introduce `StreamMuxerEvent::map_inbound_stream`. See [PR 2691].
+
+[PR 2691]: https://github.com/libp2p/rust-libp2p/pull/2691
+
 # 0.33.0
 
 - Have methods on `Transport` take `&mut self` instead of `self`. See [PR 2529].

diff --git a/core/Cargo.toml b/core/Cargo.toml
@@ -3,7 +3,7 @@ name = "libp2p-core"
 edition = "2021"
 rust-version = "1.56.1"
 description = "Core traits and structs of libp2p"
-version = "0.33.0"
+version = "0.33.1"
 authors = ["Parity Technologies <[email protected]>"]
 license = "MIT"
 repository = "https://github.com/libp2p/rust-libp2p"

diff --git a/core/src/either.rs b/core/src/either.rs
@@ -210,22 +210,14 @@ where
         cx: &mut Context<'_>,
     ) -> Poll<Result<StreamMuxerEvent<Self::Substream>, Self::Error>> {
         match self {
-            EitherOutput::First(inner) => inner.poll_event(cx).map(|result| {
-                result.map_err(|e| e.into()).map(|event| match event {
-                    StreamMuxerEvent::AddressChange(addr) => StreamMuxerEvent::AddressChange(addr),
-                    StreamMuxerEvent::InboundSubstream(substream) => {
-                        StreamMuxerEvent::InboundSubstream(EitherOutput::First(substream))
-                    }
-                })
-            }),
-            EitherOutput::Second(inner) => inner.poll_event(cx).map(|result| {
-                result.map_err(|e| e.into()).map(|event| match event {
-                    StreamMuxerEvent::AddressChange(addr) => StreamMuxerEvent::AddressChange(addr),
-                    StreamMuxerEvent::InboundSubstream(substream) => {
-                        StreamMuxerEvent::InboundSubstream(EitherOutput::Second(substream))
-                    }
-                })
-            }),
+            EitherOutput::First(inner) => inner
+                .poll_event(cx)
+                .map_err(|e| e.into())
+                .map_ok(|event| event.map_inbound_stream(EitherOutput::First)),
+            EitherOutput::Second(inner) => inner
+                .poll_event(cx)
+                .map_err(|e| e.into())
+                .map_ok(|event| event.map_inbound_stream(EitherOutput::Second)),
         }
     }
 

diff --git a/core/src/muxing.rs b/core/src/muxing.rs
@@ -236,6 +236,16 @@ impl<T> StreamMuxerEvent<T> {
             None
         }
     }
+
+    /// Map the stream within [`StreamMuxerEvent::InboundSubstream`] to a new type.
+    pub fn map_inbound_stream<O>(self, map: impl FnOnce(T) -> O) -> StreamMuxerEvent<O> {
+        match self {
+            StreamMuxerEvent::InboundSubstream(stream) => {
+                StreamMuxerEvent::InboundSubstream(map(stream))
+            }
+            StreamMuxerEvent::AddressChange(addr) => StreamMuxerEvent::AddressChange(addr),
+        }
+    }
 }
 
 /// Polls for an event from the muxer and, if an inbound substream, wraps this substream in an