arena_voxel_tree/
arena_node.rs

1//! [`Arena`] is defined here.
2//! * This library is inspired by [r3bl-org work](https://github.com/r3bl-org/)
3
4use std::{
5    collections::{HashMap, VecDeque},
6    fmt::Debug,
7    sync::{atomic::AtomicUsize, Arc, RwLock},
8};
9
10use super::arena_types::{ArenaMap, FilterFn, HasId, NodeRef, ResultUidList, WeakNodeRef};
11use crate::utils::{
12    call_if_some, unwrap_arc_read_lock_and_call, unwrap_arc_write_lock_and_call, with_mut,
13};
14
15/// This struct represents a node in a tree. It may have a parent. It can hold multiple children.
16/// And it has a payload. It also has an id that uniquely identifies it. An [`Arena`] or
17/// [`super::MTArena`] is used to hold nodes.
18#[derive(Debug)]
19pub struct Node<T>
20where
21    T: Debug + Clone + Send + Sync,
22{
23    pub id: usize,
24    pub parent_id: Option<usize>,
25    pub children_ids: VecDeque<usize>,
26    pub payload: T,
27}
28
29impl<T> HasId for Node<T>
30where
31    T: Debug + Clone + Send + Sync,
32{
33    type IdType = usize;
34
35    /// Delegate this to `self.id`, which is type `usize`.
36    fn get_id(&self) -> Self::IdType {
37        self.id.get_id()
38    }
39}
40
41#[derive(Debug)]
42pub struct Arena<T>
43where
44    T: Debug + Clone + Send + Sync,
45{
46    map: RwLock<ArenaMap<T>>,
47    atomic_counter: AtomicUsize,
48}
49
50impl<T> Arena<T>
51where
52    T: Debug + Clone + Send + Sync,
53{
54    /// If no matching nodes can be found returns `None`.
55    pub fn filter_all_nodes_by(&self, filter_fn: &FilterFn<T>) -> ResultUidList {
56        self.map
57            .read()
58            .map_or(None, |map /* ReadGuarded<'_, ArenaMap<T>> */| {
59                let filtered_map = map
60                    .iter()
61                    .filter(|(id, node_ref)| {
62                        filter_fn(**id, node_ref.read().unwrap().payload.clone())
63                    })
64                    .map(|(id, _node_ref)| *id)
65                    .collect::<VecDeque<usize>>();
66                match filtered_map.len() {
67                    0 => None,
68                    _ => Some(filtered_map),
69                }
70            })
71    }
72
73    /// If `node_id` can't be found, returns `None`.
74    pub fn get_children_of(&self, node_id: usize) -> ResultUidList {
75        if !self.node_exists(node_id) {
76            return None;
77        }
78
79        let node_to_lookup = self.get_node_arc(node_id)?;
80
81        let result =
82            if let Ok(node_to_lookup /* ReadGuarded<'_, Node<T>> */) = node_to_lookup.read() {
83                if node_to_lookup.children_ids.is_empty() {
84                    return None;
85                }
86                Some(node_to_lookup.children_ids.clone())
87            } else {
88                None
89            };
90
91        result
92    }
93
94    /// If `node_id` can't be found, returns `None`.
95    pub fn get_parent_of(&self, node_id: usize) -> Option<usize> {
96        if !self.node_exists(node_id) {
97            return None;
98        }
99
100        let node_to_lookup = self.get_node_arc(node_id)?;
101        node_to_lookup
102            .read()
103            .map_or(None, |node_to_lookup /* ReadGuarded<'_, Node<T>> */| {
104                node_to_lookup.parent_id
105            })
106    }
107
108    pub fn node_exists(&self, node_id: usize) -> bool {
109        self.map.read().unwrap().contains_key(&node_id.get_id())
110    }
111
112    pub fn has_parent(&self, node_id: usize) -> bool {
113        if self.node_exists(node_id) {
114            let parent_id_opt = self.get_parent_of(node_id);
115            if let Some(parent_id) = parent_id_opt {
116                return self.node_exists(parent_id);
117            }
118        }
119        false
120    }
121
122    /// If `node_id` can't be found, returns `None`.
123    pub fn delete_node(&self, node_id: usize) -> ResultUidList {
124        if !self.node_exists(node_id) {
125            return None;
126        }
127        let deletion_list = self.tree_walk_dfs(node_id)?;
128
129        // Note - this lambda expects that `parent_id` exists.
130        let remove_node_id_from_parent = |parent_id: usize| {
131            let parent_node_arc_opt = self.get_node_arc(parent_id);
132            if let Some(parent_node_arc) = parent_node_arc_opt {
133                if let Ok(mut parent_node /* WriteGuarded<'_, Node<T>> */) = parent_node_arc.write()
134                {
135                    parent_node
136                        .children_ids
137                        .retain(|child_id| *child_id != node_id);
138                }
139            }
140        };
141
142        // If `node_id` has a parent, remove `node_id` its children, otherwise skip this
143        // step.
144        if self.has_parent(node_id) {
145            if let Some(parent_id) = self.get_parent_of(node_id) {
146                remove_node_id_from_parent(parent_id);
147            }
148        }
149
150        // Actually delete the nodes in the deletion list.
151        if let Ok(mut map /* WriteGuarded<'_, ArenaMap<T>> */) = self.map.write() {
152            for node_id in &deletion_list {
153                map.remove(node_id);
154            }
155        }
156        // Pass the deletion list back.
157        deletion_list.into()
158    }
159
160    /// - [DFS graph walking](https://developerlife.com/2018/08/16/algorithms-in-kotlin-5/)
161    /// - [DFS tree walking](https://stephenweiss.dev/algorithms-depth-first-search-dfs#handling-non-binary-trees)
162    pub fn tree_walk_dfs(&self, node_id: usize) -> ResultUidList {
163        if !self.node_exists(node_id) {
164            return None;
165        }
166
167        let mut stack: VecDeque<usize> = VecDeque::from([node_id.get_id()]);
168
169        let mut it: VecDeque<usize> = VecDeque::new();
170
171        while let Some(node_id) = stack.pop_back() {
172            // Question mark operator works below, since it returns a `Option` to `while let ...`.
173            // Basically skip to the next item in the `stack` if `node_id` can't be found.
174            let node_ref = self.get_node_arc(node_id)?;
175            unwrap_arc_read_lock_and_call(&node_ref, &mut |node| {
176                it.push_back(node.get_id());
177                // Note that the children ordering has to be flipped! You want to perform the
178                // traversal from RIGHT -> LEFT (not LEFT -> RIGHT).
179                // PTAL: <https://developerlife.com/assets/algo-ts-2-images/depth-first-search.svg>
180                for child_id in node.children_ids.iter().rev() {
181                    stack.push_back(*child_id);
182                }
183            });
184        }
185
186        match it.len() {
187            0 => None,
188            _ => Some(it),
189        }
190    }
191
192    /// - [BFS graph walking](https://developerlife.com/2018/08/16/algorithms-in-kotlin-5/)
193    /// - [BFS tree walking](https://stephenweiss.dev/algorithms-depth-first-search-dfs#handling-non-binary-trees)
194    pub fn tree_walk_bfs(&self, node_id: usize) -> ResultUidList {
195        if !self.node_exists(node_id) {
196            return None;
197        }
198
199        let mut queue: VecDeque<usize> = VecDeque::from([node_id.get_id()]);
200
201        let mut it: VecDeque<usize> = VecDeque::new();
202
203        while let Some(node_id) = queue.pop_front() {
204            // Question mark operator works below, since it returns a `Option` to `while let ...`.
205            // Basically skip to the next item in the `stack` if `node_id` can't be found.
206            let node_ref = self.get_node_arc(node_id)?;
207            unwrap_arc_read_lock_and_call(&node_ref, &mut |node| {
208                it.push_back(node.get_id());
209                for child_id in node.children_ids.iter() {
210                    queue.push_back(*child_id);
211                }
212            });
213        }
214
215        match it.len() {
216            0 => None,
217            _ => Some(it),
218        }
219    }
220
221    /// If `node_id` can't be found, returns `None`. More info on
222    /// [`Option.map()`](https://play.rust-lang.org/?version=stable&mode=debug&edition=2021&gist=d5a54a042fea085ef8c9122b7ea47c6a)
223    pub fn get_node_arc_weak(&self, node_id: usize) -> Option<WeakNodeRef<T>> {
224        if !self.node_exists(node_id) {
225            return None;
226        }
227        self.map.read().map_or_else(
228            |_| None, // Runs if `node_ref` is some, else returns `None`.
229            |map| {
230                map.get(&node_id.get_id()) // Returns `None` if `node_id` doesn't exist.
231                    .map(Arc::downgrade)
232            },
233        )
234    }
235
236    /// If `node_id` can't be found, returns `None`. More info on
237    /// [`Option.map()`](https://play.rust-lang.org/?version=stable&mode=debug&edition=2021&gist=d5a54a042fea085ef8c9122b7ea47c6a)
238    pub fn get_node_arc(&self, node_id: usize) -> Option<NodeRef<T>> {
239        if !self.node_exists(node_id) {
240            return None;
241        }
242        self.map.read().map_or_else(
243            |_| None, // Runs if `node_ref` is some, else returns `None`.
244            |map| {
245                map.get(&node_id.get_id()) // Returns `None` if `node_id` doesn't exist.
246                    .map(Arc::clone)
247            },
248        )
249    }
250
251    /// Note `data` is cloned to avoid `data` being moved. If `parent_id` can't be found, it panics.
252    pub fn add_new_node(&mut self, payload: T, maybe_parent_id: Option<usize>) -> usize {
253        let parent_id_arg_provided = maybe_parent_id.is_some();
254
255        // Check to see if `parent_id` exists.
256        if parent_id_arg_provided {
257            let parent_id = maybe_parent_id.unwrap();
258            if !self.node_exists(parent_id) {
259                panic!("Parent node doesn't exist.");
260            }
261        }
262
263        let new_node_id = self.generate_uid();
264
265        // Create a new node from payload & add it to the arena.
266        with_mut(&mut self.map.write().unwrap(), &mut |map| {
267            let value = Arc::new(RwLock::new(Node {
268                id: new_node_id,
269                parent_id: if parent_id_arg_provided {
270                    let parent_id = maybe_parent_id.unwrap();
271                    Some(parent_id.get_id())
272                } else {
273                    None
274                },
275                children_ids: VecDeque::new(),
276                payload: payload.clone(),
277            }));
278            map.insert(new_node_id, value);
279        });
280
281        // Deal w/ adding this newly created node to the parent's children list.
282        if let Some(parent_id) = maybe_parent_id {
283            let maybe_parent_node_arc = self.get_node_arc(parent_id);
284            call_if_some(&maybe_parent_node_arc, &|parent_node_arc| {
285                unwrap_arc_write_lock_and_call(parent_node_arc, &mut |parent_node| {
286                    // Preserve the natural order of insertion.
287                    parent_node.children_ids.push_back(new_node_id);
288                });
289            });
290        }
291
292        // Return the node identifier.
293        new_node_id
294    }
295
296    fn generate_uid(&self) -> usize {
297        self.atomic_counter
298            .fetch_add(1, std::sync::atomic::Ordering::SeqCst)
299    }
300
301    pub fn new() -> Self {
302        Arena {
303            map: RwLock::new(HashMap::new()),
304            atomic_counter: AtomicUsize::new(0),
305        }
306    }
307}
308
309impl<T> Default for Arena<T>
310where
311    T: Debug + Clone + Send + Sync,
312{
313    fn default() -> Self {
314        Self::new()
315    }
316}