1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165

use std::fmt::{Debug, Display};
use std::ops::{Index, IndexMut};

use crate::types::{Expr, Node};

#[derive(Clone, Copy, Hash, PartialEq, PartialOrd, Ord, Eq)]
/// Identifier for [`Node`]s in a [`NodePool`].
///
/// NodeIDs are guaranteed to be unique to each node since they are assigned
/// sequentially and cannot re-used.
pub struct NodeID(u32);

impl Display for NodeID {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_fmt(format_args!("{}", self.0))
    }
}

impl std::fmt::Debug for NodeID {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_fmt(format_args!("{}", self.0))
    }
}

/// Arena-based container storing [`Node`]s
///
/// A [`NodePool`] is essentially just a [`Vec<Node>`] that is indexed by [`NodeID`]s.
/// Each parsing construct stores [`NodeID`]s which refer to indices within the [`NodePool`].
/// It acts like a flattened version of the AST where each node can be easily iterated over
/// and tracked.
///
/// [`NodePool::iter`] and [`NodePool::iter_mut`] can be used for iterating over the contents of
/// the pool.
///
/// # Safety
///
/// Removing or inserting an element at a position other than the end might invalidate every other NodeID
/// and therefore the AST. Removing and de-allocating an individual [`Node`] is not possible due to
/// this effect. However, the result can be feigned in the tree by "deleting" the node from its parent.
/// See [`NodePool::delete_node`].
///
///
#[derive(Debug)]
pub struct NodePool<'a> {
    pub inner_vec: Vec<Node<'a>>,
    pub counter: u32,
}

impl<'a> NodePool<'a> {
    pub(crate) fn new() -> Self {
        Self {
            inner_vec: Vec::new(),
            /// The next free index in the pool.
            counter: 0,
        }
    }

    pub(crate) fn alloc<T>(
        &mut self,
        obj: T,
        start: usize,
        end: usize,
        parent: Option<NodeID>,
    ) -> NodeID
    where
        Expr<'a>: From<T>,
    {
        let prev_id = self.counter;
        self.inner_vec.push(Node::new(obj, start, end, parent));
        self.counter += 1;
        NodeID(prev_id)
    }

    /// Allocates a node in the pool at a given location.
    ///
    /// Returns the index that was allocated.
    ///
    /// Works well with [`NodePool::reserve_id`].
    ///
    /// # Safety:
    ///
    /// Must refer to an ID that already exists in the pool.
    /// Will panic at runtime otherwise.
    pub(crate) fn alloc_with_id<T>(
        &mut self,
        obj: T,
        start: usize,
        end: usize,
        parent: Option<NodeID>,
        target_id: NodeID,
    ) -> NodeID
    where
        Expr<'a>: From<T>,
    {
        self.inner_vec[target_id.0 as usize] = Node::new(obj, start, end, parent);

        target_id
    }

    pub fn get(&self, id: NodeID) -> Option<&'a Node> {
        self.inner_vec.get(id.0 as usize)
    }

    pub fn get_mut(&mut self, id: NodeID) -> Option<&'a mut Node> {
        self.inner_vec.get_mut(id.0 as usize)
    }

    /// Allocates a default Node at in index and returns its index.
    ///
    /// To be used when intending to replace the Node at the index
    /// in conjunction with `alloc_from_id`.
    pub(crate) fn reserve_id(&mut self) -> NodeID {
        self.inner_vec.push(Node::default());
        let old_counter = self.counter;
        self.counter += 1;
        NodeID(old_counter)
    }

    pub fn iter(&self) -> impl Iterator<Item = &Node<'a>> + DoubleEndedIterator<Item = &Node<'a>> {
        self.inner_vec.iter()
    }

    pub fn iter_mut(
        &mut self,
    ) -> impl Iterator<Item = &mut Node<'a>> + DoubleEndedIterator<Item = &mut Node<'a>> {
        self.inner_vec.iter_mut()
    }

    /// Outputs a (somewhat) legible representation of the tree to stdout.
    pub fn print_tree(&self) {
        self.inner_vec[0].print_tree(self);
    }

    /// Returns the [`NodeID`] of the first element in the pool.
    pub fn root_id(&self) -> NodeID {
        NodeID(0)
    }

    /// Removes a [`Node`] from its parents' "children".
    ///
    /// This action mimicks the effect of a deletion, but does
    /// *not* actually deallocate or remove the node from the pool.
    pub fn delete_node(&mut self, index_id: NodeID) {
        let par_id = self[index_id].parent.unwrap();
        let par_node = &mut self[par_id];

        let children = par_node.obj.children_mut().unwrap();
        let index = children.iter().position(|&x| x == index_id).unwrap();
        children.remove(index);
    }
}

impl<'a> Index<NodeID> for NodePool<'a> {
    type Output = Node<'a>;

    fn index(&self, index: NodeID) -> &Self::Output {
        &self.inner_vec[index.0 as usize]
    }
}

impl<'a> IndexMut<NodeID> for NodePool<'a> {
    fn index_mut(&mut self, index: NodeID) -> &mut Self::Output {
        &mut self.inner_vec[index.0 as usize]
    }
}