From 187cf495c9f46ac059d5e7ef6e722849a7a870b2 Mon Sep 17 00:00:00 2001
From: Konrad Borowski <konrad@borowski.pw>
Date: Mon, 20 Feb 2023 10:47:18 +0100
Subject: [PATCH] use poll_fn in copy_bidirectional

---
 tokio/src/io/util/copy_bidirectional.rs | 57 ++++++-------------------
 1 file changed, 14 insertions(+), 43 deletions(-)

diff --git a/tokio/src/io/util/copy_bidirectional.rs b/tokio/src/io/util/copy_bidirectional.rs
index c93060b361a..e1a7db127a7 100644
--- a/tokio/src/io/util/copy_bidirectional.rs
+++ b/tokio/src/io/util/copy_bidirectional.rs
@@ -1,8 +1,8 @@
 use super::copy::CopyBuffer;
 
+use crate::future::poll_fn;
 use crate::io::{AsyncRead, AsyncWrite};
 
-use std::future::Future;
 use std::io;
 use std::pin::Pin;
 use std::task::{Context, Poll};
@@ -13,13 +13,6 @@ enum TransferState {
     Done(u64),
 }
 
-struct CopyBidirectional<'a, A: ?Sized, B: ?Sized> {
-    a: &'a mut A,
-    b: &'a mut B,
-    a_to_b: TransferState,
-    b_to_a: TransferState,
-}
-
 fn transfer_one_direction<A, B>(
     cx: &mut Context<'_>,
     state: &mut TransferState,
@@ -48,35 +41,6 @@ where
         }
     }
 }
-
-impl<'a, A, B> Future for CopyBidirectional<'a, A, B>
-where
-    A: AsyncRead + AsyncWrite + Unpin + ?Sized,
-    B: AsyncRead + AsyncWrite + Unpin + ?Sized,
-{
-    type Output = io::Result<(u64, u64)>;
-
-    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        // Unpack self into mut refs to each field to avoid borrow check issues.
-        let CopyBidirectional {
-            a,
-            b,
-            a_to_b,
-            b_to_a,
-        } = &mut *self;
-
-        let a_to_b = transfer_one_direction(cx, a_to_b, &mut *a, &mut *b)?;
-        let b_to_a = transfer_one_direction(cx, b_to_a, &mut *b, &mut *a)?;
-
-        // It is not a problem if ready! returns early because transfer_one_direction for the
-        // other direction will keep returning TransferState::Done(count) in future calls to poll
-        let a_to_b = ready!(a_to_b);
-        let b_to_a = ready!(b_to_a);
-
-        Poll::Ready(Ok((a_to_b, b_to_a)))
-    }
-}
-
 /// Copies data in both directions between `a` and `b`.
 ///
 /// This function returns a future that will read from both streams,
@@ -110,11 +74,18 @@ where
     A: AsyncRead + AsyncWrite + Unpin + ?Sized,
     B: AsyncRead + AsyncWrite + Unpin + ?Sized,
 {
-    CopyBidirectional {
-        a,
-        b,
-        a_to_b: TransferState::Running(CopyBuffer::new()),
-        b_to_a: TransferState::Running(CopyBuffer::new()),
-    }
+    let mut a_to_b = TransferState::Running(CopyBuffer::new());
+    let mut b_to_a = TransferState::Running(CopyBuffer::new());
+    poll_fn(|cx| {
+        let a_to_b = transfer_one_direction(cx, &mut a_to_b, a, b)?;
+        let b_to_a = transfer_one_direction(cx, &mut b_to_a, b, a)?;
+
+        // It is not a problem if ready! returns early because transfer_one_direction for the
+        // other direction will keep returning TransferState::Done(count) in future calls to poll
+        let a_to_b = ready!(a_to_b);
+        let b_to_a = ready!(b_to_a);
+
+        Poll::Ready(Ok((a_to_b, b_to_a)))
+    })
     .await
 }