@@ -27,10 +27,27 @@ pub struct FileAttr {
2727
2828pub struct ReadDir {
2929 inner : Arc < ReadDirInner > ,
30- cookie : Option < wasi:: Dircookie > ,
31- buf : Vec < u8 > ,
32- offset : usize ,
33- cap : usize ,
30+ state : ReadDirState ,
31+ }
32+
33+ enum ReadDirState {
34+ /// Fill `buf` with `buf.len()` bytes starting from `next_read_offset`.
35+ FillBuffer {
36+ next_read_offset : wasi:: Dircookie ,
37+ buf : Vec < u8 > ,
38+ } ,
39+ ProcessEntry {
40+ buf : Vec < u8 > ,
41+ next_read_offset : Option < wasi:: Dircookie > ,
42+ offset : usize ,
43+ } ,
44+ /// There is no more data to get in [`Self::FillBuffer`]; keep returning
45+ /// entries via ProcessEntry until `buf` is exhausted.
46+ RunUntilExhaustion {
47+ buf : Vec < u8 > ,
48+ offset : usize ,
49+ } ,
50+ Done ,
3451}
3552
3653struct ReadDirInner {
@@ -147,11 +164,8 @@ impl FileType {
147164impl ReadDir {
148165 fn new ( dir : File , root : PathBuf ) -> ReadDir {
149166 ReadDir {
150- cookie : Some ( 0 ) ,
151- buf : vec ! [ 0 ; 128 ] ,
152- offset : 0 ,
153- cap : 0 ,
154167 inner : Arc :: new ( ReadDirInner { dir, root } ) ,
168+ state : ReadDirState :: FillBuffer { next_read_offset : 0 , buf : vec ! [ 0 ; 128 ] } ,
155169 }
156170 }
157171}
@@ -162,78 +176,99 @@ impl fmt::Debug for ReadDir {
162176 }
163177}
164178
179+ impl core:: iter:: FusedIterator for ReadDir { }
180+
165181impl Iterator for ReadDir {
166182 type Item = io:: Result < DirEntry > ;
167183
168184 fn next ( & mut self ) -> Option < io:: Result < DirEntry > > {
169- loop {
170- // If we've reached the capacity of our buffer then we need to read
171- // some more from the OS, otherwise we pick up at our old offset.
172- let offset = if self . offset == self . cap {
173- let cookie = self . cookie . take ( ) ?;
174- match self . inner . dir . fd . readdir ( & mut self . buf , cookie) {
175- Ok ( bytes) => self . cap = bytes,
176- Err ( e) => return Some ( Err ( e) ) ,
177- }
178- self . offset = 0 ;
179- self . cookie = Some ( cookie) ;
180-
181- // If we didn't actually read anything, this is in theory the
182- // end of the directory.
183- if self . cap == 0 {
184- self . cookie = None ;
185- return None ;
186- }
187-
188- 0
189- } else {
190- self . offset
191- } ;
192- let data = & self . buf [ offset..self . cap ] ;
193-
194- // If we're not able to read a directory entry then that means it
195- // must have been truncated at the end of the buffer, so reset our
196- // offset so we can go back and reread into the buffer, picking up
197- // where we last left off.
198- let dirent_size = mem:: size_of :: < wasi:: Dirent > ( ) ;
199- if data. len ( ) < dirent_size {
200- assert ! ( self . cookie. is_some( ) ) ;
201- assert ! ( self . buf. len( ) >= dirent_size) ;
202- self . offset = self . cap ;
203- continue ;
204- }
205- let ( dirent, data) = data. split_at ( dirent_size) ;
206- let dirent = unsafe { ptr:: read_unaligned ( dirent. as_ptr ( ) as * const wasi:: Dirent ) } ;
207-
208- // If the file name was truncated, then we need to reinvoke
209- // `readdir` so we truncate our buffer to start over and reread this
210- // descriptor. Note that if our offset is 0 that means the file name
211- // is massive and we need a bigger buffer.
212- if data. len ( ) < dirent. d_namlen as usize {
213- if offset == 0 {
214- let amt_to_add = self . buf . capacity ( ) ;
215- self . buf . extend ( iter:: repeat ( 0 ) . take ( amt_to_add) ) ;
185+ match & mut self . state {
186+ ReadDirState :: FillBuffer { next_read_offset, ref mut buf } => {
187+ let result = self . inner . dir . fd . readdir ( buf, * next_read_offset) ;
188+ match result {
189+ Ok ( read_bytes) => {
190+ if read_bytes < buf. len ( ) {
191+ buf. truncate ( read_bytes) ;
192+ self . state =
193+ ReadDirState :: RunUntilExhaustion { buf : mem:: take ( buf) , offset : 0 } ;
194+ } else {
195+ debug_assert_eq ! ( read_bytes, buf. len( ) ) ;
196+ self . state = ReadDirState :: ProcessEntry {
197+ buf : mem:: take ( buf) ,
198+ offset : 0 ,
199+ next_read_offset : Some ( * next_read_offset) ,
200+ } ;
201+ }
202+ self . next ( )
203+ }
204+ Err ( e) => {
205+ self . state = ReadDirState :: Done ;
206+ return Some ( Err ( e) ) ;
207+ }
216208 }
217- assert ! ( self . cookie. is_some( ) ) ;
218- self . offset = self . cap ;
219- continue ;
220209 }
221- self . cookie = Some ( dirent. d_next ) ;
222- self . offset = offset + dirent_size + dirent. d_namlen as usize ;
210+ ReadDirState :: ProcessEntry { ref mut buf, next_read_offset, offset } => {
211+ let contents = & buf[ * offset..] ;
212+ const DIRENT_SIZE : usize = crate :: mem:: size_of :: < wasi:: Dirent > ( ) ;
213+ if contents. len ( ) >= DIRENT_SIZE {
214+ let ( dirent, data) = contents. split_at ( DIRENT_SIZE ) ;
215+ let dirent =
216+ unsafe { ptr:: read_unaligned ( dirent. as_ptr ( ) as * const wasi:: Dirent ) } ;
217+ // If the file name was truncated, then we need to reinvoke
218+ // `readdir` so we truncate our buffer to start over and reread this
219+ // descriptor.
220+ if data. len ( ) < dirent. d_namlen as usize {
221+ if buf. len ( ) < dirent. d_namlen as usize + DIRENT_SIZE {
222+ buf. resize ( dirent. d_namlen as usize + DIRENT_SIZE , 0 ) ;
223+ }
224+ if let Some ( next_read_offset) = * next_read_offset {
225+ self . state =
226+ ReadDirState :: FillBuffer { next_read_offset, buf : mem:: take ( buf) } ;
227+ } else {
228+ self . state = ReadDirState :: Done ;
229+ }
230+
231+ return self . next ( ) ;
232+ }
233+ next_read_offset. as_mut ( ) . map ( |cookie| {
234+ * cookie = dirent. d_next ;
235+ } ) ;
236+ * offset = * offset + DIRENT_SIZE + dirent. d_namlen as usize ;
223237
224- let name = & data[ ..( dirent. d_namlen as usize ) ] ;
238+ let name = & data[ ..( dirent. d_namlen as usize ) ] ;
239+
240+ // These names are skipped on all other platforms, so let's skip
241+ // them here too
242+ if name == b"." || name == b".." {
243+ return self . next ( ) ;
244+ }
225245
226- // These names are skipped on all other platforms, so let's skip
227- // them here too
228- if name == b"." || name == b".." {
229- continue ;
246+ return Some ( Ok ( DirEntry {
247+ meta : dirent,
248+ name : name. to_vec ( ) ,
249+ inner : self . inner . clone ( ) ,
250+ } ) ) ;
251+ } else if let Some ( next_read_offset) = * next_read_offset {
252+ self . state = ReadDirState :: FillBuffer { next_read_offset, buf : mem:: take ( buf) } ;
253+ } else {
254+ self . state = ReadDirState :: Done ;
255+ }
256+ self . next ( )
230257 }
258+ ReadDirState :: RunUntilExhaustion { buf, offset } => {
259+ if * offset >= buf. len ( ) {
260+ self . state = ReadDirState :: Done ;
261+ } else {
262+ self . state = ReadDirState :: ProcessEntry {
263+ buf : mem:: take ( buf) ,
264+ offset : * offset,
265+ next_read_offset : None ,
266+ } ;
267+ }
231268
232- return Some ( Ok ( DirEntry {
233- meta : dirent,
234- name : name. to_vec ( ) ,
235- inner : self . inner . clone ( ) ,
236- } ) ) ;
269+ self . next ( )
270+ }
271+ ReadDirState :: Done => None ,
237272 }
238273 }
239274}
0 commit comments