org_rust_parser/
node_pool.rs

1use std::fmt::{Debug, Display};
2use std::ops::{Index, IndexMut};
3
4use crate::types::{Expr, Node};
5
6#[derive(Clone, Copy, Hash, PartialEq, PartialOrd, Ord, Eq)]
7/// Identifier for [`Node`]s in a [`NodePool`].
8///
9/// NodeIDs are guaranteed to be unique to each node since they are assigned
10/// sequentially and cannot re-used.
11pub struct NodeID(u32);
12
13/// This exists ONLY for testing purposes
14pub(crate) fn make_node_id(id: u32) -> NodeID {
15    NodeID(id)
16}
17
18impl Display for NodeID {
19    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
20        f.write_fmt(format_args!("{}", self.0))
21    }
22}
23
24impl std::fmt::Debug for NodeID {
25    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
26        f.write_fmt(format_args!("{}", self.0))
27    }
28}
29
30/// Arena-based container storing [`Node`]s
31///
32/// A [`NodePool`] is essentially just a [`Vec<Node>`] that is indexed by [`NodeID`]s.
33/// Each parsing construct stores [`NodeID`]s which refer to indices within the [`NodePool`].
34/// It acts like a flattened version of the AST where each node can be easily iterated over
35/// and tracked.
36///
37/// [`NodePool::iter`] and [`NodePool::iter_mut`] can be used for iterating over the contents of
38/// the pool.
39///
40/// # Safety
41///
42/// Removing or inserting an element at a position other than the end might invalidate every other NodeID
43/// and therefore the AST. Removing and de-allocating an individual [`Node`] is not possible due to
44/// this effect. However, the result can be feigned in the tree by "deleting" the node from its parent.
45/// See [`NodePool::delete_node`].
46///
47///
48#[derive(Debug)]
49pub struct NodePool<'a> {
50    pub inner_vec: Vec<Node<'a>>,
51    pub counter: u32,
52}
53
54impl<'a> NodePool<'a> {
55    pub(crate) fn new() -> Self {
56        Self {
57            inner_vec: Vec::new(),
58            // The next free index in the pool.
59            counter: 0,
60        }
61    }
62
63    pub(crate) fn alloc<T>(
64        &mut self,
65        obj: T,
66        start: usize,
67        end: usize,
68        parent: Option<NodeID>,
69    ) -> NodeID
70    where
71        Expr<'a>: From<T>,
72    {
73        let prev_id = self.counter;
74        self.inner_vec.push(Node::new(obj, start, end, parent));
75        self.counter += 1;
76        NodeID(prev_id)
77    }
78
79    /// Allocates a node in the pool at a given location.
80    ///
81    /// Returns the index that was allocated.
82    ///
83    /// Works well with [`NodePool::reserve_id`].
84    ///
85    /// # Safety:
86    ///
87    /// Must refer to an ID that already exists in the pool.
88    /// Will panic at runtime otherwise.
89    pub(crate) fn alloc_with_id<T>(
90        &mut self,
91        obj: T,
92        start: usize,
93        end: usize,
94        parent: Option<NodeID>,
95        target_id: NodeID,
96    ) -> NodeID
97    where
98        Expr<'a>: From<T>,
99    {
100        self.inner_vec[target_id.0 as usize] = Node::new(obj, start, end, parent);
101
102        target_id
103    }
104
105    pub fn get(&self, id: NodeID) -> Option<&'a Node> {
106        self.inner_vec.get(id.0 as usize)
107    }
108
109    pub fn get_mut(&mut self, id: NodeID) -> Option<&'a mut Node> {
110        self.inner_vec.get_mut(id.0 as usize)
111    }
112
113    /// Allocates a default Node at in index and returns its index.
114    ///
115    /// To be used when intending to replace the Node at the index
116    /// in conjunction with `alloc_from_id`.
117    pub(crate) fn reserve_id(&mut self) -> NodeID {
118        self.inner_vec.push(Node::default());
119        let old_counter = self.counter;
120        self.counter += 1;
121        NodeID(old_counter)
122    }
123
124    pub fn iter(&self) -> impl Iterator<Item = &Node<'a>> + DoubleEndedIterator<Item = &Node<'a>> {
125        self.inner_vec.iter()
126    }
127
128    pub fn iter_mut(
129        &mut self,
130    ) -> impl Iterator<Item = &mut Node<'a>> + DoubleEndedIterator<Item = &mut Node<'a>> {
131        self.inner_vec.iter_mut()
132    }
133
134    /// Outputs a (somewhat) legible representation of the tree to stdout.
135    pub fn print_tree(&self) {
136        self.inner_vec[0].print_tree(self);
137    }
138
139    /// Returns the [`NodeID`] of the first element in the pool.
140    pub fn root_id(&self) -> NodeID {
141        NodeID(0)
142    }
143
144    /// Removes a [`Node`] from its parents' "children".
145    ///
146    /// This action mimicks the effect of a deletion, but does
147    /// *not* actually deallocate or remove the node from the pool.
148    pub fn delete_node(&mut self, index_id: NodeID) {
149        let par_id = self[index_id].parent.unwrap();
150        let par_node = &mut self[par_id];
151
152        let children = par_node.obj.children_mut().unwrap();
153        let index = children.iter().position(|&x| x == index_id).unwrap();
154        children.remove(index);
155    }
156}
157
158impl<'a> Index<NodeID> for NodePool<'a> {
159    type Output = Node<'a>;
160
161    fn index(&self, index: NodeID) -> &Self::Output {
162        &self.inner_vec[index.0 as usize]
163    }
164}
165
166impl<'a> IndexMut<NodeID> for NodePool<'a> {
167    fn index_mut(&mut self, index: NodeID) -> &mut Self::Output {
168        &mut self.inner_vec[index.0 as usize]
169    }
170}