dust_mail/
tree.rs

1#[cfg(feature = "serde")]
2use serde::{Deserialize, Serialize};
3
4#[derive(Debug, PartialEq)]
5#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
6pub enum Node<T: Default> {
7    Branch { data: T, children: Vec<Node<T>> },
8    Root(Vec<Node<T>>),
9    Leaf(T),
10}
11
12impl<T: Default> Default for Node<T> {
13    fn default() -> Self {
14        Self::Leaf(T::default())
15    }
16}
17
18impl<T: Default> From<T> for Node<T> {
19    fn from(value: T) -> Self {
20        Self::Leaf(value)
21    }
22}
23
24pub trait Find<T> {
25    fn find(&self, item: &T) -> bool;
26}
27
28impl<T: Default> Node<T> {
29    pub fn branch<C: IntoIterator<Item = Node<T>>>(data: T, children: C) -> Self {
30        Node::Branch {
31            data,
32            children: children.into_iter().collect(),
33        }
34    }
35
36    pub fn empty_root() -> Self {
37        Node::Root(Vec::new())
38    }
39
40    pub fn empty_branch(data: T) -> Self {
41        Node::branch(data, Vec::new())
42    }
43
44    pub fn leaf(data: T) -> Self {
45        Node::Leaf(data)
46    }
47
48    pub fn data(&self) -> Option<&T> {
49        match &self {
50            Node::Branch { data, .. } | Node::Leaf(data) => Some(data),
51            _ => None,
52        }
53    }
54
55    pub fn data_mut(&mut self) -> Option<&mut T> {
56        match self {
57            Node::Branch { data, .. } | Node::Leaf(data) => Some(data),
58            _ => None,
59        }
60    }
61
62    pub fn into_data(self) -> Option<T> {
63        match self {
64            Node::Branch { data, .. } | Node::Leaf(data) => Some(data),
65            _ => None,
66        }
67    }
68
69    /// Inserts a node if this node is capable of containing children
70    pub fn insert(&mut self, node: Node<T>) -> bool {
71        match self {
72            Node::Root(children) | Node::Branch { children, .. } => {
73                children.push(node);
74
75                true
76            }
77            _ => false,
78        }
79    }
80
81    /// Turns all (sub)branches with no children into leaves
82    pub fn create_leaves(self) -> Self {
83        match self {
84            Node::Root(mut children) if children.len() == 1 => {
85                Node::create_leaves(children.remove(0))
86            }
87            Node::Branch { data, children } if children.is_empty() => Node::Leaf(data),
88            Node::Branch { children, data } => Node::Branch {
89                data,
90                children: children.into_iter().map(Node::create_leaves).collect(),
91            },
92            Node::Root(children) => {
93                Node::Root(children.into_iter().map(Node::create_leaves).collect())
94            }
95            _ => self,
96        }
97    }
98
99    pub fn find<P: Find<T>>(&self, predicate: &P) -> Option<&Self> {
100        match self {
101            Node::Leaf(data) | Node::Branch { data, .. } if predicate.find(data) => Some(self),
102            Node::Root(children) | Node::Branch { children, .. } => {
103                for child in children {
104                    if let Some(found) = Self::find(child, predicate) {
105                        return Some(found);
106                    }
107                }
108
109                None
110            }
111            _ => None,
112        }
113    }
114
115    pub fn find_mut<P: Find<T>>(&mut self, predicate: &P) -> Option<&mut Self> {
116        match self {
117            Node::Leaf(data) | Node::Branch { data, .. } if predicate.find(&data) => Some(self),
118            Node::Root(children) | Node::Branch { children, .. } => {
119                for child in children {
120                    if let Some(found) = Self::find_mut(child, predicate) {
121                        return Some(found);
122                    }
123                }
124
125                None
126            }
127            _ => None,
128        }
129    }
130
131    pub fn into_find<P: Find<T>>(self, predicate: &P) -> Option<Self> {
132        match self {
133            Node::Leaf(ref data) | Node::Branch { ref data, .. } if predicate.find(data) => {
134                Some(self)
135            }
136            Node::Root(children) | Node::Branch { children, .. } => {
137                for child in children {
138                    if let Some(found) = Self::into_find(child, predicate) {
139                        return Some(found);
140                    }
141                }
142
143                None
144            }
145            _ => None,
146        }
147    }
148}
149
150#[cfg(test)]
151mod test {
152    use super::*;
153
154    struct GreaterThanThree;
155
156    impl Find<i32> for GreaterThanThree {
157        fn find(&self, item: &i32) -> bool {
158            *item > 3
159        }
160    }
161
162    struct GreaterThanFour;
163
164    impl Find<i32> for GreaterThanFour {
165        fn find(&self, item: &i32) -> bool {
166            *item > 4
167        }
168    }
169
170    #[test]
171    fn test_find() {
172        let test_tree = Node::branch(1, vec![2.into(), 3.into(), Node::branch(4, Vec::new())]);
173
174        assert_eq!(Some(&4), test_tree.find(&GreaterThanThree).unwrap().data());
175
176        assert_eq!(None, test_tree.find(&GreaterThanFour));
177    }
178}
dust_mail/tree.rs

dust_mail/
tree.rs
