fs: update to use std::future (#1269)

tokio-fs/src/write.rs

@@ -1,7 +1,11 @@
 use crate::{file, File};
-use futures::{try_ready, Async, Future, Poll};
+use std::future::Future;
+use std::pin::Pin;
+use std::task::Context;
+use std::task::Poll;
 use std::{fmt, io, mem, path::Path};
 use tokio_io;
+use tokio_io::AsyncWrite;
 
 /// Creates a future that will open a file for writing and write the entire
 /// contents of `contents` to it.
@@ -25,9 +29,9 @@ use tokio_io;
 ///
 /// tokio::run(task);
 /// ```
-pub fn write<P, C: AsRef<[u8]>>(path: P, contents: C) -> WriteFile<P, C>
+pub fn write<P, C: AsRef<[u8]> + Unpin>(path: P, contents: C) -> WriteFile<P, C>
 where
-    P: AsRef<Path> + Send + 'static,
+    P: AsRef<Path> + Send + Unpin + 'static,
 {
     WriteFile {
         state: State::Create(File::create(path), Some(contents)),
@@ -37,35 +41,62 @@ where
 /// A future used to open a file for writing and write the entire contents
 /// of some data to it.
 #[derive(Debug)]
-pub struct WriteFile<P: AsRef<Path> + Send + 'static, C: AsRef<[u8]>> {
+pub struct WriteFile<P: AsRef<Path> + Send + Unpin + 'static, C: AsRef<[u8]> + Unpin> {
     state: State<P, C>,
 }
 
 #[derive(Debug)]
-enum State<P: AsRef<Path> + Send + 'static, C: AsRef<[u8]>> {
+enum State<P: AsRef<Path> + Send + Unpin + 'static, C: AsRef<[u8]> + Unpin> {
     Create(file::CreateFuture<P>, Option<C>),
-    Write(tokio_io::io::WriteAll<File, C>),
+    Writing { f: File, buf: C, pos: usize },
+    Empty,
 }
 
-impl<P: AsRef<Path> + Send + 'static, C: AsRef<[u8]> + fmt::Debug> Future for WriteFile<P, C> {
-    type Item = C;
-    type Error = io::Error;
+fn zero_write() -> io::Error {
+    io::Error::new(io::ErrorKind::WriteZero, "zero-length write")
+}
 
-    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
-        let new_state = match &mut self.state {
-            State::Create(ref mut create_file, contents) => {
-                let file = try_ready!(create_file.poll());
-                let write = tokio_io::io::write_all(file, contents.take().unwrap());
-                State::Write(write)
-            }
-            State::Write(ref mut write) => {
-                let (_, contents) = try_ready!(write.poll());
-                return Ok(Async::Ready(contents));
-            }
-        };
+impl<P: AsRef<Path> + Send + Unpin + 'static, C: AsRef<[u8]> + Unpin + fmt::Debug> Future
+    for WriteFile<P, C>
+{
+    type Output = io::Result<C>;
 
-        mem::replace(&mut self.state, new_state);
-        // We just entered the Write state, need to poll it before returning.
-        self.poll()
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let inner = Pin::get_mut(self);
+        match &mut inner.state {
+            State::Create(create_file, contents) => {
+                let file = ready!(Pin::new(create_file).poll(cx))?;
+                let contents = contents.take().unwrap();
+                let new_state = State::Writing {
+                    f: file,
+                    buf: contents,
+                    pos: 0,
+                };
+                mem::replace(&mut inner.state, new_state);
+                // We just entered the Write state, need to poll it before returning.
+                return Pin::new(inner).poll(cx);
+            }
+            State::Empty => panic!("poll a WriteFile after it's done"),
+            _ => {}
+        }
+
+        match mem::replace(&mut inner.state, State::Empty) {
+            State::Writing {
+                mut f,
+                buf,
+                mut pos,
+            } => {
+                let buf_ref = buf.as_ref();
+                while pos < buf_ref.len() {
+                    let n = ready!(Pin::new(&mut f).poll_write(cx, &buf_ref[pos..]))?;
+                    pos += n;
+                    if n == 0 {
+                        return Poll::Ready(Err(zero_write()));
+                    }
+                }
+                Poll::Ready(Ok(buf))
+            }
+            _ => panic!(),
+        }
     }
 }