From 3c120ef971d909f09ec0773d7ea91392aab250a5 Mon Sep 17 00:00:00 2001
From: Thomas Eizinger <thomas@eizinger.io>
Date: Wed, 15 Jun 2022 21:33:32 +0200
Subject: [PATCH] core/muxing: Introduce `StreamMuxerEvent::map_inbound_stream`
 (#2691)

Co-authored-by: Max Inden <mail@max-inden.de>
---
 core/CHANGELOG.md  |  6 ++++++
 core/Cargo.toml    |  2 +-
 core/src/either.rs | 24 ++++++++----------------
 core/src/muxing.rs | 10 ++++++++++
 4 files changed, 25 insertions(+), 17 deletions(-)

diff --git a/core/CHANGELOG.md b/core/CHANGELOG.md
index bebe025d468..cdd252351c5 100644
--- 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
index 41e0c2a1763..b6a45248a56 100644
--- 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 <admin@parity.io>"]
 license = "MIT"
 repository = "https://github.com/libp2p/rust-libp2p"
diff --git a/core/src/either.rs b/core/src/either.rs
index e1c5316fd6a..a9e51a47e79 100644
--- 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
index 969c56c782e..baabf656bcc 100644
--- 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