1use std::{fs::Metadata, io, path::PathBuf};
4
5use tokio::fs::{self};
6
7pub struct DirectoryWalker<T>
8where
9 T: Fn(&str) -> bool,
10{
11 stack: Vec<(PathBuf, Option<Metadata>)>,
12 filter: T,
13 recursive: bool,
14 eat_root: bool,
15}
16
17impl<T> DirectoryWalker<T>
18where
19 T: Fn(&str) -> bool,
20{
21 pub fn new(root: PathBuf, filter: T) -> Self {
22 Self {
23 stack: vec![(root, None)],
24 filter,
25 recursive: false,
26 eat_root: true,
27 }
28 }
29
30 pub fn set_recursive(&mut self, recursive: bool) {
31 self.recursive = recursive;
32 }
33
34 pub async fn walk(&mut self) -> io::Result<Option<(PathBuf, Metadata)>> {
35 if self.eat_root {
36 self.eat_root = false;
37 let (dir, _) = self.stack.pop().unwrap();
38 self.append_stack(&dir).await?;
39 }
40 if let Some((entry, metadata)) = self.stack.pop() {
41 let metadata = metadata.unwrap();
42 if self.recursive && metadata.is_dir() {
43 self.append_stack(&entry).await?;
44 }
45
46 Ok(Some((entry, metadata)))
47 } else {
48 Ok(None)
49 }
50 }
51
52 pub fn walk_sync(&mut self) -> io::Result<Option<(PathBuf, Metadata)>> {
53 if self.eat_root {
54 self.eat_root = false;
55 let (dir, _) = self.stack.pop().unwrap();
56 self.append_stack_sync(&dir)?;
57 }
58 if let Some((entry, metadata)) = self.stack.pop() {
59 let metadata = metadata.unwrap();
60 if self.recursive && metadata.is_dir() {
61 self.append_stack_sync(&entry)?;
62 }
63
64 Ok(Some((entry, metadata)))
65 } else {
66 Ok(None)
67 }
68 }
69
70 async fn append_stack(&mut self, dir: &PathBuf) -> io::Result<()> {
71 let mut stream = fs::read_dir(dir).await?;
72
73 while let Some(entry) = stream.next_entry().await? {
74 let name = entry.file_name();
75 let name = name.to_string_lossy();
76 if !(self.filter)(name.as_ref()) {
77 continue;
78 }
79 let entry_path = entry.path();
80 let metadata = fs::symlink_metadata(&entry_path).await?;
81
82 self.stack.push((entry_path, Some(metadata)));
83 }
84 Ok(())
85 }
86
87 fn append_stack_sync(&mut self, dir: &PathBuf) -> io::Result<()> {
88 let dir = std::fs::read_dir(dir)?;
89
90 for entry in dir.flatten() {
91 let name = entry.file_name();
92 let name = name.to_string_lossy();
93 if !(self.filter)(name.as_ref()) {
94 continue;
95 }
96 let entry_path = entry.path();
97 let metadata = entry_path.symlink_metadata()?;
98 self.stack.push((entry_path, Some(metadata)))
99 }
100
101 Ok(())
102 }
103}