From b0637668cb94427f86414d08840cdb20674ac67d Mon Sep 17 00:00:00 2001
From: Wim Looman <wim@nemo157.com>
Date: Fri, 23 Aug 2019 20:23:40 +0200
Subject: [PATCH] Add extension trait for getting futures from AsyncDatagram

---
 src/ext.rs | 79 ++++++++++++++++++++++++++++++++++++++++++++++++++++++
 src/lib.rs | 11 ++++++--
 2 files changed, 88 insertions(+), 2 deletions(-)
 create mode 100644 src/ext.rs

diff --git a/src/ext.rs b/src/ext.rs
new file mode 100644
index 0000000..971e108
--- /dev/null
+++ b/src/ext.rs
@@ -0,0 +1,79 @@
+use core::{
+  future::Future,
+  pin::Pin,
+  task::{Context, Poll},
+};
+
+use crate::AsyncDatagram;
+
+/// An extension trait that adds utility methods to [`AsyncDatagram`] types.
+pub trait AsyncDatagramExt: AsyncDatagram {
+  /// Creates a future that will send a given message to the specified target.
+  ///
+  /// The returned future will resolve to the number of bytes actually sent once the operation
+  /// completes.
+  fn send_to<'a>(
+    &'a mut self,
+    buf: &'a [u8],
+    receiver: &'a Self::Receiver,
+  ) -> SendTo<'a, Self>
+  where
+    Self: Unpin,
+  {
+    SendTo {
+      socket: self,
+      buf,
+      receiver,
+    }
+  }
+
+  /// Creates a future that will receive a message into the provided buffer.
+  ///
+  /// The returned future will resolve to the number of bytes received and target from whence the
+  /// data came  once the operation completes.
+  fn recv_from<'a>(&'a mut self, buf: &'a mut [u8]) -> RecvFrom<'a, Self>
+  where
+    Self: Unpin,
+  {
+    RecvFrom { socket: self, buf }
+  }
+}
+
+impl<T: AsyncDatagram> AsyncDatagramExt for T {}
+
+/// Future for the [`send_to`](AsyncDatagramExt::send_to) method.
+#[derive(Debug)]
+pub struct SendTo<'a, T: AsyncDatagram + ?Sized> {
+  socket: &'a mut T,
+  buf: &'a [u8],
+  receiver: &'a T::Receiver,
+}
+
+impl<T: AsyncDatagram + Unpin + ?Sized> Future for SendTo<'_, T> {
+  type Output = Result<usize, T::Err>;
+
+  fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+    let Self {
+      socket,
+      buf,
+      receiver,
+    } = self.get_mut();
+    Pin::new(&mut **socket).poll_send_to(cx, buf, receiver)
+  }
+}
+
+/// Future for the [`recv_from`](AsyncDatagramExt::recv_from) method.
+#[derive(Debug)]
+pub struct RecvFrom<'a, T: AsyncDatagram + ?Sized> {
+  socket: &'a mut T,
+  buf: &'a mut [u8],
+}
+
+impl<T: AsyncDatagram + Unpin + ?Sized> Future for RecvFrom<'_, T> {
+  type Output = Result<(usize, T::Sender), T::Err>;
+
+  fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+    let Self { socket, buf } = self.get_mut();
+    Pin::new(&mut **socket).poll_recv_from(cx, buf)
+  }
+}
diff --git a/src/lib.rs b/src/lib.rs
index 3425807..b4682f6 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -2,6 +2,7 @@
 #![deny(missing_debug_implementations, bad_style)]
 #![deny(missing_docs)]
 #![cfg_attr(test, deny(warnings))]
+#![no_std]
 
 //! Async datagram traits.
 //!
@@ -38,8 +39,14 @@
 //! }
 //! ```
 
-use std::pin::Pin;
-use std::task::{Context, Poll};
+mod ext;
+
+pub use ext::{AsyncDatagramExt, RecvFrom, SendTo};
+
+use core::{
+  pin::Pin,
+  task::{Context, Poll},
+};
 
 /// Implement a datagram protocol.
 pub trait AsyncDatagram {