1use std::{
2 collections::{BTreeMap, btree_map::Entry},
3 ffi::OsStr,
4 ops::ControlFlow,
5 path::{Path, PathBuf},
6 sync::Arc,
7};
8
9/// [RootPathTrie] is a workhorse of [super::ManifestTree]. It is responsible for determining the closest known entry for a given path.
10/// It also determines how much of a given path is unexplored, thus letting callers fill in that gap if needed.
11/// Conceptually, it allows one to annotate Worktree entries with arbitrary extra metadata and run closest-ancestor searches.
12///
13/// A path is unexplored when the closest ancestor of a path is not the path itself; that means that we have not yet ran the scan on that path.
14/// For example, if there's a project root at path `python/project` and we query for a path `python/project/subdir/another_subdir/file.py`, there is
15/// a known root at `python/project` and the unexplored part is `subdir/another_subdir` - we need to run a scan on these 2 directories.
16pub(super) struct RootPathTrie<Label> {
17 worktree_relative_path: Arc<Path>,
18 labels: BTreeMap<Label, LabelPresence>,
19 children: BTreeMap<Arc<OsStr>, RootPathTrie<Label>>,
20}
21
22/// Label presence is a marker that allows to optimize searches within [RootPathTrie]; node label can be:
23/// - Present; we know there's definitely a project root at this node.
24/// - Known Absent - we know there's definitely no project root at this node and none of it's ancestors are Present (descendants can be present though!).
25/// The distinction is there to optimize searching; when we encounter a node with unknown status, we don't need to look at it's full path
26/// to the root of the worktree; it's sufficient to explore only the path between last node with a KnownAbsent state and the directory of a path, since we run searches
27/// from the leaf up to the root of the worktree.
28///
29/// In practical terms, it means that by storing label presence we don't need to do a project discovery on a given folder more than once
30/// (unless the node is invalidated, which can happen when FS entries are renamed/removed).
31///
32/// Storing absent nodes allows us to recognize which paths have already been scanned for a project root unsuccessfully. This way we don't need to run
33/// such scan more than once.
34#[derive(Clone, Copy, Debug, PartialOrd, PartialEq, Ord, Eq)]
35pub(super) enum LabelPresence {
36 KnownAbsent,
37 Present,
38}
39
40impl<Label: Ord + Clone> RootPathTrie<Label> {
41 pub(super) fn new() -> Self {
42 Self::new_with_key(Arc::from(Path::new("")))
43 }
44 fn new_with_key(worktree_relative_path: Arc<Path>) -> Self {
45 RootPathTrie {
46 worktree_relative_path,
47 labels: Default::default(),
48 children: Default::default(),
49 }
50 }
51 // Internal implementation of inner that allows one to visit descendants of insertion point for a node.
52 fn insert_inner(
53 &mut self,
54 path: &TriePath,
55 value: Label,
56 presence: LabelPresence,
57 ) -> &mut Self {
58 let mut current = self;
59
60 let mut path_so_far = PathBuf::new();
61 for key in path.0.iter() {
62 path_so_far.push(Path::new(key));
63 current = match current.children.entry(key.clone()) {
64 Entry::Vacant(vacant_entry) => vacant_entry
65 .insert(RootPathTrie::new_with_key(Arc::from(path_so_far.as_path()))),
66 Entry::Occupied(occupied_entry) => occupied_entry.into_mut(),
67 };
68 }
69 let _previous_value = current.labels.insert(value, presence);
70 debug_assert_eq!(_previous_value, None);
71 current
72 }
73 pub(super) fn insert(&mut self, path: &TriePath, value: Label, presence: LabelPresence) {
74 self.insert_inner(path, value, presence);
75 }
76
77 pub(super) fn walk<'a>(
78 &'a self,
79 path: &TriePath,
80 callback: &mut dyn for<'b> FnMut(
81 &'b Arc<Path>,
82 &'a BTreeMap<Label, LabelPresence>,
83 ) -> ControlFlow<()>,
84 ) {
85 let mut current = self;
86 for key in path.0.iter() {
87 if !current.labels.is_empty()
88 && (callback)(¤t.worktree_relative_path, ¤t.labels).is_break() {
89 return;
90 };
91 current = match current.children.get(key) {
92 Some(child) => child,
93 None => return,
94 };
95 }
96 if !current.labels.is_empty() {
97 let _ = (callback)(¤t.worktree_relative_path, ¤t.labels);
98 }
99 }
100
101 pub(super) fn remove(&mut self, path: &TriePath) {
102 let mut current = self;
103 for path in path.0.iter().take(path.0.len().saturating_sub(1)) {
104 current = match current.children.get_mut(path) {
105 Some(child) => child,
106 None => return,
107 };
108 }
109 if let Some(final_entry_name) = path.0.last() {
110 current.children.remove(final_entry_name);
111 }
112 }
113}
114
115/// [TriePath] is a [Path] preprocessed for amortizing the cost of doing multiple lookups in distinct [RootPathTrie]s.
116#[derive(Clone)]
117pub(super) struct TriePath(Arc<[Arc<OsStr>]>);
118
119impl From<&Path> for TriePath {
120 fn from(value: &Path) -> Self {
121 TriePath(value.components().map(|c| c.as_os_str().into()).collect())
122 }
123}
124
125#[cfg(test)]
126mod tests {
127 use std::collections::BTreeSet;
128
129 use super::*;
130
131 #[test]
132 fn test_insert_and_lookup() {
133 let mut trie = RootPathTrie::<()>::new();
134 trie.insert(
135 &TriePath::from(Path::new("a/b/c")),
136 (),
137 LabelPresence::Present,
138 );
139
140 trie.walk(&TriePath::from(Path::new("a/b/c")), &mut |path, nodes| {
141 assert_eq!(nodes.get(&()), Some(&LabelPresence::Present));
142 assert_eq!(path.as_ref(), Path::new("a/b/c"));
143 ControlFlow::Continue(())
144 });
145 // Now let's annotate a parent with "Known missing" node.
146 trie.insert(
147 &TriePath::from(Path::new("a")),
148 (),
149 LabelPresence::KnownAbsent,
150 );
151
152 // Ensure that we walk from the root to the leaf.
153 let mut visited_paths = BTreeSet::new();
154 trie.walk(&TriePath::from(Path::new("a/b/c")), &mut |path, nodes| {
155 if path.as_ref() == Path::new("a/b/c") {
156 assert_eq!(
157 visited_paths,
158 BTreeSet::from_iter([Arc::from(Path::new("a/"))])
159 );
160 assert_eq!(nodes.get(&()), Some(&LabelPresence::Present));
161 } else if path.as_ref() == Path::new("a/") {
162 assert!(visited_paths.is_empty());
163 assert_eq!(nodes.get(&()), Some(&LabelPresence::KnownAbsent));
164 } else {
165 panic!("Unknown path");
166 }
167 // Assert that we only ever visit a path once.
168 assert!(visited_paths.insert(path.clone()));
169 ControlFlow::Continue(())
170 });
171
172 // One can also pass a path whose prefix is in the tree, but not that path itself.
173 let mut visited_paths = BTreeSet::new();
174 trie.walk(
175 &TriePath::from(Path::new("a/b/c/d/e/f/g")),
176 &mut |path, nodes| {
177 if path.as_ref() == Path::new("a/b/c") {
178 assert_eq!(
179 visited_paths,
180 BTreeSet::from_iter([Arc::from(Path::new("a/"))])
181 );
182 assert_eq!(nodes.get(&()), Some(&LabelPresence::Present));
183 } else if path.as_ref() == Path::new("a/") {
184 assert!(visited_paths.is_empty());
185 assert_eq!(nodes.get(&()), Some(&LabelPresence::KnownAbsent));
186 } else {
187 panic!("Unknown path");
188 }
189 // Assert that we only ever visit a path once.
190 assert!(visited_paths.insert(path.clone()));
191 ControlFlow::Continue(())
192 },
193 );
194
195 // Test breaking from the tree-walk.
196 let mut visited_paths = BTreeSet::new();
197 trie.walk(&TriePath::from(Path::new("a/b/c")), &mut |path, nodes| {
198 if path.as_ref() == Path::new("a/") {
199 assert!(visited_paths.is_empty());
200 assert_eq!(nodes.get(&()), Some(&LabelPresence::KnownAbsent));
201 } else {
202 panic!("Unknown path");
203 }
204 // Assert that we only ever visit a path once.
205 assert!(visited_paths.insert(path.clone()));
206 ControlFlow::Break(())
207 });
208 assert_eq!(visited_paths.len(), 1);
209
210 // Entry removal.
211 trie.insert(
212 &TriePath::from(Path::new("a/b")),
213 (),
214 LabelPresence::KnownAbsent,
215 );
216 let mut visited_paths = BTreeSet::new();
217 trie.walk(&TriePath::from(Path::new("a/b/c")), &mut |path, _nodes| {
218 // Assert that we only ever visit a path once.
219 assert!(visited_paths.insert(path.clone()));
220 ControlFlow::Continue(())
221 });
222 assert_eq!(visited_paths.len(), 3);
223 trie.remove(&TriePath::from(Path::new("a/b/")));
224 let mut visited_paths = BTreeSet::new();
225 trie.walk(&TriePath::from(Path::new("a/b/c")), &mut |path, _nodes| {
226 // Assert that we only ever visit a path once.
227 assert!(visited_paths.insert(path.clone()));
228 ControlFlow::Continue(())
229 });
230 assert_eq!(visited_paths.len(), 1);
231 assert_eq!(
232 visited_paths.into_iter().next().unwrap().as_ref(),
233 Path::new("a/")
234 );
235 }
236
237 #[test]
238 fn path_to_a_root_can_contain_multiple_known_nodes() {
239 let mut trie = RootPathTrie::<()>::new();
240 trie.insert(
241 &TriePath::from(Path::new("a/b")),
242 (),
243 LabelPresence::Present,
244 );
245 trie.insert(&TriePath::from(Path::new("a")), (), LabelPresence::Present);
246 let mut visited_paths = BTreeSet::new();
247 trie.walk(&TriePath::from(Path::new("a/b/c")), &mut |path, nodes| {
248 assert_eq!(nodes.get(&()), Some(&LabelPresence::Present));
249 if path.as_ref() != Path::new("a") && path.as_ref() != Path::new("a/b") {
250 panic!("Unexpected path: {}", path.as_ref().display());
251 }
252 assert!(visited_paths.insert(path.clone()));
253 ControlFlow::Continue(())
254 });
255 assert_eq!(visited_paths.len(), 2);
256 }
257}