tenant.openagents/omega

1mod path_trie {
2    use std::{collections::BTreeSet, ops::ControlFlow};
3
4    use util::rel_path::rel_path;
5
6    use project::manifest_tree::path_trie::*;
7
8    #[test]
9    fn test_insert_and_lookup() {
10        let mut trie = RootPathTrie::<()>::new();
11        trie.insert(
12            &TriePath::new(rel_path("a/b/c")),
13            (),
14            LabelPresence::Present,
15        );
16
17        trie.walk(&TriePath::new(rel_path("a/b/c")), &mut |path, nodes| {
18            assert_eq!(nodes.get(&()), Some(&LabelPresence::Present));
19            assert_eq!(path.as_unix_str(), "a/b/c");
20            ControlFlow::Continue(())
21        });
22        // Now let's annotate a parent with "Known missing" node.
23        trie.insert(
24            &TriePath::new(rel_path("a")),
25            (),
26            LabelPresence::KnownAbsent,
27        );
28
29        // Ensure that we walk from the root to the leaf.
30        let mut visited_paths = BTreeSet::new();
31        trie.walk(&TriePath::new(rel_path("a/b/c")), &mut |path, nodes| {
32            if path.as_unix_str() == "a/b/c" {
33                assert_eq!(visited_paths, BTreeSet::from_iter([rel_path("a").into()]));
34                assert_eq!(nodes.get(&()), Some(&LabelPresence::Present));
35            } else if path.as_unix_str() == "a" {
36                assert!(visited_paths.is_empty());
37                assert_eq!(nodes.get(&()), Some(&LabelPresence::KnownAbsent));
38            } else {
39                panic!("Unknown path");
40            }
41            // Assert that we only ever visit a path once.
42            assert!(visited_paths.insert(path.clone()));
43            ControlFlow::Continue(())
44        });
45
46        // One can also pass a path whose prefix is in the tree, but not that path itself.
47        let mut visited_paths = BTreeSet::new();
48        trie.walk(
49            &TriePath::new(rel_path("a/b/c/d/e/f/g")),
50            &mut |path, nodes| {
51                if path.as_unix_str() == "a/b/c" {
52                    assert_eq!(visited_paths, BTreeSet::from_iter([rel_path("a").into()]));
53                    assert_eq!(nodes.get(&()), Some(&LabelPresence::Present));
54                } else if path.as_unix_str() == "a" {
55                    assert!(visited_paths.is_empty());
56                    assert_eq!(nodes.get(&()), Some(&LabelPresence::KnownAbsent));
57                } else {
58                    panic!("Unknown path");
59                }
60                // Assert that we only ever visit a path once.
61                assert!(visited_paths.insert(path.clone()));
62                ControlFlow::Continue(())
63            },
64        );
65
66        // Test breaking from the tree-walk.
67        let mut visited_paths = BTreeSet::new();
68        trie.walk(&TriePath::new(rel_path("a/b/c")), &mut |path, nodes| {
69            if path.as_unix_str() == "a" {
70                assert!(visited_paths.is_empty());
71                assert_eq!(nodes.get(&()), Some(&LabelPresence::KnownAbsent));
72            } else {
73                panic!("Unknown path");
74            }
75            // Assert that we only ever visit a path once.
76            assert!(visited_paths.insert(path.clone()));
77            ControlFlow::Break(())
78        });
79        assert_eq!(visited_paths.len(), 1);
80
81        // Entry removal.
82        trie.insert(
83            &TriePath::new(rel_path("a/b")),
84            (),
85            LabelPresence::KnownAbsent,
86        );
87        let mut visited_paths = BTreeSet::new();
88        trie.walk(&TriePath::new(rel_path("a/b/c")), &mut |path, _nodes| {
89            // Assert that we only ever visit a path once.
90            assert!(visited_paths.insert(path.clone()));
91            ControlFlow::Continue(())
92        });
93        assert_eq!(visited_paths.len(), 3);
94        trie.remove(&TriePath::new(rel_path("a/b")));
95        let mut visited_paths = BTreeSet::new();
96        trie.walk(&TriePath::new(rel_path("a/b/c")), &mut |path, _nodes| {
97            // Assert that we only ever visit a path once.
98            assert!(visited_paths.insert(path.clone()));
99            ControlFlow::Continue(())
100        });
101        assert_eq!(visited_paths.len(), 1);
102        assert_eq!(
103            visited_paths.into_iter().next().unwrap(),
104            rel_path("a").into()
105        );
106    }
107
108    #[test]
109    fn path_to_a_root_can_contain_multiple_known_nodes() {
110        let mut trie = RootPathTrie::<()>::new();
111        trie.insert(&TriePath::new(rel_path("a/b")), (), LabelPresence::Present);
112        trie.insert(&TriePath::new(rel_path("a")), (), LabelPresence::Present);
113        let mut visited_paths = BTreeSet::new();
114        trie.walk(&TriePath::new(rel_path("a/b/c")), &mut |path, nodes| {
115            assert_eq!(nodes.get(&()), Some(&LabelPresence::Present));
116            if path.as_unix_str() != "a" && path.as_unix_str() != "a/b" {
117                panic!("Unexpected path: {}", path.as_unix_str());
118            }
119            assert!(visited_paths.insert(path.clone()));
120            ControlFlow::Continue(())
121        });
122        assert_eq!(visited_paths.len(), 2);
123    }
124}
125