TcpStream.recv(): wrap callback in RefCell to allow FnMut

This commit is contained in:
Astro 2020-07-23 18:58:25 +02:00
parent 0aa75d3544
commit b74cc5f3aa

View File

@ -3,7 +3,9 @@
//! TODO: implement futures AsyncRead/AsyncWrite/Stream/Sink interfaces //! TODO: implement futures AsyncRead/AsyncWrite/Stream/Sink interfaces
use core::{ use core::{
cell::RefCell,
future::Future, future::Future,
ops::DerefMut,
pin::Pin, pin::Pin,
task::{Context, Poll}, task::{Context, Poll},
}; };
@ -108,17 +110,18 @@ impl TcpStream {
/// number of bytes it consumed, and a user-defined return value of type R. /// number of bytes it consumed, and a user-defined return value of type R.
pub async fn recv<F, R>(&self, f: F) -> Result<R> pub async fn recv<F, R>(&self, f: F) -> Result<R>
where where
F: Fn(&[u8]) -> (usize, R), F: FnMut(&[u8]) -> (usize, R),
{ {
struct Recv<'a, F: FnOnce(&[u8]) -> (usize, R), R> { struct Recv<'a, F: FnMut(&[u8]) -> (usize, R), R> {
stream: &'a TcpStream, stream: &'a TcpStream,
f: F, f: RefCell<F>,
} }
impl<'a, F: Fn(&[u8]) -> (usize, R), R> Future for Recv<'a, F, R> { impl<'a, F: FnMut(&[u8]) -> (usize, R), R> Future for Recv<'a, F, R> {
type Output = Result<R>; type Output = Result<R>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let mut f = self.f.borrow_mut();
let result = self.stream.with_socket(|mut socket| { let result = self.stream.with_socket(|mut socket| {
if socket_is_handhshaking(&socket) { if socket_is_handhshaking(&socket) {
return Ok(Poll::Pending); return Ok(Poll::Pending);
@ -126,7 +129,7 @@ impl TcpStream {
socket.recv(|buf| { socket.recv(|buf| {
if buf.len() > 0 { if buf.len() > 0 {
let (amount, result) = (self.f)(buf); let (amount, result) = (f.deref_mut())(buf);
assert!(amount > 0); assert!(amount > 0);
(amount, Poll::Ready(Ok(result))) (amount, Poll::Ready(Ok(result)))
} else { } else {
@ -150,7 +153,7 @@ impl TcpStream {
Recv { Recv {
stream: self, stream: self,
f, f: RefCell::new(f),
}.await }.await
} }