indextree/
arena.rs

1//! Arena.
2
3#[cfg(not(feature = "std"))]
4use alloc::vec::Vec;
5
6#[cfg(not(feature = "std"))]
7use core::{
8    mem,
9    num::NonZeroUsize,
10    ops::{Index, IndexMut},
11    slice,
12};
13
14#[cfg(feature = "par_iter")]
15use rayon::prelude::*;
16
17#[cfg(feature = "deser")]
18use serde::{Deserialize, Serialize};
19
20#[cfg(feature = "std")]
21use std::{
22    mem,
23    num::NonZeroUsize,
24    ops::{Index, IndexMut},
25    slice,
26};
27
28use crate::{node::NodeData, Node, NodeId};
29
30#[derive(PartialEq, Eq, Clone, Debug)]
31#[cfg_attr(feature = "deser", derive(Deserialize, Serialize))]
32/// An `Arena` structure containing certain [`Node`]s.
33///
34/// [`Node`]: struct.Node.html
35pub struct Arena<T> {
36    nodes: Vec<Node<T>>,
37    first_free_slot: Option<usize>,
38    last_free_slot: Option<usize>,
39}
40
41impl<T> Arena<T> {
42    /// Creates a new empty `Arena`.
43    pub fn new() -> Arena<T> {
44        Self::default()
45    }
46
47    /// Creates a new empty `Arena` with enough capacity to store `n` nodes.
48    pub fn with_capacity(n: usize) -> Self {
49        Self {
50            nodes: Vec::with_capacity(n),
51            first_free_slot: None,
52            last_free_slot: None,
53        }
54    }
55
56    /// Returns the number of nodes the arena can hold without reallocating.
57    pub fn capacity(&self) -> usize {
58        self.nodes.capacity()
59    }
60
61    /// Reserves capacity for `additional` more nodes to be inserted.
62    ///
63    /// The arena may reserve more space to avoid frequent reallocations.
64    ///
65    /// # Panics
66    ///
67    /// Panics if the new capacity exceeds isize::MAX bytes.
68    pub fn reserve(&mut self, additional: usize) {
69        self.nodes.reserve(additional);
70    }
71
72    /// Retrieves the `NodeId` corresponding to a `Node` in the `Arena`.
73    ///
74    /// # Examples
75    ///
76    /// ```
77    /// # use indextree::Arena;
78    /// let mut arena = Arena::new();
79    /// let foo = arena.new_node("foo");
80    /// let node = arena.get(foo).unwrap();
81    ///
82    /// let node_id = arena.get_node_id(node).unwrap();
83    /// assert_eq!(*arena[node_id].get(), "foo");
84    /// ```
85    pub fn get_node_id(&self, node: &Node<T>) -> Option<NodeId> {
86        let nodes_range = self.nodes.as_ptr_range();
87        let p = node as *const Node<T>;
88
89        if !nodes_range.contains(&p) {
90            return None;
91        }
92
93        let node_index = (p as usize - nodes_range.start as usize) / mem::size_of::<Node<T>>();
94        let node_id = NonZeroUsize::new(node_index.wrapping_add(1))?;
95
96        Some(NodeId::from_non_zero_usize(
97            node_id,
98            self.nodes[node_index].stamp,
99        ))
100    }
101
102    /// Retrieves the `NodeId` corresponding to the `Node` at `index` in the `Arena`, if it exists.
103    ///
104    /// Note: We use 1 based indexing, so the first element is at `1` and not `0`.
105    ///
106    /// # Examples
107    ///
108    /// ```
109    /// # use indextree::Arena;
110    /// # use std::num::NonZeroUsize;
111    /// let mut arena = Arena::new();
112    /// let foo = arena.new_node("foo");
113    /// let node = arena.get(foo).unwrap();
114    /// let index: NonZeroUsize = foo.into();
115    ///
116    /// let new_foo = arena.get_node_id_at(index).unwrap();
117    /// assert_eq!(foo, new_foo);
118    ///
119    /// foo.remove(&mut arena);
120    /// let new_foo = arena.get_node_id_at(index);
121    /// assert!(new_foo.is_none(), "must be none if the node at the index doesn't exist");
122    /// ```
123    pub fn get_node_id_at(&self, index: NonZeroUsize) -> Option<NodeId> {
124        let index0 = index.get() - 1; // we use 1 based indexing.
125        self.nodes
126            .get(index0)
127            .filter(|n| !n.is_removed())
128            .map(|node| NodeId::from_non_zero_usize(index, node.stamp))
129    }
130
131    /// Creates a new node from its associated data.
132    ///
133    /// # Panics
134    ///
135    /// Panics if the arena already has `usize::max_value()` nodes.
136    ///
137    /// # Examples
138    ///
139    /// ```
140    /// # use indextree::Arena;
141    /// let mut arena = Arena::new();
142    /// let foo = arena.new_node("foo");
143    ///
144    /// assert_eq!(*arena[foo].get(), "foo");
145    /// ```
146    pub fn new_node(&mut self, data: T) -> NodeId {
147        let (index, stamp) = if let Some(index) = self.pop_front_free_node() {
148            let node = &mut self.nodes[index];
149            node.reuse(data);
150            (index, node.stamp)
151        } else {
152            let index = self.nodes.len();
153            let node = Node::new(data);
154            let stamp = node.stamp;
155            self.nodes.push(node);
156            (index, stamp)
157        };
158        let next_index1 =
159            NonZeroUsize::new(index.wrapping_add(1)).expect("Too many nodes in the arena");
160        NodeId::from_non_zero_usize(next_index1, stamp)
161    }
162
163    /// Counts the number of nodes in arena and returns it.
164    ///
165    /// # Examples
166    ///
167    /// ```
168    /// # use indextree::Arena;
169    /// let mut arena = Arena::new();
170    /// let foo = arena.new_node("foo");
171    /// let _bar = arena.new_node("bar");
172    /// assert_eq!(arena.count(), 2);
173    ///
174    /// foo.remove(&mut arena);
175    /// assert_eq!(arena.count(), 2);
176    /// ```
177    pub fn count(&self) -> usize {
178        self.nodes.len()
179    }
180
181    /// Returns `true` if arena has no nodes, `false` otherwise.
182    ///
183    /// # Examples
184    ///
185    /// ```
186    /// # use indextree::Arena;
187    /// let mut arena = Arena::new();
188    /// assert!(arena.is_empty());
189    ///
190    /// let foo = arena.new_node("foo");
191    /// assert!(!arena.is_empty());
192    ///
193    /// foo.remove(&mut arena);
194    /// assert!(!arena.is_empty());
195    /// ```
196    pub fn is_empty(&self) -> bool {
197        self.count() == 0
198    }
199
200    /// Returns a reference to the node with the given id if in the arena.
201    ///
202    /// Returns `None` if not available.
203    ///
204    /// # Examples
205    ///
206    /// ```
207    /// # use indextree::{Arena, NodeId};
208    /// let mut arena = Arena::new();
209    /// let foo = arena.new_node("foo");
210    /// assert_eq!(arena.get(foo).map(|node| *node.get()), Some("foo"));
211    /// ```
212    ///
213    /// Note that this does not check whether the given node ID is created by
214    /// the arena.
215    ///
216    /// ```
217    /// # use indextree::Arena;
218    /// let mut arena = Arena::new();
219    /// let foo = arena.new_node("foo");
220    /// let bar = arena.new_node("bar");
221    /// assert_eq!(arena.get(foo).map(|node| *node.get()), Some("foo"));
222    ///
223    /// let mut another_arena = Arena::new();
224    /// let _ = another_arena.new_node("Another arena");
225    /// assert_eq!(another_arena.get(foo).map(|node| *node.get()), Some("Another arena"));
226    /// assert!(another_arena.get(bar).is_none());
227    /// ```
228    pub fn get(&self, id: NodeId) -> Option<&Node<T>> {
229        self.nodes.get(id.index0())
230    }
231
232    /// Returns a mutable reference to the node with the given id if in the
233    /// arena.
234    ///
235    /// Returns `None` if not available.
236    ///
237    /// # Examples
238    ///
239    /// ```
240    /// # use indextree::{Arena, NodeId};
241    /// let mut arena = Arena::new();
242    /// let foo = arena.new_node("foo");
243    /// assert_eq!(arena.get(foo).map(|node| *node.get()), Some("foo"));
244    ///
245    /// *arena.get_mut(foo).expect("The `foo` node exists").get_mut() = "FOO!";
246    /// assert_eq!(arena.get(foo).map(|node| *node.get()), Some("FOO!"));
247    /// ```
248    pub fn get_mut(&mut self, id: NodeId) -> Option<&mut Node<T>> {
249        self.nodes.get_mut(id.index0())
250    }
251
252    /// Returns an iterator of all nodes in the arena in storage-order.
253    ///
254    /// Note that this iterator returns also removed elements, which can be
255    /// tested with the [`is_removed()`] method on the node.
256    ///
257    /// # Examples
258    ///
259    /// ```
260    /// # use indextree::Arena;
261    /// let mut arena = Arena::new();
262    /// let _foo = arena.new_node("foo");
263    /// let _bar = arena.new_node("bar");
264    ///
265    /// let mut iter = arena.iter();
266    /// assert_eq!(iter.next().map(|node| *node.get()), Some("foo"));
267    /// assert_eq!(iter.next().map(|node| *node.get()), Some("bar"));
268    /// assert_eq!(iter.next().map(|node| *node.get()), None);
269    /// ```
270    ///
271    /// ```
272    /// # use indextree::Arena;
273    /// let mut arena = Arena::new();
274    /// let _foo = arena.new_node("foo");
275    /// let bar = arena.new_node("bar");
276    /// bar.remove(&mut arena);
277    ///
278    /// let mut iter = arena.iter();
279    /// assert_eq!(iter.next().map(|node| (*node.get(), node.is_removed())), Some(("foo", false)));
280    /// assert_eq!(iter.next().map_or(false, |node| node.is_removed()), true);
281    /// assert_eq!(iter.next().map(|node| (*node.get(), node.is_removed())), None);
282    /// ```
283    ///
284    /// [`is_removed()`]: struct.Node.html#method.is_removed
285    pub fn iter(&self) -> slice::Iter<Node<T>> {
286        self.nodes.iter()
287    }
288
289    /// Returns a mutable iterator of all nodes in the arena in storage-order.
290    ///
291    /// Note that this iterator returns also removed elements, which can be
292    /// tested with the [`is_removed()`] method on the node.
293    ///
294    /// # Example
295    ///
296    /// ```
297    /// # use indextree::Arena;
298    /// let arena: &mut Arena<i64> = &mut Arena::new();
299    /// let a = arena.new_node(1);
300    /// let b = arena.new_node(2);
301    /// assert!(a.checked_append(b, arena).is_ok());
302    ///
303    /// for node in arena.iter_mut() {
304    ///     let data = node.get_mut();
305    ///     *data = data.wrapping_add(4);
306    /// }
307    ///
308    /// let node_refs = arena.iter().map(|i| i.get().clone()).collect::<Vec<_>>();
309    /// assert_eq!(node_refs, vec![5, 6]);
310    /// ```
311    /// [`is_removed()`]: struct.Node.html#method.is_removed
312    pub fn iter_mut(&mut self) -> slice::IterMut<Node<T>> {
313        self.nodes.iter_mut()
314    }
315
316    /// Clears all the nodes in the arena, but retains its allocated capacity.
317    ///
318    /// Note that this does not marks all nodes as removed, but completely
319    /// removes them from the arena storage, thus invalidating all the node ids
320    /// that were previously created.
321    ///
322    /// Any attempt to call the [`is_removed()`] method on the node id will
323    /// result in panic behavior.
324    ///
325    /// [`is_removed()`]: struct.NodeId.html#method.is_removed
326    pub fn clear(&mut self) {
327        self.nodes.clear();
328        self.first_free_slot = None;
329        self.last_free_slot = None;
330    }
331
332    /// Returns a slice of the inner nodes collection.
333    ///
334    /// Note that this **does not** return root elements, it simply
335    /// returns a slice into the internal representation of the arena.
336    pub fn as_slice(&self) -> &[Node<T>] {
337        self.nodes.as_slice()
338    }
339
340    pub(crate) fn free_node(&mut self, id: NodeId) {
341        let node = &mut self[id];
342        node.data = NodeData::NextFree(None);
343        node.stamp.as_removed();
344        let stamp = node.stamp;
345        if stamp.reuseable() {
346            if let Some(index) = self.last_free_slot {
347                let new_last = id.index0();
348                self.nodes[index].data = NodeData::NextFree(Some(new_last));
349                self.last_free_slot = Some(new_last);
350            } else {
351                debug_assert!(self.first_free_slot.is_none());
352                debug_assert!(self.last_free_slot.is_none());
353                self.first_free_slot = Some(id.index0());
354                self.last_free_slot = Some(id.index0());
355            }
356        }
357    }
358
359    fn pop_front_free_node(&mut self) -> Option<usize> {
360        let first = self.first_free_slot.take();
361        if let Some(index) = first {
362            if let NodeData::NextFree(next_free) = self.nodes[index].data {
363                self.first_free_slot = next_free;
364            } else {
365                unreachable!("A data node consider as a freed node");
366            }
367            if self.first_free_slot.is_none() {
368                self.last_free_slot = None;
369            }
370        }
371
372        first
373    }
374}
375
376#[cfg(feature = "par_iter")]
377impl<T: Sync> Arena<T> {
378    /// Returns an parallel iterator over the whole arena.
379    ///
380    /// Note that this iterator returns also removed elements, which can be
381    /// tested with the [`is_removed()`] method on the node.
382    ///
383    /// [`is_removed()`]: struct.Node.html#method.is_removed
384    pub fn par_iter(&self) -> rayon::slice::Iter<'_, Node<T>> {
385        self.nodes.par_iter()
386    }
387}
388
389impl<T> Default for Arena<T> {
390    fn default() -> Self {
391        Self {
392            nodes: Vec::new(),
393            first_free_slot: None,
394            last_free_slot: None,
395        }
396    }
397}
398
399impl<T> Index<NodeId> for Arena<T> {
400    type Output = Node<T>;
401
402    fn index(&self, node: NodeId) -> &Node<T> {
403        &self.nodes[node.index0()]
404    }
405}
406
407impl<T> IndexMut<NodeId> for Arena<T> {
408    fn index_mut(&mut self, node: NodeId) -> &mut Node<T> {
409        &mut self.nodes[node.index0()]
410    }
411}
412
413#[test]
414fn reuse_node() {
415    let mut arena = Arena::new();
416    let n1_id = arena.new_node("1");
417    let n2_id = arena.new_node("2");
418    let n3_id = arena.new_node("3");
419    n1_id.remove(&mut arena);
420    n2_id.remove(&mut arena);
421    n3_id.remove(&mut arena);
422    let n1_id = arena.new_node("1");
423    let n2_id = arena.new_node("2");
424    let n3_id = arena.new_node("3");
425    assert_eq!(n1_id.index0(), 0);
426    assert_eq!(n2_id.index0(), 1);
427    assert_eq!(n3_id.index0(), 2);
428    assert_eq!(arena.nodes.len(), 3);
429}
430
431#[test]
432fn conserve_capacity() {
433    let mut arena = Arena::with_capacity(5);
434    let cap = arena.capacity();
435    assert!(cap >= 5);
436    for i in 0..cap {
437        arena.new_node(i);
438    }
439    arena.clear();
440    assert!(arena.is_empty());
441    let n1_id = arena.new_node(1);
442    let n2_id = arena.new_node(2);
443    let n3_id = arena.new_node(3);
444    assert_eq!(n1_id.index0(), 0);
445    assert_eq!(n2_id.index0(), 1);
446    assert_eq!(n3_id.index0(), 2);
447    assert_eq!(arena.count(), 3);
448    assert_eq!(arena.capacity(), cap);
449}