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