1 files changed, 30 insertions, 25 deletions

diff --git a/src/fs/read_dir.rs b/src/fs/read_dir.rs
index 9471e8c..def735b 100644
--- a/src/fs/read_dir.rs
+++ b/src/fs/read_dir.rs
@@ -35,9 +35,9 @@ pub async fn read_dir(path: impl AsRef<Path>) -> io::Result<ReadDir> {
     asyncify(|| -> io::Result<ReadDir> {
         let mut std = std::fs::read_dir(path)?;
         let mut buf = VecDeque::with_capacity(CHUNK_SIZE);
-        ReadDir::next_chunk(&mut buf, &mut std);
+        let remain = ReadDir::next_chunk(&mut buf, &mut std);
 
-        Ok(ReadDir(State::Idle(Some((buf, std)))))
+        Ok(ReadDir(State::Idle(Some((buf, std, remain)))))
     })
     .await
 }
@@ -66,8 +66,8 @@ pub struct ReadDir(State);
 
 #[derive(Debug)]
 enum State {
-    Idle(Option<(VecDeque<io::Result<DirEntry>>, std::fs::ReadDir)>),
-    Pending(JoinHandle<(VecDeque<io::Result<DirEntry>>, std::fs::ReadDir)>),
+    Idle(Option<(VecDeque<io::Result<DirEntry>>, std::fs::ReadDir, bool)>),
+    Pending(JoinHandle<(VecDeque<io::Result<DirEntry>>, std::fs::ReadDir, bool)>),
 }
 
 impl ReadDir {
@@ -103,38 +103,35 @@ impl ReadDir {
         loop {
             match self.0 {
                 State::Idle(ref mut data) => {
-                    let (buf, _) = data.as_mut().unwrap();
+                    let (buf, _, ref remain) = data.as_mut().unwrap();
 
                     if let Some(ent) = buf.pop_front() {
                         return Poll::Ready(ent.map(Some));
-                    };
+                    } else if !remain {
+                        return Poll::Ready(Ok(None));
+                    }
 
-                    let (mut buf, mut std) = data.take().unwrap();
+                    let (mut buf, mut std, _) = data.take().unwrap();
 
                     self.0 = State::Pending(spawn_blocking(move || {
-                        ReadDir::next_chunk(&mut buf, &mut std);
-                        (buf, std)
+                        let remain = ReadDir::next_chunk(&mut buf, &mut std);
+                        (buf, std, remain)
                     }));
                 }
                 State::Pending(ref mut rx) => {
-                    let (mut buf, std) = ready!(Pin::new(rx).poll(cx))?;
-
-                    let ret = match buf.pop_front() {
-                        Some(Ok(x)) => Ok(Some(x)),
-                        Some(Err(e)) => Err(e),
-                        None => Ok(None),
-                    };
-
-                    self.0 = State::Idle(Some((buf, std)));
-
-                    return Poll::Ready(ret);
+                    self.0 = State::Idle(Some(ready!(Pin::new(rx).poll(cx))?));
                 }
             }
         }
     }
 
-    fn next_chunk(buf: &mut VecDeque<io::Result<DirEntry>>, std: &mut std::fs::ReadDir) {
-        for ret in std.by_ref().take(CHUNK_SIZE) {
+    fn next_chunk(buf: &mut VecDeque<io::Result<DirEntry>>, std: &mut std::fs::ReadDir) -> bool {
+        for _ in 0..CHUNK_SIZE {
+            let ret = match std.next() {
+                Some(ret) => ret,
+                None => return false,
+            };
+
             let success = ret.is_ok();
 
             buf.push_back(ret.map(|std| DirEntry {
@@ -142,7 +139,9 @@ impl ReadDir {
                     target_os = "solaris",
                     target_os = "illumos",
                     target_os = "haiku",
-                    target_os = "vxworks"
+                    target_os = "vxworks",
+                    target_os = "nto",
+                    target_os = "vita",
                 )))]
                 file_type: std.file_type().ok(),
                 std: Arc::new(std),
@@ -152,6 +151,8 @@ impl ReadDir {
                 break;
             }
         }
+
+        true
     }
 }
 
@@ -201,7 +202,9 @@ pub struct DirEntry {
         target_os = "solaris",
         target_os = "illumos",
         target_os = "haiku",
-        target_os = "vxworks"
+        target_os = "vxworks",
+        target_os = "nto",
+        target_os = "vita",
     )))]
     file_type: Option<FileType>,
     std: Arc<std::fs::DirEntry>,
@@ -332,7 +335,9 @@ impl DirEntry {
             target_os = "solaris",
             target_os = "illumos",
             target_os = "haiku",
-            target_os = "vxworks"
+            target_os = "vxworks",
+            target_os = "nto",
+            target_os = "vita",
         )))]
         if let Some(file_type) = self.file_type {
             return Ok(file_type);