transact/scheduler/parallel/
tree.rs

1use std::cell::RefCell;
2use std::collections::BTreeMap;
3use std::collections::VecDeque;
4use std::error::Error as StdError;
5use std::rc::Rc;
6
7#[derive(Debug)]
8pub enum RadixTreeError {
9    AddressNotInTree(String),
10}
11
12impl StdError for RadixTreeError {
13    fn description(&self) -> &str {
14        match *self {
15            RadixTreeError::AddressNotInTree(ref msg) => msg,
16        }
17    }
18}
19
20impl std::fmt::Display for RadixTreeError {
21    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
22        match *self {
23            RadixTreeError::AddressNotInTree(ref s) => write!(f, "AddressNotInTree: {}", s),
24        }
25    }
26}
27
28#[derive(Debug, Clone, Default)]
29pub struct Node<T> {
30    address: String,
31    children: BTreeMap<String, Rc<RefCell<Node<T>>>>,
32    data: Option<T>,
33}
34
35#[derive(Default, Debug, Clone)]
36pub struct RadixTree<T> {
37    root: Rc<RefCell<Node<T>>>,
38}
39
40/// This radix tree is a prefix tree: a node's address is always a strict prefix of the addresses
41/// of its children, and every node either has data or has multiple children.
42impl<T: Clone> RadixTree<T> {
43    pub fn new() -> Self {
44        RadixTree {
45            root: Rc::new(RefCell::new(Node {
46                address: "".to_string(),
47                children: BTreeMap::new(),
48                data: None,
49            })),
50        }
51    }
52
53    fn get_child(
54        node: &Rc<RefCell<Node<T>>>,
55        address: &str,
56    ) -> Result<Rc<RefCell<Node<T>>>, RadixTreeError> {
57        node.borrow()
58            .children
59            .values()
60            .find(|child| address.starts_with(&child.borrow().address))
61            .map(Rc::clone)
62            .ok_or_else(|| RadixTreeError::AddressNotInTree("Address Not In Tree".to_string()))
63    }
64
65    /// Collects all children below ADDRESS, as well as the children's descendants
66    fn walk_to_address(&self, address: &str) -> Vec<Rc<RefCell<Node<T>>>> {
67        let mut current_node = Rc::clone(&self.root);
68        let mut results = vec![Rc::clone(&current_node)];
69
70        // A node's address is always a proper prefix of the addresses of its children
71        while address != current_node.borrow().address.as_str()
72            && address.starts_with(current_node.borrow().address.as_str())
73        {
74            match RadixTree::get_child(&current_node, address) {
75                Ok(child) => {
76                    results.push(Rc::clone(&child));
77                    current_node = child;
78                }
79                Err(_) => break,
80            }
81        }
82        results
83    }
84
85    /// Return a vector of tuple pairs of node addresses and data
86    /// The Node address is the key, the data is the value.
87    /// First the ancestors of ADDRESS (including self) are yielded, earliest to latest, and
88    /// then the descendants of ADDRESS are yielded
89    pub fn walk(&self, address: &str) -> Vec<(String, Option<T>)> {
90        let mut return_nodes = Vec::new();
91        let accumulated_nodes = self.walk_to_address(address);
92        for node in accumulated_nodes.iter() {
93            return_nodes.push((node.borrow().address.clone(), node.borrow().data.clone()));
94        }
95
96        if let Some(node) = accumulated_nodes.iter().last() {
97            let mut to_process = VecDeque::new();
98            let descendants = node.borrow().children.clone();
99            for descendant in descendants.values() {
100                to_process.push_front(Rc::clone(descendant));
101            }
102            while let Some(current_child) = to_process.pop_front() {
103                return_nodes.push((
104                    current_child.borrow().address.clone(),
105                    current_child.borrow().data.clone(),
106                ));
107                let additional_descendants = &current_child.borrow().children;
108                for child in additional_descendants.values() {
109                    to_process.push_front(Rc::clone(child));
110                }
111            }
112        }
113        return_nodes
114    }
115
116    /// Walk as far down the tree as possible. If the desired address is reached, return that node.
117    /// Otherwise, add a new one.
118    fn get_or_create(&self, address: &str) -> Rc<RefCell<Node<T>>> {
119        let accumulated_nodes = self.walk_to_address(address);
120        let first_ancestor = accumulated_nodes
121            .iter()
122            .last()
123            .expect("Node ancestors not formed correctly");
124
125        if first_ancestor.borrow().address == address {
126            return Rc::clone(first_ancestor);
127        }
128
129        // The address isn't in the tree, so a new node will need to be added
130        let new_node = Rc::new(RefCell::new(Node {
131            address: address.to_string(),
132            children: BTreeMap::new(),
133            data: None,
134        }));
135
136        // Attempt to get the next child with a matching prefix.
137        let prefix_len = first_ancestor.borrow().address.len();
138        let option_ancestor_child = first_ancestor
139            .borrow()
140            .children
141            .values()
142            .find(|child| {
143                let child_address = &child.borrow().address;
144                let child_address_prefix: String =
145                    child_address.chars().skip(prefix_len).collect::<String>();
146                let address_prefix: String =
147                    address.chars().skip(prefix_len).take(1).collect::<String>();
148                child_address_prefix.starts_with(&address_prefix)
149            })
150            .map(Rc::clone);
151
152        // Checks if the next child with a matching prefix was found, else just adds the new
153        // address as a child.
154        let ancestor_child = match option_ancestor_child {
155            Some(child) => child,
156            None => {
157                first_ancestor
158                    .borrow_mut()
159                    .children
160                    .insert(address.to_string(), Rc::clone(&new_node));
161                return new_node;
162            }
163        };
164
165        // If node address is 'rustic' and the address being added is 'rust',
166        // then 'rust' will be the intermediate node taking 'rustic' as a child.
167        if ancestor_child.borrow().address.starts_with(address) {
168            first_ancestor
169                .borrow_mut()
170                .children
171                .insert(address.to_string(), Rc::clone(&new_node));
172
173            new_node
174                .borrow_mut()
175                .children
176                .insert(address.to_string(), Rc::clone(&ancestor_child));
177            first_ancestor
178                .borrow_mut()
179                .children
180                .remove(&ancestor_child.borrow().address);
181            return new_node;
182        };
183
184        // The address and the match address share a common prefix, so
185        // an intermediate node with the prefix as its address will
186        // take them both as children.
187        let ancestor_child_address = ancestor_child.borrow().address.clone();
188        let shorter = if address.len() < ancestor_child_address.len() {
189            address
190        } else {
191            &ancestor_child_address
192        };
193        let intermediate_node = Rc::new(RefCell::new(Node {
194            address: String::from(""),
195            children: BTreeMap::new(),
196            data: None,
197        }));
198
199        for i in 0..shorter.len() {
200            if address.chars().nth(i) != ancestor_child_address.chars().nth(i) {
201                intermediate_node.borrow_mut().address = shorter[..i].to_string();
202                break;
203            }
204        }
205        let mut new_children_map = BTreeMap::new();
206        new_children_map.insert(new_node.borrow().address.clone(), Rc::clone(&new_node));
207        new_children_map.insert(
208            ancestor_child.borrow().address.clone(),
209            Rc::clone(&ancestor_child),
210        );
211        intermediate_node
212            .borrow_mut()
213            .children
214            .append(&mut new_children_map);
215        first_ancestor.borrow_mut().children.insert(
216            intermediate_node.borrow().address.clone(),
217            Rc::clone(&intermediate_node),
218        );
219        first_ancestor
220            .borrow_mut()
221            .children
222            .remove(&ancestor_child.borrow().address.clone());
223
224        new_node
225    }
226
227    /// Walk to ADDRESS, creating nodes if necessary, and set the data there to
228    /// UPDATER(data)
229    pub fn update(&self, address: &str, updater: &dyn Fn(Option<T>) -> Option<T>, prune: bool) {
230        let node = self.get_or_create(address);
231        let existing_data = node.borrow_mut().data.take();
232        node.borrow_mut().data = updater(existing_data);
233
234        if prune {
235            node.borrow_mut().children.clear();
236        }
237    }
238
239    /// Remove all children (and descendants) below ADDRESS
240    pub fn prune(&self, address: &str) {
241        let accumulated_nodes = self.walk_to_address(address);
242        if let Some(node) = accumulated_nodes.iter().last() {
243            node.borrow_mut().children.clear()
244        }
245    }
246}
247
248#[cfg(test)]
249mod tests {
250    use super::*;
251
252    #[test]
253    fn tree_creation() {
254        let tree: RadixTree<i32> = RadixTree::new();
255        assert_eq!(tree.root.borrow().children.len(), 0);
256        assert_eq!(tree.root.borrow().address, "".to_string());
257    }
258
259    #[test]
260    fn tree_insert_children() {
261        // R
262        //  A
263        //   D
264        //    I
265        //     S
266        //      H
267        //     X
268        //    O
269        //     N
270        let tree: RadixTree<i32> = RadixTree::new();
271        tree.get_or_create("radix");
272        tree.get_or_create("radish");
273        tree.get_or_create("radon");
274
275        assert_eq!(tree.root.borrow().children.len(), 1);
276
277        let found_node_rad = tree.get_or_create("rad");
278        assert_eq!(found_node_rad.borrow().address, "rad".to_string());
279        assert_eq!(found_node_rad.borrow().children.len(), 2);
280
281        let found_node_radi = tree.get_or_create("radi");
282        assert_eq!(found_node_radi.borrow().address, "radi".to_string());
283        assert_eq!(found_node_radi.borrow().children.len(), 2);
284
285        let found_node_radix = tree.get_or_create("radix");
286        assert_eq!(found_node_radix.borrow().address, "radix".to_string());
287        assert_eq!(found_node_radix.borrow().children.len(), 0);
288
289        let found_node_radish = tree.get_or_create("radish");
290        assert_eq!(found_node_radish.borrow().address, "radish".to_string());
291        assert_eq!(found_node_radish.borrow().children.len(), 0);
292
293        let found_node_radon = tree.get_or_create("radon");
294        assert_eq!(found_node_radon.borrow().address, "radon".to_string());
295        assert_eq!(found_node_radon.borrow().children.len(), 0);
296    }
297
298    #[test]
299    fn tree_walk_to_address() {
300        // R
301        //  A
302        //   D
303        //    I
304        //     S
305        //      H
306        //     X
307        //    O
308        //     N
309        let tree: RadixTree<i32> = RadixTree::new();
310        tree.get_or_create("radix");
311        tree.get_or_create("radish");
312        tree.get_or_create("radon");
313
314        let walk_to_results_rad = tree.walk_to_address("rad");
315        assert_eq!(walk_to_results_rad.len(), 2);
316
317        let walk_to_results_radon = tree.walk_to_address("radon");
318        assert_eq!(walk_to_results_radon.len(), 3);
319
320        let walk_to_results_radix = tree.walk_to_address("radix");
321        assert_eq!(walk_to_results_radix.len(), 4);
322    }
323
324    #[test]
325    fn tree_walk() {
326        // R
327        //  A
328        //   D
329        //    I
330        //     S
331        //      H
332        //     X
333        //    O
334        //     N
335        let tree: RadixTree<i32> = RadixTree::new();
336        tree.get_or_create("radix");
337        tree.get_or_create("radish");
338        tree.get_or_create("radon");
339
340        let walk_results_radix = tree.walk("radix");
341        assert_eq!(walk_results_radix.len(), 4);
342        assert!(walk_results_radix.contains(&("radix".to_string(), None)));
343
344        let walk_results_rad = tree.walk("rad");
345        assert_eq!(walk_results_rad.len(), 6);
346        assert!(walk_results_rad.contains(&("rad".to_string(), None)));
347    }
348
349    fn update_data(data: Option<i32>) -> Option<i32> {
350        if data.is_none() {
351            return Some(1);
352        } else {
353            return Some(data.unwrap() + 1);
354        }
355    }
356
357    #[test]
358    fn tree_node_update() {
359        // R
360        //  A
361        //   D
362        //    I
363        //     S
364        //      H
365        //     X
366        //    O
367        //     N
368        let tree: RadixTree<i32> = RadixTree::new();
369        tree.get_or_create("radix");
370        tree.get_or_create("radish");
371        tree.get_or_create("radon");
372
373        tree.update("radix", &update_data, false);
374        let updated_node = tree.get_or_create("radix");
375        assert_eq!(updated_node.borrow().data, Some(1));
376        tree.update("radix", &update_data, false);
377        assert_eq!(updated_node.borrow().data, Some(2));
378    }
379
380    #[test]
381    fn tree_prune() {
382        // R
383        //  A
384        //   D
385        //    I
386        //     S
387        //      H
388        //     X
389        //    O
390        //     N
391        let tree: RadixTree<i32> = RadixTree::new();
392        tree.get_or_create("radix");
393        tree.get_or_create("radish");
394        tree.get_or_create("radon");
395
396        tree.prune("rad");
397        let parent_node = tree.get_or_create("rad");
398        let mut parent_node_walk = tree.walk("rad");
399
400        assert_eq!(parent_node.borrow().children.len(), 0);
401        assert!(!parent_node_walk.contains(&("radi".to_string(), None)));
402        assert!(parent_node_walk.contains(&("rad".to_string(), None)));
403
404        tree.get_or_create("radish");
405        tree.get_or_create("radix");
406        parent_node_walk = tree.walk("rad");
407
408        assert_eq!(parent_node.borrow().children.len(), 1);
409        assert!(parent_node_walk.contains(&("radix".to_string(), None)));
410        assert!(parent_node_walk.contains(&("radish".to_string(), None)));
411    }
412}
transact/scheduler/parallel/tree.rs

transact/scheduler/parallel/
tree.rs
