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}