catlog/dbl/
tree_algorithms.rs

1//! Generic algorithms on [trees](Tree).
2
3use ego_tree::iter::{Descendants, Edge};
4use ego_tree::{NodeId, NodeRef, Tree};
5use std::collections::VecDeque;
6
7/// Extension trait adding traversal algorithms on [trees](Tree).
8pub trait TreeTraversal<T> {
9    /// Iterates over nodes of a tree in depth-first order.
10    fn dfs<'a>(&'a self) -> Descendants<'a, T>
11    where
12        T: 'a;
13
14    /// Iterates over the nodes in a tree in breadth-first order.
15    fn bfs<'a>(&'a self) -> BreadthFirstTraversal<'a, T>
16    where
17        T: 'a;
18}
19
20/// Iterator for traversing a tree in breadth-first order.
21pub struct BreadthFirstTraversal<'a, T: 'a> {
22    tree: &'a Tree<T>,
23    queue: VecDeque<(NodeId, usize)>,
24    current_level: usize,
25}
26
27impl<'a, T: 'a> BreadthFirstTraversal<'a, T> {
28    /// Initialize a breadth-first traversal at the given node.
29    pub fn starting_at(root: NodeRef<'a, T>) -> Self {
30        let tree = root.tree();
31        let mut queue = VecDeque::new();
32        queue.push_back((root.id(), 1));
33        Self {
34            tree,
35            queue,
36            current_level: 0,
37        }
38    }
39
40    /// Peeks at the next node, if it's at the same level as the previous one.
41    pub fn peek_at_same_level(&self) -> Option<NodeRef<'a, T>> {
42        self.queue.front().and_then(|(id, level)| {
43            if *level == self.current_level {
44                self.tree.get(*id)
45            } else {
46                None
47            }
48        })
49    }
50}
51
52impl<'a, T: 'a> Iterator for BreadthFirstTraversal<'a, T> {
53    type Item = NodeRef<'a, T>;
54
55    fn next(&mut self) -> Option<Self::Item> {
56        let (id, level) = self.queue.pop_front()?;
57        self.current_level = level;
58        let node = self.tree.get(id).unwrap();
59        for child in node.children() {
60            self.queue.push_back((child.id(), level + 1));
61        }
62        Some(node)
63    }
64}
65
66impl<'a, T: 'a> std::iter::FusedIterator for BreadthFirstTraversal<'a, T> {}
67
68impl<T> TreeTraversal<T> for Tree<T> {
69    /// Uses the built-in traversal algorithm, which is depth-first, though that
70    /// is not documented: <https://github.com/rust-scraper/ego-tree/issues/38>
71    fn dfs<'a>(&'a self) -> Descendants<'a, T>
72    where
73        T: 'a,
74    {
75        self.root().descendants()
76    }
77
78    /// Implements the standard BFS algorithm using a queue.
79    fn bfs<'a>(&'a self) -> BreadthFirstTraversal<'a, T>
80    where
81        T: 'a,
82    {
83        BreadthFirstTraversal::starting_at(self.root())
84    }
85}
86
87/// Extension trait adding isomorphism checking on [trees](Tree).
88pub trait TreeIsomorphism<T> {
89    /** Is the tree isomorphic to another?
90
91    In the standard data structure for trees based on pointers, there is only
92    one notion of sameness that makes sense, but for vector-backed trees with
93    node IDs, trees can be isomorphic (logically the same) without having
94    underlying data that is equal.
95     */
96    fn is_isomorphic_to(&self, other: &Self) -> bool;
97}
98
99impl<T> TreeIsomorphism<T> for Tree<T>
100where
101    T: Eq,
102{
103    fn is_isomorphic_to(&self, other: &Self) -> bool {
104        let mut self_traversal = self.root().traverse();
105        let mut other_traversal = other.root().traverse();
106        loop {
107            match (self_traversal.next(), other_traversal.next()) {
108                (Some(Edge::Open(n1)), Some(Edge::Open(n2))) if n1.value() == n2.value() => {}
109                (Some(Edge::Close(n1)), Some(Edge::Close(n2))) if n1.value() == n2.value() => {}
110                (None, None) => {
111                    break;
112                }
113                _ => {
114                    return false;
115                }
116            }
117        }
118        true
119    }
120}
121
122#[cfg(test)]
123mod tests {
124    use super::*;
125    use ego_tree::tree;
126
127    #[test]
128    fn dfs() {
129        let tree = tree!('a' => { 'b' => { 'd', 'e' }, 'c' });
130        let values: Vec<_> = tree.dfs().map(|node| *node.value()).collect();
131        assert_eq!(values, vec!['a', 'b', 'd', 'e', 'c']);
132    }
133
134    #[test]
135    fn bfs() {
136        let tree = tree!('a' => { 'b' => { 'd', 'e' }, 'c' });
137        let values: Vec<_> = tree.bfs().map(|node| *node.value()).collect();
138        assert_eq!(values, vec!['a', 'b', 'c', 'd', 'e']);
139
140        let tree = tree!('a' => { 'b' => {'d'}, 'c' => {'e'} });
141        let mut traverse = tree.bfs();
142        traverse.next();
143        assert!(traverse.peek_at_same_level().is_none());
144        assert_eq!(traverse.nth(2).map(|node| *node.value()), Some('d'));
145        assert_eq!(traverse.peek_at_same_level().map(|node| *node.value()), Some('e'));
146    }
147
148    #[test]
149    fn isomorphism() {
150        let tree = tree!('a' => { 'b' => { 'd', 'e' }, 'c' });
151        assert!(tree.is_isomorphic_to(&tree));
152
153        let other = tree!('a' => { 'b' => { 'd' }, 'e' => { 'c' }});
154        let tree_dfs_values: Vec<_> = tree.dfs().map(|node| *node.value()).collect();
155        let other_dfs_values: Vec<_> = other.dfs().map(|node| *node.value()).collect();
156        assert_eq!(tree_dfs_values, other_dfs_values);
157        assert!(!tree.is_isomorphic_to(&other));
158    }
159}
catlog/dbl/tree_algorithms.rs

catlog/dbl/
tree_algorithms.rs
