dce_util/
arena_tree.rs

1use std::collections::BTreeMap;
2use std::ops::{Deref, DerefMut};
3
4pub trait KeyFactory<K> {
5    fn id(&self) -> K;
6    fn child_of(&self, parent: &Self) -> bool;
7    fn new_parent(&self) -> Self;
8}
9
10pub struct ArenaNode<E> {
11    element: E,
12    parent: Option<usize>,
13    children: Vec<usize>,
14}
15
16impl <E> ArenaNode<E> {
17    fn new(element: E, parent: Option<usize>) -> Self {
18        Self { element, parent, children: Default::default() }
19    }
20
21    pub fn element(&self) -> &E {
22        &self.element
23    }
24
25    pub fn element_mut(&mut self) -> &mut E {
26        &mut self.element
27    }
28
29    pub fn parent_index(&self) -> &Option<usize> {
30        &self.parent
31    }
32
33    pub fn parent_mut<'a, K>(&self, tree: &'a mut ArenaTree<E, K>) -> Option<&'a mut ArenaNode<E>> {
34        if let Some(parent_index) = self.parent {
35            return tree.nodes.get_mut(parent_index);
36        }
37        None
38    }
39
40    pub fn child_indexes(&self) -> &Vec<usize> {
41        &self.children
42    }
43}
44
45pub struct ArenaTree<E, K> {
46    nodes: Vec<ArenaNode<E>>,
47    mapping: BTreeMap<K, usize>,
48}
49
50impl<E: KeyFactory<K>, K: Ord + Clone> ArenaTree<E, K> {
51    pub fn new(element: E) -> Self {
52        let id = element.id();
53        let index = 0;
54        let node = ArenaNode::new(element, None);
55        Self { nodes: vec![node], mapping: BTreeMap::from([(id, index)]) }
56    }
57
58    pub fn insert(&mut self, element: E) -> (K, usize) {
59        let id = element.id();
60        let index = self.nodes.len();
61        let parent_index = self.find_parent(&element);
62        let node = ArenaNode::new(element, parent_index);
63        self.nodes.push(node);
64        self.mapping.insert(id.clone(), index);
65        (id, index)
66    }
67
68    fn find_parent(&self, element: &E) -> Option<usize> {
69        self.nodes.iter().enumerate()
70            .find(|&(_, node)| element.child_of(&node.element))
71            .map(|(index, _)| index)
72    }
73
74    pub fn by_id(&self, id: &K) -> Option<&ArenaNode<E>> {
75        self.mapping.get(id).into_iter()
76            .flat_map(|i| self.by_index(*i))
77            .next()
78    }
79
80    pub fn by_id_mut(&mut self, id: &K) -> Option<&mut ArenaNode<E>> {
81        if let Some(index) = self.mapping.get(id) {
82            return self.by_index_mut(*index);
83        }
84        None
85    }
86
87    pub fn by_index(&self, index: usize) -> Option<&ArenaNode<E>> {
88        self.nodes.get(index)
89    }
90
91    pub fn by_index_mut(&mut self, index: usize) -> Option<&mut ArenaNode<E>> {
92        self.nodes.get_mut(index)
93    }
94
95    pub fn fill(&mut self, mut elements: Vec<E>) {
96        // 1. Insert all
97        while elements.len() > 0 {
98            self.insert(elements.remove(0));
99        }
100        // 2. Find all hanging indexes
101        let mut hanging_indexes = self.nodes.iter().enumerate()
102            .filter(|&(i, n)| i > 0 && matches!(n.parent, None))
103            .map(|(i, _)| i).collect::<Vec<_>>();
104        // 3. Fill the gap branches
105        while hanging_indexes.len() > 0 {
106            let hanging_index = hanging_indexes.remove(0);
107            let hanging_node = self.by_index(hanging_index).unwrap();
108            if matches!(hanging_node.parent, Some(_)) {
109            } else if let Some(parent_index) = self.find_parent(&hanging_node.element) {
110                self.by_index_mut(hanging_index).iter_mut().for_each(|n| n.parent = Some(parent_index));
111            } else {
112                let parent_element = hanging_node.element.new_parent();
113                let (_, parent_index) = self.insert(parent_element);
114                // The parent may not bound a grand, just push into the hanging_indexes
115                hanging_indexes.push(parent_index);
116                self.by_index_mut(hanging_index).iter_mut().for_each(|n| n.parent = Some(parent_index));
117            }
118        }
119    }
120}
121
122impl <E, K> Deref for ArenaTree<E, K> {
123    type Target = Vec<ArenaNode<E>>;
124
125    fn deref(&self) -> &Self::Target {
126        &self.nodes
127    }
128}
129
130impl <E, K> DerefMut for ArenaTree<E, K> {
131    fn deref_mut(&mut self) -> &mut Self::Target {
132        &mut self.nodes
133    }
134}
dce_util/arena_tree.rs

dce_util/
arena_tree.rs
