1 use std::iter::Peekable;
2 
3 use super::*;
4 use crate::Result;
5 
6 /// DirEntry iterator from `WalkDir.into_iter()`.
7 ///
8 /// Yeilds entries from recursive traversal of filesystem.
9 pub struct DirEntryIter<C: ClientState> {
10     min_depth: usize,
11     // iterator yeilding next ReadDir results when needed
12     read_dir_iter: Peekable<ReadDirIter<C>>,
13     // stack of ReadDir results, track location in filesystem traversal
14     read_dir_results_stack: Vec<vec::IntoIter<Result<DirEntry<C>>>>,
15 }
16 
17 impl<C: ClientState> DirEntryIter<C> {
new( root_entry_results: Vec<Result<DirEntry<C>>>, parallelism: Parallelism, min_depth: usize, root_read_dir_state: C::ReadDirState, core_read_dir_callback: Arc<ReadDirCallback<C>>, ) -> DirEntryIter<C>18     pub(crate) fn new(
19         root_entry_results: Vec<Result<DirEntry<C>>>,
20         parallelism: Parallelism,
21         min_depth: usize,
22         root_read_dir_state: C::ReadDirState,
23         core_read_dir_callback: Arc<ReadDirCallback<C>>,
24     ) -> DirEntryIter<C> {
25         // 1. Gather read_dir_specs from root level
26         let read_dir_specs: Vec<_> = root_entry_results
27             .iter()
28             .flat_map(|dir_entry_result| {
29                 dir_entry_result
30                     .as_ref()
31                     .ok()?
32                     .read_children_spec(root_read_dir_state.clone())
33             })
34             .collect();
35 
36         // 2. Init new read_dir_iter from those specs
37         let read_dir_iter = ReadDirIter::new(read_dir_specs, parallelism, core_read_dir_callback);
38 
39         // 3. Return DirEntryIter that will return initial root entries and then
40         //    fill and process read_dir_iter until complete
41         DirEntryIter {
42             min_depth,
43             read_dir_iter: read_dir_iter.peekable(),
44             read_dir_results_stack: vec![root_entry_results.into_iter()],
45         }
46     }
47 
push_next_read_dir_results(&mut self) -> Result<()>48     fn push_next_read_dir_results(&mut self) -> Result<()> {
49         // Push next read dir results or return error if read failed
50         let read_dir_result = self.read_dir_iter.next().unwrap();
51         let read_dir = match read_dir_result {
52             Ok(read_dir) => read_dir,
53             Err(err) => return Err(err),
54         };
55 
56         let ReadDir { results_list, .. } = read_dir;
57 
58         self.read_dir_results_stack.push(results_list.into_iter());
59 
60         Ok(())
61     }
62 }
63 
64 impl<C: ClientState> Iterator for DirEntryIter<C> {
65     type Item = Result<DirEntry<C>>;
next(&mut self) -> Option<Self::Item>66     fn next(&mut self) -> Option<Self::Item> {
67         loop {
68             if self.read_dir_results_stack.is_empty() {
69                 return None;
70             }
71 
72             // 1. Get current read dir results iter from top of stack
73             let top_read_dir_results = self.read_dir_results_stack.last_mut().unwrap();
74 
75             // 2. If more results in current read dir then process
76             if let Some(dir_entry_result) = top_read_dir_results.next() {
77                 // 2.1 Handle error case
78                 let mut dir_entry = match dir_entry_result {
79                     Ok(dir_entry) => dir_entry,
80                     Err(err) => return Some(Err(err)),
81                 };
82                 // 2.2 If dir_entry has a read_children_path means we need to read a new
83                 // directory and push those results onto read_dir_results_stack
84                 if dir_entry.read_children_path.is_some() {
85                     if let Err(err) = self.push_next_read_dir_results() {
86                         dir_entry.read_children_error = Some(err);
87                     }
88                 }
89 
90                 if dir_entry.depth >= self.min_depth {
91                     // 2.3 Finished, return dir_entry
92                     return Some(Ok(dir_entry));
93                 }
94             } else {
95                 // If no more results in current then pop stack
96                 self.read_dir_results_stack.pop();
97             }
98         }
99     }
100 }
101