synctree/
arena.rs

1use std::iter::{ExactSizeIterator, FusedIterator};
2use std::mem;
3use std::num::NonZeroUsize;
4use std::ops::{Index, IndexMut};
5
6use fnv::FnvHashSet;
7use onevec::OneVec;
8
9/// A unique identifier associated with `Arena` element.
10///
11/// It is defined as a wrapper around NonZeroUsize, so size of
12/// `Option<Id>` should be exactly the same as the size of `usize`.
13#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
14pub struct Id(NonZeroUsize);
15
16impl Id {
17    fn inc(self) -> Id {
18        Id(unsafe { NonZeroUsize::new_unchecked(self.0.get() + 1) })
19    }
20
21    fn dec(self) -> Id {
22        Id(unsafe { NonZeroUsize::new_unchecked(self.0.get() - 1) })
23    }
24}
25
26/// A map, providing fast indexing, insertion and removal of elements,
27/// but not preserving element order.
28///
29/// Every element is associated with a unique [id](::arena::Id).
30///
31/// Inside, it uses a fast `FnvHashSet` with free ids.
32///
33/// Every time an element is removed, it's id is being written there,
34/// so there is no need to shift other elements to the empty space.
35#[derive(Debug, Clone)]
36pub struct Arena<T> {
37    elements: OneVec<T>,
38    vacant: FnvHashSet<Id>,
39}
40
41impl<T> Arena<T> {
42    /// Constructs an empty `Arena`.
43    ///
44    /// The arena won't allocate until elements are pushed onto it.
45    pub fn new() -> Arena<T> {
46        Arena {
47            elements: OneVec::new(),
48            vacant: FnvHashSet::default(),
49        }
50    }
51
52    /// Constructs a new `Arena` with specified capacity. If capacity
53    /// is `0`, the arena won't allocate.
54    pub fn with_capacity(capacity: usize) -> Arena<T> {
55        Arena {
56            elements: OneVec::with_capacity(capacity),
57            vacant: FnvHashSet::default(),
58        }
59    }
60
61    /// Returns the capacity of `Arena`.
62    pub fn capacity(&self) -> usize {
63        self.elements.capacity()
64    }
65
66    /// Returns the number of elements in `Arena`.
67    pub fn len(&self) -> usize {
68        self.elements.len() - self.vacant.len()
69    }
70
71    /// Returns true if the arena contains no elements.
72    pub fn is_empty(&self) -> bool {
73        self.len() == 0
74    }
75
76    /// Inserts a new element into the arena and returns
77    /// an id associated with it.
78    pub fn insert(&mut self, element: T) -> Id {
79        let mut iter = self.vacant.iter();
80
81        if let Some(&id) = iter.next() {
82            self.elements[id.0] = element;
83            id
84        } else {
85            self.elements.push(element);
86            Id(unsafe { NonZeroUsize::new_unchecked(self.elements.len()) })
87        }
88    }
89
90    /// Returns true if the arena contains given id.
91    pub fn contains(&self, id: Id) -> bool {
92        !self.vacant.contains(&id) && id.0.get() <= self.elements.len()
93    }
94
95    /// Returns a reference to an element in the arena.
96    ///
97    /// If the specified id does not exists, returns `None`.
98    pub fn get(&self, id: Id) -> Option<&T> {
99        if self.contains(id) {
100            Some(&self.elements[id.0])
101        } else {
102            None
103        }
104    }
105
106    /// Returns a mutable reference to an element in the arena.
107    ///
108    /// If the specified id does not exists, returns `None`.
109    pub fn get_mut(&mut self, id: Id) -> Option<&mut T> {
110        if self.contains(id) {
111            Some(&mut self.elements[id.0])
112        } else {
113            None
114        }
115    }
116
117    /// Removes and returns the element with specified id.
118    ///
119    /// If the id does not exists, returns `None`.
120    pub fn remove(&mut self, id: Id) -> Option<T> {
121        if self.contains(id) {
122            self.vacant.insert(id);
123            let old = mem::replace(&mut self.elements[id.0], unsafe { mem::uninitialized() });
124            Some(old)
125        } else {
126            None
127        }
128    }
129
130    /// Clears the arena.
131    ///
132    /// This method does not change the capacity of the arena.
133    pub fn clear(&mut self) {
134        self.elements.clear();
135        self.vacant.clear();
136    }
137
138    /// Returns an iterator over the arena.
139    pub fn iter<'a>(&'a self) -> Iter<'a, T> {
140        Iter {
141            arena: &self,
142            first: Id(unsafe { NonZeroUsize::new_unchecked(1) }),
143            last: Id(unsafe { NonZeroUsize::new_unchecked(self.elements.len() + 1) }),
144            len: self.len(),
145        }
146    }
147
148    /// Returns a mutable iterator over the arena.
149    pub fn iter_mut<'a>(&'a mut self) -> IterMut<'a, T> {
150        let last = Id(unsafe { NonZeroUsize::new_unchecked(self.elements.len() + 1) });
151        let len = self.len();
152
153        IterMut {
154            last,
155            len,
156            arena: self,
157            first: Id(unsafe { NonZeroUsize::new_unchecked(1) }),
158        }
159    }
160}
161
162impl<T> Index<Id> for Arena<T> {
163    type Output = T;
164
165    fn index(&self, index: Id) -> &T {
166        self.get(index).unwrap()
167    }
168}
169
170impl<T> IndexMut<Id> for Arena<T> {
171    fn index_mut(&mut self, index: Id) -> &mut T {
172        self.get_mut(index).unwrap()
173    }
174}
175
176/// An immutable iterator over the arena.
177pub struct Iter<'a, T: 'a> {
178    arena: &'a Arena<T>,
179    first: Id,
180    last: Id,
181    len: usize,
182}
183
184impl<'a, T: 'a> Iterator for Iter<'a, T> {
185    type Item = &'a T;
186
187    fn next(&mut self) -> Option<Self::Item> {
188        if self.first.0.get() >= self.last.0.get() {
189            return None;
190        }
191
192        let mut cur = self.first;
193
194        while self.arena.vacant.contains(&self.first) {
195            cur = self.first.inc();
196        }
197
198        self.first = cur.inc();
199
200        self.len -= 1;
201        self.arena.get(cur)
202    }
203
204    fn size_hint(&self) -> (usize, Option<usize>) {
205        (self.len, Some(self.len))
206    }
207}
208
209impl<'a, T: 'a> DoubleEndedIterator for Iter<'a, T> {
210    fn next_back(&mut self) -> Option<Self::Item> {
211        if self.last.0.get() <= self.first.0.get() {
212            return None;
213        }
214
215        while self.arena.vacant.contains(&self.last.dec()) {
216            if let Some(id) = NonZeroUsize::new(self.last.0.get() - 1).map(|v| Id(v)) {
217                self.last = id;
218            } else {
219                return None;
220            }
221        }
222
223        self.last = self.last.dec();
224
225        self.len -= 1;
226        self.arena.get(self.last)
227    }
228}
229
230impl<'a, T: 'a> FusedIterator for Iter<'a, T> {}
231impl<'a, T: 'a> ExactSizeIterator for Iter<'a, T> {}
232
233/// A mutable iterator over the arena.
234pub struct IterMut<'a, T: 'a> {
235    arena: &'a mut Arena<T>,
236    first: Id,
237    last: Id,
238    len: usize,
239}
240
241impl<'a, T: 'a> Iterator for IterMut<'a, T> {
242    type Item = &'a mut T;
243
244    fn next(&mut self) -> Option<Self::Item> {
245        if self.first.0.get() >= self.last.0.get() {
246            return None;
247        }
248
249        let mut cur = self.first;
250
251        while self.arena.vacant.contains(&self.first) {
252            cur = self.first.inc();
253        }
254
255        self.first = cur.inc();
256
257        self.len -= 1;
258        self.arena
259            .get_mut(cur)
260            .map(|v| unsafe { &mut *(v as *mut T) })
261    }
262
263    fn size_hint(&self) -> (usize, Option<usize>) {
264        (self.len, Some(self.len))
265    }
266}
267
268impl<'a, T: 'a> DoubleEndedIterator for IterMut<'a, T> {
269    fn next_back(&mut self) -> Option<Self::Item> {
270        if self.last.0.get() <= self.first.0.get() {
271            return None;
272        }
273
274        while self.arena.vacant.contains(&self.last.dec()) {
275            if let Some(id) = NonZeroUsize::new(self.last.0.get() - 1).map(|v| Id(v)) {
276                self.last = id;
277            } else {
278                return None;
279            }
280        }
281
282        self.last = self.last.dec();
283
284        self.len -= 1;
285        self.arena
286            .get_mut(self.last)
287            .map(|v| unsafe { &mut *(v as *mut T) })
288    }
289}
290
291impl<'a, T: 'a> FusedIterator for IterMut<'a, T> {}
292impl<'a, T: 'a> ExactSizeIterator for IterMut<'a, T> {}
293
294#[cfg(test)]
295mod tests {
296    use super::*;
297
298    #[test]
299    fn create() {
300        let a = Arena::<i32>::new();
301        assert_eq!(a.capacity(), 0);
302
303        let a = Arena::<i32>::with_capacity(10);
304        assert_eq!(a.capacity(), 10);
305        assert!(a.is_empty());
306    }
307
308    #[test]
309    fn simple() {
310        let mut a = Arena::new();
311        let id = a.insert(10);
312
313        assert!(a.contains(id));
314        assert!(a.capacity() >= 1);
315        assert_eq!(a.len(), 1);
316        assert!(!a.is_empty());
317        assert_eq!(*a.get(id).unwrap(), 10);
318
319        *a.get_mut(id).unwrap() = 20;
320
321        assert_eq!(a[id], 20);
322
323        a[id] = 30;
324
325        assert_eq!(a.remove(id).unwrap(), 30);
326        assert!(a.remove(id).is_none());
327        assert!(a.is_empty());
328    }
329
330    #[test]
331    fn iter() {
332        let mut a = Arena::new();
333        a.insert(1);
334        a.insert(2);
335        a.insert(3);
336        a.insert(4);
337
338        let mut iter = a.iter();
339        assert_eq!(iter.len(), 4);
340        assert_eq!(*iter.next().unwrap(), 1);
341        assert_eq!(iter.len(), 3);
342        assert_eq!(*iter.next_back().unwrap(), 4);
343        assert_eq!(iter.len(), 2);
344        assert_eq!(*iter.next().unwrap(), 2);
345        assert_eq!(iter.len(), 1);
346        assert_eq!(*iter.next_back().unwrap(), 3);
347        assert_eq!(iter.len(), 0);
348        assert!(iter.next().is_none());
349        assert!(iter.next().is_none());
350        assert!(iter.next_back().is_none());
351        assert!(iter.next_back().is_none());
352    }
353
354    #[test]
355    fn iter_mut() {
356        let mut a = Arena::new();
357        a.insert(1);
358        a.insert(2);
359        a.insert(3);
360        a.insert(4);
361
362        let mut iter = a.iter_mut();
363        assert_eq!(iter.len(), 4);
364        assert_eq!(*iter.next().unwrap(), 1);
365        assert_eq!(iter.len(), 3);
366        assert_eq!(*iter.next_back().unwrap(), 4);
367        assert_eq!(iter.len(), 2);
368        assert_eq!(*iter.next().unwrap(), 2);
369        assert_eq!(iter.len(), 1);
370        assert_eq!(*iter.next_back().unwrap(), 3);
371        assert_eq!(iter.len(), 0);
372        assert!(iter.next().is_none());
373        assert!(iter.next().is_none());
374        assert!(iter.next_back().is_none());
375        assert!(iter.next_back().is_none());
376    }
377}