duvet_core/dir/
walk.rs

1// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4use super::{DirIter, Path};
5use crate::{glob::Glob, vfs};
6use core::future::Future;
7use futures::{stream, Stream};
8use std::{future::ready, marker::PhantomData};
9
10pub fn dir(root: Path) -> impl Stream<Item = Path> {
11    filtered(root, |depth, _path| ready((depth <= 100).into()))
12}
13
14pub fn glob(root: Path, include: Glob, ignore: Glob) -> impl Stream<Item = Path> {
15    filtered(root, move |depth, path| {
16        let path = path.clone();
17        let include = include.clone();
18        let ignore = ignore.clone();
19        async move {
20            let mut is_ok = !ignore.is_match(&path);
21
22            is_ok &= depth <= 100;
23
24            if is_ok {
25                match vfs::read_metadata(path.clone()).await.ok() {
26                    Some(meta) if meta.is_dir() => return ControlFlow::Skip,
27                    _ => {}
28                }
29            }
30
31            is_ok &= include.is_match(&path);
32
33            is_ok.into()
34        }
35    })
36}
37
38pub fn filtered<F, Fut>(root: Path, filter: F) -> impl Stream<Item = Path>
39where
40    F: FnMut(usize, &Path) -> Fut + Send,
41    Fut: Future<Output = ControlFlow> + Send,
42{
43    let stream = State {
44        start: Some(root),
45        stack_list: vec![],
46        depth: 0,
47        filter,
48        _f: PhantomData,
49    };
50
51    stream.walk_dir()
52}
53
54#[derive(Clone, Copy, Debug, PartialEq, Eq)]
55pub enum ControlFlow {
56    Yield,
57    Skip,
58    Break,
59}
60
61impl From<bool> for ControlFlow {
62    fn from(value: bool) -> Self {
63        if value {
64            Self::Yield
65        } else {
66            Self::Break
67        }
68    }
69}
70
71struct State<F, Fut> {
72    /// The start path.
73    ///
74    /// This is only `Some(...)` at the beginning. After the first iteration,
75    /// this is always `None`.
76    start: Option<Path>,
77    stack_list: Vec<DirIter>,
78    /// The current depth of iteration (the length of the stack at the
79    /// beginning of each iteration).
80    depth: usize,
81
82    filter: F,
83    _f: PhantomData<Fut>,
84}
85
86impl<F, Fut> State<F, Fut>
87where
88    F: FnMut(usize, &Path) -> Fut + Send,
89    Fut: Future<Output = ControlFlow> + Send,
90{
91    fn walk_dir(self) -> impl Stream<Item = Path> {
92        stream::unfold(self, move |mut state| async move {
93            if let Some(path) = state.start.take() {
94                if let Some(entry) = state.handle_entry(path).await {
95                    return Some((entry, state));
96                }
97            }
98
99            while !state.stack_list.is_empty() {
100                state.depth = state.stack_list.len();
101
102                let next = state
103                    .stack_list
104                    .last_mut()
105                    .expect("BUG: stack should be non-empty")
106                    .next();
107
108                match next {
109                    None => state.pop(),
110                    Some(path) => {
111                        if let Some(entry) = state.handle_entry(path).await {
112                            return Some((entry, state));
113                        }
114                    }
115                }
116            }
117
118            None
119        })
120    }
121
122    async fn handle_entry(&mut self, path: Path) -> Option<Path> {
123        let result = (self.filter)(self.depth, &path).await;
124
125        let should_yield = match result {
126            ControlFlow::Break => return None,
127            ControlFlow::Skip => false,
128            ControlFlow::Yield => true,
129        };
130
131        let meta = vfs::read_metadata(path.clone()).await.ok()?;
132
133        if meta.is_dir() {
134            if let Ok(dir) = vfs::read_dir(path.clone()).await {
135                self.stack_list.push(dir.into_iter());
136            }
137        }
138
139        if should_yield {
140            Some(path)
141        } else {
142            None
143        }
144    }
145
146    fn pop(&mut self) {
147        self.stack_list
148            .pop()
149            .expect("BUG: cannot pop from empty stack");
150    }
151}