tenant.openagents/omega

1use std::{
2    collections::{BTreeMap, btree_map::Entry},
3    ops::ControlFlow,
4    sync::Arc,
5};
6
7use path::rel_path::RelPathBuf;
8use util::rel_path::RelPath;
9
10/// [RootPathTrie] is a workhorse of [super::ManifestTree]. It is responsible for determining the closest known entry for a given path.
11/// It also determines how much of a given path is unexplored, thus letting callers fill in that gap if needed.
12/// Conceptually, it allows one to annotate Worktree entries with arbitrary extra metadata and run closest-ancestor searches.
13///
14/// 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.
15/// 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
16/// a known root at `python/project` and the unexplored part is `subdir/another_subdir` - we need to run a scan on these 2 directories.
17pub struct RootPathTrie<Label> {
18    worktree_relative_path: Arc<RelPath>,
19    labels: BTreeMap<Label, LabelPresence>,
20    children: BTreeMap<Arc<str>, RootPathTrie<Label>>,
21}
22
23/// Label presence is a marker that allows to optimize searches within [RootPathTrie]; node label can be:
24/// - Present; we know there's definitely a project root at this node.
25/// - 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!).
26///   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
27///   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
28///   from the leaf up to the root of the worktree.
29///
30/// 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
31/// (unless the node is invalidated, which can happen when FS entries are renamed/removed).
32///
33/// 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
34/// such scan more than once.
35#[derive(Clone, Copy, Debug, PartialOrd, PartialEq, Ord, Eq)]
36pub enum LabelPresence {
37    KnownAbsent,
38    Present,
39}
40
41impl<Label: Ord + Clone> RootPathTrie<Label> {
42    pub fn new() -> Self {
43        Self::new_with_key(Arc::from(RelPath::empty()))
44    }
45
46    fn new_with_key(worktree_relative_path: Arc<RelPath>) -> Self {
47        RootPathTrie {
48            worktree_relative_path,
49            labels: Default::default(),
50            children: Default::default(),
51        }
52    }
53
54    // Internal implementation of inner that allows one to visit descendants of insertion point for a node.
55    fn insert_inner(
56        &mut self,
57        path: &TriePath,
58        value: Label,
59        presence: LabelPresence,
60    ) -> &mut Self {
61        let mut current = self;
62
63        let mut path_so_far = RelPathBuf::new();
64        for key in path.0.iter() {
65            path_so_far = path_so_far.join(RelPath::from_unix_str(key.as_ref()).unwrap());
66            current = match current.children.entry(key.clone()) {
67                Entry::Vacant(vacant_entry) => {
68                    vacant_entry.insert(RootPathTrie::new_with_key(path_so_far.clone().into()))
69                }
70                Entry::Occupied(occupied_entry) => occupied_entry.into_mut(),
71            };
72        }
73        let _previous_value = current.labels.insert(value, presence);
74        debug_assert_eq!(_previous_value, None);
75        current
76    }
77
78    pub fn insert(&mut self, path: &TriePath, value: Label, presence: LabelPresence) {
79        self.insert_inner(path, value, presence);
80    }
81
82    pub fn walk<'a>(
83        &'a self,
84        path: &TriePath,
85        callback: &mut dyn for<'b> FnMut(
86            &'b Arc<RelPath>,
87            &'a BTreeMap<Label, LabelPresence>,
88        ) -> ControlFlow<()>,
89    ) {
90        let mut current = self;
91        for key in path.0.iter() {
92            if !current.labels.is_empty()
93                && (callback)(&current.worktree_relative_path, &current.labels).is_break()
94            {
95                return;
96            };
97            current = match current.children.get(key) {
98                Some(child) => child,
99                None => return,
100            };
101        }
102        if !current.labels.is_empty() {
103            let _ = (callback)(&current.worktree_relative_path, &current.labels);
104        }
105    }
106
107    pub fn remove(&mut self, path: &TriePath) {
108        let mut current = self;
109        for path in path.0.iter().take(path.0.len().saturating_sub(1)) {
110            current = match current.children.get_mut(path) {
111                Some(child) => child,
112                None => return,
113            };
114        }
115        if let Some(final_entry_name) = path.0.last() {
116            current.children.remove(final_entry_name);
117        }
118    }
119}
120
121/// [TriePath] is a [Path] preprocessed for amortizing the cost of doing multiple lookups in distinct [RootPathTrie]s.
122#[derive(Clone)]
123pub struct TriePath(Arc<[Arc<str>]>);
124
125impl TriePath {
126    pub fn new(value: &RelPath) -> Self {
127        TriePath(
128            value
129                .components()
130                .map(|component| component.into())
131                .collect(),
132        )
133    }
134}
135
136impl From<&RelPath> for TriePath {
137    fn from(value: &RelPath) -> Self {
138        Self::new(value)
139    }
140}
141