Skip to main content

weavatrix_scan/
stateful_walk.rs

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/// A walk entry carrying caller-owned state assigned by `process_read_dir`.
20#[derive(Debug)]
21pub struct StatefulWalkEntry<E> {
22    entry: WalkEntry,
23    /// State retained with this entry after the directory callback returns.
24    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    /// Controls whether traversal descends into this directory.
60    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
70/// Builder for an iterative DFS walk with per-directory and per-entry state.
71///
72/// `R` is cloned from a processed directory into each accepted child
73/// directory. `E` is attached independently to every yielded entry.
74pub 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    /// Sets parallel directory workers. Zero selects runtime parallelism.
107    #[must_use]
108    pub const fn with_parallelism(mut self, parallelism: usize) -> Self {
109        self.parallelism = parallelism;
110        self
111    }
112
113    /// Selects the executor used by parallel stateful traversal.
114    #[must_use]
115    pub fn runtime(mut self, runtime: ParallelRuntime) -> Self {
116        self.runtime = runtime;
117        self
118    }
119
120    /// Processes the complete immediate child batch before entries are yielded.
121    ///
122    /// The callback may sort or retain the vector, remove local errors, mutate
123    /// the inherited read-directory state, attach state to entries, and call
124    /// [`StatefulWalkEntry::set_read_children`] for per-entry pruning.
125    #[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    /// Builds the stateful iterator after validating the root.
138    ///
139    /// # Errors
140    ///
141    /// Returns an error when the root cannot be resolved or inspected.
142    pub fn build(self) -> Result<StatefulWalker<R, E>, WalkError> {
143        StatefulWalker::new(self)
144    }
145
146    /// Builds a bounded parallel iterator with strict depth-first output.
147    ///
148    /// Directory callbacks run on the selected executor over complete child
149    /// batches. Their mutated directory state is cloned into accepted child
150    /// tasks, while yielded entries retain callback-assigned entry state.
151    ///
152    /// # Errors
153    ///
154    /// Returns a root-validation, coordinator-startup, or worker-submission
155    /// error.
156    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
178/// Strict depth-first iterator produced by [`StatefulWalkBuilder`].
179pub 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}