1use crate::ParallelRuntime;
2use crate::walk_platform::{DirectoryIdentity, FileSystemId};
3use crate::walker::{WalkEntry, WalkError, WalkOptions, Walker};
4use std::collections::HashSet;
5use std::path::{Path, PathBuf};
6use std::sync::Arc;
7
8mod parallel;
9
10pub use parallel::ParallelStatefulWalker;
11
12type DirectoryProcessor<R, E> = Arc<
13 dyn Fn(usize, &Path, &mut R, &mut Vec<Result<StatefulWalkEntry<E>, WalkError>>)
14 + Send
15 + Sync
16 + 'static,
17>;
18
19#[derive(Debug)]
21pub struct StatefulWalkEntry<E> {
22 entry: WalkEntry,
23 pub state: E,
25 read_children: bool,
26}
27
28impl<E> StatefulWalkEntry<E> {
29 #[must_use]
30 pub const fn entry(&self) -> &WalkEntry {
31 &self.entry
32 }
33
34 #[must_use]
35 pub fn path(&self) -> &Path {
36 self.entry.path()
37 }
38
39 #[must_use]
40 pub const fn depth(&self) -> usize {
41 self.entry.depth()
42 }
43
44 #[must_use]
45 pub const fn is_file(&self) -> bool {
46 self.entry.is_file()
47 }
48
49 #[must_use]
50 pub const fn is_dir(&self) -> bool {
51 self.entry.is_dir()
52 }
53
54 #[must_use]
55 pub const fn read_children(&self) -> bool {
56 self.read_children
57 }
58
59 pub const fn set_read_children(&mut self, enabled: bool) {
61 self.read_children = enabled && self.entry.is_dir() && self.entry.skip_reason().is_none();
62 }
63
64 #[must_use]
65 pub fn into_parts(self) -> (WalkEntry, E) {
66 (self.entry, self.state)
67 }
68}
69
70pub struct StatefulWalkBuilder<R, E> {
75 root: PathBuf,
76 options: WalkOptions,
77 root_read_dir_state: R,
78 processor: Option<DirectoryProcessor<R, E>>,
79 parallelism: usize,
80 runtime: ParallelRuntime,
81}
82
83impl<R, E> StatefulWalkBuilder<R, E>
84where
85 R: Clone + Send + 'static,
86 E: Default + Send + 'static,
87{
88 #[must_use]
89 pub fn new(root: impl Into<PathBuf>, root_read_dir_state: R) -> Self {
90 Self {
91 root: root.into(),
92 options: WalkOptions::default(),
93 root_read_dir_state,
94 processor: None,
95 parallelism: 0,
96 runtime: ParallelRuntime::global(),
97 }
98 }
99
100 #[must_use]
101 pub const fn options(mut self, options: WalkOptions) -> Self {
102 self.options = options;
103 self
104 }
105
106 #[must_use]
108 pub const fn with_parallelism(mut self, parallelism: usize) -> Self {
109 self.parallelism = parallelism;
110 self
111 }
112
113 #[must_use]
115 pub fn runtime(mut self, runtime: ParallelRuntime) -> Self {
116 self.runtime = runtime;
117 self
118 }
119
120 #[must_use]
126 pub fn process_read_dir<F>(mut self, processor: F) -> Self
127 where
128 F: Fn(usize, &Path, &mut R, &mut Vec<Result<StatefulWalkEntry<E>, WalkError>>)
129 + Send
130 + Sync
131 + 'static,
132 {
133 self.processor = Some(Arc::new(processor));
134 self
135 }
136
137 pub fn build(self) -> Result<StatefulWalker<R, E>, WalkError> {
143 StatefulWalker::new(self)
144 }
145
146 pub fn build_parallel_ordered(
157 self,
158 capacity: usize,
159 ) -> Result<ParallelStatefulWalker<E>, WalkError> {
160 ParallelStatefulWalker::start(self, capacity)
161 }
162}
163
164struct DirectoryTask<R> {
165 path: PathBuf,
166 depth: usize,
167 identity: Option<DirectoryIdentity>,
168 ancestors: HashSet<DirectoryIdentity>,
169 read_state: R,
170}
171
172struct DirectoryFrame<R, E> {
173 entries: std::vec::IntoIter<Result<StatefulWalkEntry<E>, WalkError>>,
174 child_state: R,
175 ancestors: HashSet<DirectoryIdentity>,
176}
177
178pub struct StatefulWalker<R, E> {
180 root: Arc<PathBuf>,
181 root_file_system: Option<FileSystemId>,
182 options: WalkOptions,
183 processor: Option<DirectoryProcessor<R, E>>,
184 root_entry: Option<StatefulWalkEntry<E>>,
185 pending: Option<DirectoryTask<R>>,
186 frames: Vec<DirectoryFrame<R, E>>,
187}
188
189impl<R, E> StatefulWalker<R, E>
190where
191 R: Clone + Send + 'static,
192 E: Default + Send + 'static,
193{
194 fn new(builder: StatefulWalkBuilder<R, E>) -> Result<Self, WalkError> {
195 let options = builder.options.normalized();
196 let mut root_options = options;
197 root_options.min_depth = 0;
198 let mut walker = Walker::with_options(&builder.root, root_options)?;
199 let root_file_system = walker.root_file_system;
200 let root = Arc::clone(&walker.root);
201 let root_entry = walker.next().expect("a validated root yields one entry")?;
202 let identity = root_entry.directory_identity();
203 let mut ancestors = HashSet::new();
204 if let Some(identity) = identity {
205 ancestors.insert(identity);
206 }
207 let can_descend = root_entry.is_dir() && root_entry.skip_reason().is_none();
208 let root_entry = (root_entry.depth() >= options.min_depth).then(|| StatefulWalkEntry {
209 read_children: can_descend,
210 entry: root_entry,
211 state: E::default(),
212 });
213 let pending = can_descend.then(|| DirectoryTask {
214 path: root.as_ref().clone(),
215 depth: 0,
216 identity,
217 ancestors,
218 read_state: builder.root_read_dir_state,
219 });
220 Ok(Self {
221 root,
222 root_file_system,
223 options,
224 processor: builder.processor,
225 root_entry,
226 pending,
227 frames: Vec::new(),
228 })
229 }
230
231 fn read_directory(&self, mut task: DirectoryTask<R>) -> DirectoryFrame<R, E> {
232 let mut worker_options = self.options;
233 worker_options.error_policy = crate::ErrorPolicy::Continue;
234 worker_options.min_depth = 0;
235 worker_options.max_open = 1;
236 worker_options.max_depth = Some(
237 self.options
238 .max_depth
239 .unwrap_or(task.depth.saturating_add(1))
240 .min(task.depth.saturating_add(1)),
241 );
242 let mut walker = Walker::from_known_directory_with_ancestry(
243 &self.root,
244 task.path.clone(),
245 task.depth,
246 worker_options,
247 self.root_file_system,
248 task.identity,
249 task.ancestors.clone(),
250 );
251 let mut entries = Vec::new();
252 while let Some(item) = walker.next() {
253 match item {
254 Ok(mut entry) => {
255 if entry.is_dir()
256 && entry.skip_reason() == Some(crate::WalkSkipReason::MaxDepth)
257 && self
258 .options
259 .max_depth
260 .is_none_or(|maximum| entry.depth() < maximum)
261 {
262 entry.clear_depth_skip();
263 }
264 if entry.is_dir() {
265 walker.skip_current_dir();
266 }
267 entries.push(Ok(StatefulWalkEntry {
268 read_children: entry.is_dir() && entry.skip_reason().is_none(),
269 entry,
270 state: E::default(),
271 }));
272 }
273 Err(error) => entries.push(Err(error)),
274 }
275 }
276 if let Some(processor) = self.processor.as_ref() {
277 processor(task.depth, &task.path, &mut task.read_state, &mut entries);
278 }
279 DirectoryFrame {
280 entries: entries.into_iter(),
281 child_state: task.read_state,
282 ancestors: task.ancestors,
283 }
284 }
285}
286
287impl<R, E> Iterator for StatefulWalker<R, E>
288where
289 R: Clone + Send + 'static,
290 E: Default + Send + 'static,
291{
292 type Item = Result<StatefulWalkEntry<E>, WalkError>;
293
294 fn next(&mut self) -> Option<Self::Item> {
295 if let Some(root_entry) = self.root_entry.take() {
296 return Some(Ok(root_entry));
297 }
298 loop {
299 if let Some(task) = self.pending.take() {
300 let frame = self.read_directory(task);
301 self.frames.push(frame);
302 }
303 let frame = self.frames.last_mut()?;
304 let Some(item) = frame.entries.next() else {
305 self.frames.pop();
306 continue;
307 };
308 if let Ok(entry) = &item
309 && entry.read_children
310 {
311 let identity = entry.entry.directory_identity();
312 let mut ancestors = frame.ancestors.clone();
313 if let Some(identity) = identity {
314 ancestors.insert(identity);
315 }
316 self.pending = Some(DirectoryTask {
317 path: entry.path().to_path_buf(),
318 depth: entry.depth(),
319 identity,
320 ancestors,
321 read_state: frame.child_state.clone(),
322 });
323 }
324 let visible = item
325 .as_ref()
326 .map_or(true, |entry| entry.depth() >= self.options.min_depth);
327 if visible {
328 return Some(item);
329 }
330 }
331 }
332}