diff --git a/library/std/src/sys/pal/wasi/fs.rs b/library/std/src/sys/pal/wasi/fs.rs index 0667eb9010100..7779d2b97d7f9 100644 --- a/library/std/src/sys/pal/wasi/fs.rs +++ b/library/std/src/sys/pal/wasi/fs.rs @@ -27,10 +27,27 @@ pub struct FileAttr { pub struct ReadDir { inner: Arc, - cookie: Option, - buf: Vec, - offset: usize, - cap: usize, + state: ReadDirState, +} + +enum ReadDirState { + /// Fill `buf` with `buf.len()` bytes starting from `next_read_offset`. + FillBuffer { + next_read_offset: wasi::Dircookie, + buf: Vec, + }, + ProcessEntry { + buf: Vec, + next_read_offset: Option, + offset: usize, + }, + /// There is no more data to get in [`Self::FillBuffer`]; keep returning + /// entries via ProcessEntry until `buf` is exhausted. + RunUntilExhaustion { + buf: Vec, + offset: usize, + }, + Done, } struct ReadDirInner { @@ -147,11 +164,8 @@ impl FileType { impl ReadDir { fn new(dir: File, root: PathBuf) -> ReadDir { ReadDir { - cookie: Some(0), - buf: vec![0; 128], - offset: 0, - cap: 0, inner: Arc::new(ReadDirInner { dir, root }), + state: ReadDirState::FillBuffer { next_read_offset: 0, buf: vec![0; 128] }, } } } @@ -162,78 +176,99 @@ impl fmt::Debug for ReadDir { } } +impl core::iter::FusedIterator for ReadDir {} + impl Iterator for ReadDir { type Item = io::Result; fn next(&mut self) -> Option> { - loop { - // If we've reached the capacity of our buffer then we need to read - // some more from the OS, otherwise we pick up at our old offset. - let offset = if self.offset == self.cap { - let cookie = self.cookie.take()?; - match self.inner.dir.fd.readdir(&mut self.buf, cookie) { - Ok(bytes) => self.cap = bytes, - Err(e) => return Some(Err(e)), - } - self.offset = 0; - self.cookie = Some(cookie); - - // If we didn't actually read anything, this is in theory the - // end of the directory. - if self.cap == 0 { - self.cookie = None; - return None; - } - - 0 - } else { - self.offset - }; - let data = &self.buf[offset..self.cap]; - - // If we're not able to read a directory entry then that means it - // must have been truncated at the end of the buffer, so reset our - // offset so we can go back and reread into the buffer, picking up - // where we last left off. - let dirent_size = mem::size_of::(); - if data.len() < dirent_size { - assert!(self.cookie.is_some()); - assert!(self.buf.len() >= dirent_size); - self.offset = self.cap; - continue; - } - let (dirent, data) = data.split_at(dirent_size); - let dirent = unsafe { ptr::read_unaligned(dirent.as_ptr() as *const wasi::Dirent) }; - - // If the file name was truncated, then we need to reinvoke - // `readdir` so we truncate our buffer to start over and reread this - // descriptor. Note that if our offset is 0 that means the file name - // is massive and we need a bigger buffer. - if data.len() < dirent.d_namlen as usize { - if offset == 0 { - let amt_to_add = self.buf.capacity(); - self.buf.extend(iter::repeat(0).take(amt_to_add)); + match &mut self.state { + ReadDirState::FillBuffer { next_read_offset, ref mut buf } => { + let result = self.inner.dir.fd.readdir(buf, *next_read_offset); + match result { + Ok(read_bytes) => { + if read_bytes < buf.len() { + buf.truncate(read_bytes); + self.state = + ReadDirState::RunUntilExhaustion { buf: mem::take(buf), offset: 0 }; + } else { + debug_assert_eq!(read_bytes, buf.len()); + self.state = ReadDirState::ProcessEntry { + buf: mem::take(buf), + offset: 0, + next_read_offset: Some(*next_read_offset), + }; + } + self.next() + } + Err(e) => { + self.state = ReadDirState::Done; + return Some(Err(e)); + } } - assert!(self.cookie.is_some()); - self.offset = self.cap; - continue; } - self.cookie = Some(dirent.d_next); - self.offset = offset + dirent_size + dirent.d_namlen as usize; + ReadDirState::ProcessEntry { ref mut buf, next_read_offset, offset } => { + let contents = &buf[*offset..]; + const DIRENT_SIZE: usize = crate::mem::size_of::(); + if contents.len() >= DIRENT_SIZE { + let (dirent, data) = contents.split_at(DIRENT_SIZE); + let dirent = + unsafe { ptr::read_unaligned(dirent.as_ptr() as *const wasi::Dirent) }; + // If the file name was truncated, then we need to reinvoke + // `readdir` so we truncate our buffer to start over and reread this + // descriptor. + if data.len() < dirent.d_namlen as usize { + if buf.len() < dirent.d_namlen as usize + DIRENT_SIZE { + buf.resize(dirent.d_namlen as usize + DIRENT_SIZE, 0); + } + if let Some(next_read_offset) = *next_read_offset { + self.state = + ReadDirState::FillBuffer { next_read_offset, buf: mem::take(buf) }; + } else { + self.state = ReadDirState::Done; + } + + return self.next(); + } + next_read_offset.as_mut().map(|cookie| { + *cookie = dirent.d_next; + }); + *offset = *offset + DIRENT_SIZE + dirent.d_namlen as usize; - let name = &data[..(dirent.d_namlen as usize)]; + let name = &data[..(dirent.d_namlen as usize)]; + + // These names are skipped on all other platforms, so let's skip + // them here too + if name == b"." || name == b".." { + return self.next(); + } - // These names are skipped on all other platforms, so let's skip - // them here too - if name == b"." || name == b".." { - continue; + return Some(Ok(DirEntry { + meta: dirent, + name: name.to_vec(), + inner: self.inner.clone(), + })); + } else if let Some(next_read_offset) = *next_read_offset { + self.state = ReadDirState::FillBuffer { next_read_offset, buf: mem::take(buf) }; + } else { + self.state = ReadDirState::Done; + } + self.next() } + ReadDirState::RunUntilExhaustion { buf, offset } => { + if *offset >= buf.len() { + self.state = ReadDirState::Done; + } else { + self.state = ReadDirState::ProcessEntry { + buf: mem::take(buf), + offset: *offset, + next_read_offset: None, + }; + } - return Some(Ok(DirEntry { - meta: dirent, - name: name.to_vec(), - inner: self.inner.clone(), - })); + self.next() + } + ReadDirState::Done => None, } } }