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

use std::ops::{Index, IndexMut};

/// An AST node ID for lookups in the arena
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct NodeId(pub usize);

/// Arenas are efficient ways of storing trees: All nodes are actually on a
/// flat surface, but have IDs to other nodes.
#[derive(Debug)]
pub enum Arena<'a, T: 'a> {
    Owned(Vec<Option<T>>),
    Borrowed(&'a mut Vec<Option<T>>)
}
impl<'a, T> Default for Arena<'a, T> {
    fn default() -> Self {
        Arena::Owned(Vec::new())
    }
}
impl<'a, T> Arena<'a, T> {
    /// Create a new arena instance
    pub fn new() -> Self {
        Self::default()
    }
    /// Create a new arena instance that reuses an existing vector
    pub fn with_arena(arena: &'a mut Vec<Option<T>>) -> Self {
        Arena::Borrowed(arena)
    }
    /// Create a new arena reference that shares the inner vector. This
    /// reference will make the existing instance invalid for use while the
    /// reference is active.
    pub fn reference<'b>(&'b mut self) -> Arena<'b, T> {
        Arena::with_arena(self.get_mut())
    }

    /// Return the arena instance
    ///
    /// # Panics
    /// Panics if this arena instance was created with Borrowed
    pub fn into_inner(self) -> Vec<Option<T>> {
        match self {
            Arena::Owned(inner) => inner,
            Arena::Borrowed(_) => panic!("can't move out of borrowed arena")
        }
    }
    /// Return a slice pointing to the inner arena representation
    pub fn get_ref(&self) -> &[Option<T>] {
        match self {
            Arena::Owned(inner) => inner,
            Arena::Borrowed(inner) => inner
        }
    }
    fn get_mut(&mut self) -> &mut Vec<Option<T>> {
        match self {
            Arena::Owned(inner) => inner,
            Arena::Borrowed(inner) => inner
        }
    }
    /// Return an iterator over the nodes in this arena
    pub fn iter(&self) -> ArenaIter<T> {
        self.into_iter()
    }
    /// Place an element into the arena
    pub fn insert(&mut self, elem: T) -> NodeId {
        let inner = self.get_mut();
        inner.push(Some(elem));
        NodeId(inner.len() - 1)
    }
    /// Move out of an element in the arena. This will make following indexing
    /// calls on that index panic. This will not deallocate any space. This is
    /// due to the fact that all indexes need to stay the same. We could use a
    /// different data type instead of a tree to prevent this, but that instead
    /// means lookup will be slightly slower.
    pub fn take(&mut self, index: NodeId) -> T {
        self.get_mut()[index.0].take().expect("this index has been taken")
    }
}
impl<'a, T> Index<NodeId> for Arena<'a, T> {
    type Output = T;

    fn index(&self, index: NodeId) -> &Self::Output {
        &self.get_ref()[index.0].as_ref().expect("this index has been taken")
    }
}
impl<'a, T> IndexMut<NodeId> for Arena<'a, T> {
    fn index_mut(&mut self, index: NodeId) -> &mut Self::Output {
        self.get_mut()[index.0].as_mut().expect("this index has been taken")
    }
}
impl<'r, 'a, T> IntoIterator for &'r Arena<'a, T> {
    type Item = &'r T;
    type IntoIter = ArenaIter<'r, T>;

    fn into_iter(self) -> Self::IntoIter {
        ArenaIter {
            slice: self.get_ref(),
            index: 0
        }
    }
}

impl<'a, T> Clone for Arena<'a, T> where T: Clone {
    fn clone(&self) -> Self {
        match self {
            Arena::Owned(inner) => Arena::Owned(inner.to_vec()),
            Arena::Borrowed(inner) => Arena::Owned(inner.to_vec())
        }
    }
}

/// An iterator over an arena, created with `iter`
pub struct ArenaIter<'a, T: 'a> {
    slice: &'a [Option<T>],
    index: usize
}
impl<'a, T> Iterator for ArenaIter<'a, T> {
    type Item = &'a T;

    fn next(&mut self) -> Option<Self::Item> {
        while let Some(None) = self.slice.get(self.index) {
            self.index += 1;
        }
        let item = self.slice.get(self.index).map(|item| item.as_ref().unwrap());
        self.index += 1;
        item
    }
}