From cbd458b1db5c9795415d9caad8e91eb52e66dcd4 Mon Sep 17 00:00:00 2001
From: Stjepan Glavina <stjepang@gmail.com>
Date: Wed, 9 Oct 2019 14:26:14 +0200
Subject: [PATCH] Cleanup ToSocketAddrs, add more net reexports

---
 src/net/addr.rs | 191 ++++++++++++++++++++++++++++++++++++------------
 src/net/mod.rs  |   5 ++
 2 files changed, 151 insertions(+), 45 deletions(-)

diff --git a/src/net/addr.rs b/src/net/addr.rs
index baa41c8..01c58fd 100644
--- a/src/net/addr.rs
+++ b/src/net/addr.rs
@@ -1,4 +1,4 @@
-use std::marker::PhantomData;
+use std::mem;
 use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
 use std::net::{SocketAddr, SocketAddrV4, SocketAddrV6};
 use std::pin::Pin;
@@ -12,24 +12,38 @@ use crate::task::{blocking, Context, JoinHandle, Poll};
 cfg_if! {
     if #[cfg(feature = "docs")] {
         #[doc(hidden)]
-        pub struct ImplFuture<'a, T>(std::marker::PhantomData<&'a T>);
+        pub struct ImplFuture<T>(std::marker::PhantomData<T>);
 
         macro_rules! ret {
-            ($a:lifetime, $f:tt, $i:ty) => (ImplFuture<$a, io::Result<$i>>);
+            (impl Future<Output = $out:ty>, $fut:ty) => (ImplFuture<$out>);
         }
     } else {
         macro_rules! ret {
-            ($a:lifetime, $f:tt, $i:ty) => ($f<$a, $i>);
+            (impl Future<Output = $out:ty>, $fut:ty) => ($fut);
         }
     }
 }
 
-/// A trait for objects which can be converted or resolved to one or more [`SocketAddr`] values.
+/// Converts or resolves addresses to [`SocketAddr`] values.
 ///
 /// This trait is an async version of [`std::net::ToSocketAddrs`].
 ///
 /// [`std::net::ToSocketAddrs`]: https://doc.rust-lang.org/std/net/trait.ToSocketAddrs.html
-/// [`SocketAddr`]: https://doc.rust-lang.org/std/net/enum.SocketAddr.html
+/// [`SocketAddr`]: enum.SocketAddr.html
+///
+/// # Examples
+///
+/// ```
+/// # fn main() -> std::io::Result<()> { async_std::task::block_on(async {
+/// #
+/// use async_std::net::ToSocketAddrs;
+///
+/// let addr1 = "localhost:8080".to_socket_addrs().await?.next().unwrap();
+/// let addr2 = ("127.0.0.1", 8080).to_socket_addrs().await?.next().unwrap();
+/// assert_eq!(addr1, addr2);
+/// #
+/// # Ok(()) }) }
+/// ```
 pub trait ToSocketAddrs {
     /// Returned iterator over socket addresses which this type may correspond to.
     type Iter: Iterator<Item = SocketAddr>;
@@ -40,28 +54,39 @@ pub trait ToSocketAddrs {
     /// resolution performed.
     ///
     /// Note that this function may block a backend thread while resolution is performed.
-    fn to_socket_addrs(&self) -> ret!('_, ToSocketAddrsFuture, Self::Iter);
+    fn to_socket_addrs(
+        &self,
+    ) -> ret!(
+        impl Future<Output = Self::Iter>,
+        ToSocketAddrsFuture<Self::Iter>
+    );
 }
 
 #[doc(hidden)]
 #[allow(missing_debug_implementations)]
-pub enum ToSocketAddrsFuture<'a, I> {
-    Phantom(PhantomData<&'a ()>),
-    Join(JoinHandle<io::Result<I>>),
-    Ready(Option<io::Result<I>>),
+pub enum ToSocketAddrsFuture<I> {
+    Resolving(JoinHandle<io::Result<I>>),
+    Ready(io::Result<I>),
+    Done,
 }
 
-impl<I: Iterator<Item = SocketAddr>> Future for ToSocketAddrsFuture<'_, I> {
+impl<I: Iterator<Item = SocketAddr>> Future for ToSocketAddrsFuture<I> {
     type Output = io::Result<I>;
 
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        match unsafe { self.get_unchecked_mut() } {
-            ToSocketAddrsFuture::Join(f) => Pin::new(&mut *f).poll(cx),
-            ToSocketAddrsFuture::Ready(res) => {
-                let res = res.take().expect("polled a completed future");
-                Poll::Ready(res)
+        let this = unsafe { self.get_unchecked_mut() };
+        let state = mem::replace(this, ToSocketAddrsFuture::Done);
+
+        match state {
+            ToSocketAddrsFuture::Resolving(mut task) => {
+                let poll = Pin::new(&mut task).poll(cx);
+                if poll.is_pending() {
+                    *this = ToSocketAddrsFuture::Resolving(task);
+                }
+                poll
             }
-            _ => unreachable!(),
+            ToSocketAddrsFuture::Ready(res) => Poll::Ready(res),
+            ToSocketAddrsFuture::Done => panic!("polled a completed future"),
         }
     }
 }
@@ -69,87 +94,158 @@ impl<I: Iterator<Item = SocketAddr>> Future for ToSocketAddrsFuture<'_, I> {
 impl ToSocketAddrs for SocketAddr {
     type Iter = std::option::IntoIter<SocketAddr>;
 
-    fn to_socket_addrs(&self) -> ret!('_, ToSocketAddrsFuture, Self::Iter) {
-        ToSocketAddrsFuture::Ready(Some(std::net::ToSocketAddrs::to_socket_addrs(self)))
+    fn to_socket_addrs(
+        &self,
+    ) -> ret!(
+        impl Future<Output = Self::Iter>,
+        ToSocketAddrsFuture<Self::Iter>
+    ) {
+        ToSocketAddrsFuture::Ready(Ok(Some(*self).into_iter()))
     }
 }
 
 impl ToSocketAddrs for SocketAddrV4 {
     type Iter = std::option::IntoIter<SocketAddr>;
 
-    fn to_socket_addrs(&self) -> ret!('_, ToSocketAddrsFuture, Self::Iter) {
-        ToSocketAddrsFuture::Ready(Some(std::net::ToSocketAddrs::to_socket_addrs(self)))
+    fn to_socket_addrs(
+        &self,
+    ) -> ret!(
+        impl Future<Output = Self::Iter>,
+        ToSocketAddrsFuture<Self::Iter>
+    ) {
+        SocketAddr::V4(*self).to_socket_addrs()
     }
 }
 
 impl ToSocketAddrs for SocketAddrV6 {
     type Iter = std::option::IntoIter<SocketAddr>;
 
-    fn to_socket_addrs(&self) -> ret!('_, ToSocketAddrsFuture, Self::Iter) {
-        ToSocketAddrsFuture::Ready(Some(std::net::ToSocketAddrs::to_socket_addrs(self)))
+    fn to_socket_addrs(
+        &self,
+    ) -> ret!(
+        impl Future<Output = Self::Iter>,
+        ToSocketAddrsFuture<Self::Iter>
+    ) {
+        SocketAddr::V6(*self).to_socket_addrs()
     }
 }
 
 impl ToSocketAddrs for (IpAddr, u16) {
     type Iter = std::option::IntoIter<SocketAddr>;
 
-    fn to_socket_addrs(&self) -> ret!('_, ToSocketAddrsFuture, Self::Iter) {
-        ToSocketAddrsFuture::Ready(Some(std::net::ToSocketAddrs::to_socket_addrs(self)))
+    fn to_socket_addrs(
+        &self,
+    ) -> ret!(
+        impl Future<Output = Self::Iter>,
+        ToSocketAddrsFuture<Self::Iter>
+    ) {
+        let (ip, port) = *self;
+        match ip {
+            IpAddr::V4(a) => (a, port).to_socket_addrs(),
+            IpAddr::V6(a) => (a, port).to_socket_addrs(),
+        }
     }
 }
 
 impl ToSocketAddrs for (Ipv4Addr, u16) {
     type Iter = std::option::IntoIter<SocketAddr>;
 
-    fn to_socket_addrs(&self) -> ret!('_, ToSocketAddrsFuture, Self::Iter) {
-        ToSocketAddrsFuture::Ready(Some(std::net::ToSocketAddrs::to_socket_addrs(self)))
+    fn to_socket_addrs(
+        &self,
+    ) -> ret!(
+        impl Future<Output = Self::Iter>,
+        ToSocketAddrsFuture<Self::Iter>
+    ) {
+        let (ip, port) = *self;
+        SocketAddrV4::new(ip, port).to_socket_addrs()
     }
 }
 
 impl ToSocketAddrs for (Ipv6Addr, u16) {
     type Iter = std::option::IntoIter<SocketAddr>;
 
-    fn to_socket_addrs(&self) -> ret!('_, ToSocketAddrsFuture, Self::Iter) {
-        ToSocketAddrsFuture::Ready(Some(std::net::ToSocketAddrs::to_socket_addrs(self)))
+    fn to_socket_addrs(
+        &self,
+    ) -> ret!(
+        impl Future<Output = Self::Iter>,
+        ToSocketAddrsFuture<Self::Iter>
+    ) {
+        let (ip, port) = *self;
+        SocketAddrV6::new(ip, port, 0, 0).to_socket_addrs()
     }
 }
 
 impl ToSocketAddrs for (&str, u16) {
     type Iter = std::vec::IntoIter<SocketAddr>;
 
-    fn to_socket_addrs(&self) -> ret!('_, ToSocketAddrsFuture, Self::Iter) {
-        let host = self.0.to_string();
-        let port = self.1;
-        let join = blocking::spawn(async move {
+    fn to_socket_addrs(
+        &self,
+    ) -> ret!(
+        impl Future<Output = Self::Iter>,
+        ToSocketAddrsFuture<Self::Iter>
+    ) {
+        let (host, port) = *self;
+
+        if let Ok(addr) = host.parse::<Ipv4Addr>() {
+            let addr = SocketAddrV4::new(addr, port);
+            return ToSocketAddrsFuture::Ready(Ok(vec![SocketAddr::V4(addr)].into_iter()));
+        }
+
+        if let Ok(addr) = host.parse::<Ipv6Addr>() {
+            let addr = SocketAddrV6::new(addr, port, 0, 0);
+            return ToSocketAddrsFuture::Ready(Ok(vec![SocketAddr::V6(addr)].into_iter()));
+        }
+
+        let host = host.to_string();
+        let task = blocking::spawn(async move {
             std::net::ToSocketAddrs::to_socket_addrs(&(host.as_str(), port))
         });
-        ToSocketAddrsFuture::Join(join)
+        ToSocketAddrsFuture::Resolving(task)
     }
 }
 
 impl ToSocketAddrs for str {
     type Iter = std::vec::IntoIter<SocketAddr>;
 
-    fn to_socket_addrs(&self) -> ret!('_, ToSocketAddrsFuture, Self::Iter) {
-        let socket_addrs = self.to_string();
-        let join =
-            blocking::spawn(async move { std::net::ToSocketAddrs::to_socket_addrs(&socket_addrs) });
-        ToSocketAddrsFuture::Join(join)
+    fn to_socket_addrs(
+        &self,
+    ) -> ret!(
+        impl Future<Output = Self::Iter>,
+        ToSocketAddrsFuture<Self::Iter>
+    ) {
+        if let Some(addr) = self.parse().ok() {
+            return ToSocketAddrsFuture::Ready(Ok(vec![addr].into_iter()));
+        }
+
+        let addr = self.to_string();
+        let task =
+            blocking::spawn(async move { std::net::ToSocketAddrs::to_socket_addrs(addr.as_str()) });
+        ToSocketAddrsFuture::Resolving(task)
     }
 }
 
 impl<'a> ToSocketAddrs for &'a [SocketAddr] {
     type Iter = std::iter::Cloned<std::slice::Iter<'a, SocketAddr>>;
 
-    fn to_socket_addrs(&self) -> ret!('_, ToSocketAddrsFuture, Self::Iter) {
-        ToSocketAddrsFuture::Ready(Some(std::net::ToSocketAddrs::to_socket_addrs(self)))
+    fn to_socket_addrs(
+        &self,
+    ) -> ret!(
+        impl Future<Output = Self::Iter>,
+        ToSocketAddrsFuture<Self::Iter>
+    ) {
+        ToSocketAddrsFuture::Ready(Ok(self.iter().cloned()))
     }
 }
 
 impl<T: ToSocketAddrs + ?Sized> ToSocketAddrs for &T {
     type Iter = T::Iter;
 
-    fn to_socket_addrs(&self) -> ret!('_, ToSocketAddrsFuture, Self::Iter) {
+    fn to_socket_addrs(
+        &self,
+    ) -> ret!(
+        impl Future<Output = Self::Iter>,
+        ToSocketAddrsFuture<Self::Iter>
+    ) {
         (**self).to_socket_addrs()
     }
 }
@@ -157,7 +253,12 @@ impl<T: ToSocketAddrs + ?Sized> ToSocketAddrs for &T {
 impl ToSocketAddrs for String {
     type Iter = std::vec::IntoIter<SocketAddr>;
 
-    fn to_socket_addrs(&self) -> ret!('_, ToSocketAddrsFuture, Self::Iter) {
-        ToSocketAddrs::to_socket_addrs(self.as_str())
+    fn to_socket_addrs(
+        &self,
+    ) -> ret!(
+        impl Future<Output = Self::Iter>,
+        ToSocketAddrsFuture<Self::Iter>
+    ) {
+        (&**self).to_socket_addrs()
     }
 }
diff --git a/src/net/mod.rs b/src/net/mod.rs
index 259dc1d..b3ae287 100644
--- a/src/net/mod.rs
+++ b/src/net/mod.rs
@@ -28,6 +28,11 @@
 //! # }) }
 //! ```
 
+pub use std::net::AddrParseError;
+pub use std::net::Shutdown;
+pub use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
+pub use std::net::{SocketAddr, SocketAddrV4, SocketAddrV6};
+
 pub use addr::ToSocketAddrs;
 pub use tcp::{Incoming, TcpListener, TcpStream};
 pub use udp::UdpSocket;