git.lucas.co / cce-ui
GPU-accelerated UI toolkit (Vulkan)
git clone https://git.lucas.co/cce-ui.git

src/scene/arena.rs (21.7K)

  1 //! Generational node arena — the ownership spine of the rebuilt cce-ui core.
  2 //!
  3 //! This is Phase 1 of the core rebuild (see `docs/rfc-core-rebuild.md`). It replaces the old
  4 //! model where the widget tree was smeared across three parallel stores kept in sync by hand
  5 //! (`root plate container.children: Vec<*mut dyn WidgetHost>`, `UiContext.layout_tree`, and
  6 //! `UiContext.widget_registry`) and traversed through raw `*mut dyn WidgetHost` pointers that
  7 //! `Drop` did not fully clear.
  8 //!
  9 //! Here there is exactly **one** store. Every node lives in the [`Arena`], addressed by a
 10 //! [`NodeId`] that carries a generation. When a node is removed its slot's generation is bumped,
 11 //! so any [`NodeId`] still pointing at the old occupant reads back as [`None`] instead of
 12 //! dereferencing freed memory. Use-after-free becomes a missed lookup, not undefined behavior —
 13 //! the single property that dissolves the dangling-pointer bug class.
 14 //!
 15 //! The arena is a **forest**: a freshly [`insert`](Arena::insert)ed node is a detached root
 16 //! (`parent == None`); [`append_child`](Arena::append_child) links nodes into trees. Children are
 17 //! stored as `Vec<NodeId>` (indices, not pointers), so traversal never aliases a `&mut`, which
 18 //! keeps the whole thing safe and borrow-checker-friendly without `unsafe`.
 19 //!
 20 //! [`Node<T>`] is generic over its payload for now. In later phases the payload grows into the
 21 //! rich per-node record from the RFC (widget + style + computed layout + animation + dirty
 22 //! flags); nothing about the identity/ownership model below changes when it does.
 23 
 24 use std::num::NonZeroU32;
 25 
 26 /// A stable handle to a node in an [`Arena`].
 27 ///
 28 /// Carries both a slot index and a generation. The generation makes the handle *safe across
 29 /// removal*: once the node it referred to is removed (and its slot possibly reused for a
 30 /// different node), every lookup with this id returns [`None`]. `NodeId` is `Copy` and cheap to
 31 /// pass around; `Option<NodeId>` is the same size as `NodeId` thanks to the `NonZero` generation.
 32 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 33 pub struct NodeId {
 34     index: u32,
 35     generation: NonZeroU32,
 36 }
 37 
 38 impl NodeId {
 39     /// Opaque index into the arena's backing storage. Exposed only for debugging/telemetry —
 40     /// do not use it to bypass generational checks.
 41     #[inline]
 42     pub fn slot_index(self) -> u32 {
 43         self.index
 44     }
 45 }
 46 
 47 /// A node in the arena: a payload plus its tree links. Links are only mutated through the
 48 /// [`Arena`] so the parent/child relationship stays symmetric.
 49 #[derive(Debug, Clone)]
 50 pub struct Node<T> {
 51     parent: Option<NodeId>,
 52     children: Vec<NodeId>,
 53     value: T,
 54 }
 55 
 56 impl<T> Node<T> {
 57     /// This node's parent, or `None` if it is a detached root.
 58     #[inline]
 59     pub fn parent(&self) -> Option<NodeId> {
 60         self.parent
 61     }
 62 
 63     /// This node's direct children, in order.
 64     #[inline]
 65     pub fn children(&self) -> &[NodeId] {
 66         &self.children
 67     }
 68 
 69     /// Shared access to the payload.
 70     #[inline]
 71     pub fn value(&self) -> &T {
 72         &self.value
 73     }
 74 
 75     /// Mutable access to the payload. Tree links are intentionally not reachable here — use the
 76     /// [`Arena`] methods so both ends of every edge stay consistent.
 77     #[inline]
 78     pub fn value_mut(&mut self) -> &mut T {
 79         &mut self.value
 80     }
 81 }
 82 
 83 enum Slot<T> {
 84     /// A live node. `generation` matches the [`NodeId`] handed out for it.
 85     Occupied { generation: NonZeroU32, node: Node<T> },
 86     /// A free slot. `generation` is the generation the *next* occupant will receive, and
 87     /// `next_free` chains the free list.
 88     Vacant { generation: NonZeroU32, next_free: Option<u32> },
 89 }
 90 
 91 impl<T> Slot<T> {
 92     #[inline]
 93     fn generation(&self) -> NonZeroU32 {
 94         match self {
 95             Slot::Occupied { generation, .. } | Slot::Vacant { generation, .. } => *generation,
 96         }
 97     }
 98 }
 99 
100 #[inline]
101 fn bump(generation: NonZeroU32) -> NonZeroU32 {
102     // Wrap while skipping 0 (which `NonZeroU32` cannot hold). A wrap only aliases a generation
103     // after 2^32-1 reuses of the same slot, which no real UI session approaches.
104     let next = generation.get().wrapping_add(1);
105     NonZeroU32::new(if next == 0 { 1 } else { next }).unwrap()
106 }
107 
108 const FIRST_GENERATION: NonZeroU32 = match NonZeroU32::new(1) {
109     Some(g) => g,
110     None => unreachable!(),
111 };
112 
113 /// A generational forest of [`Node<T>`]. See the module docs for the design.
114 pub struct Arena<T> {
115     slots: Vec<Slot<T>>,
116     free_head: Option<u32>,
117     len: usize,
118 }
119 
120 impl<T> Default for Arena<T> {
121     fn default() -> Self {
122         Self::new()
123     }
124 }
125 
126 impl<T> Arena<T> {
127     /// An empty arena.
128     pub fn new() -> Self {
129         Arena { slots: Vec::new(), free_head: None, len: 0 }
130     }
131 
132     /// An empty arena with room for `capacity` nodes before reallocating.
133     pub fn with_capacity(capacity: usize) -> Self {
134         Arena { slots: Vec::with_capacity(capacity), free_head: None, len: 0 }
135     }
136 
137     /// Number of live nodes.
138     #[inline]
139     pub fn len(&self) -> usize {
140         self.len
141     }
142 
143     /// Whether there are no live nodes.
144     #[inline]
145     pub fn is_empty(&self) -> bool {
146         self.len == 0
147     }
148 
149     /// Whether `id` still refers to a live node (generation matches).
150     #[inline]
151     pub fn contains(&self, id: NodeId) -> bool {
152         matches!(self.slots.get(id.index as usize),
153             Some(Slot::Occupied { generation, .. }) if *generation == id.generation)
154     }
155 
156     /// Insert a detached node (a new root) and return its id.
157     pub fn insert(&mut self, value: T) -> NodeId {
158         self.len += 1;
159         match self.free_head {
160             Some(index) => {
161                 let slot = &mut self.slots[index as usize];
162                 let generation = slot.generation();
163                 let next_free = match slot {
164                     Slot::Vacant { next_free, .. } => *next_free,
165                     Slot::Occupied { .. } => unreachable!("free list pointed at an occupied slot"),
166                 };
167                 self.free_head = next_free;
168                 *slot = Slot::Occupied {
169                     generation,
170                     node: Node { parent: None, children: Vec::new(), value },
171                 };
172                 NodeId { index, generation }
173             }
174             None => {
175                 let index = self.slots.len() as u32;
176                 let generation = FIRST_GENERATION;
177                 self.slots.push(Slot::Occupied {
178                     generation,
179                     node: Node { parent: None, children: Vec::new(), value },
180                 });
181                 NodeId { index, generation }
182             }
183         }
184     }
185 
186     /// Shared access to a node, or `None` if `id` is stale/out of range.
187     #[inline]
188     pub fn get(&self, id: NodeId) -> Option<&Node<T>> {
189         match self.slots.get(id.index as usize) {
190             Some(Slot::Occupied { generation, node }) if *generation == id.generation => Some(node),
191             _ => None,
192         }
193     }
194 
195     /// Mutable access to a node, or `None` if `id` is stale/out of range.
196     #[inline]
197     pub fn get_mut(&mut self, id: NodeId) -> Option<&mut Node<T>> {
198         match self.slots.get_mut(id.index as usize) {
199             Some(Slot::Occupied { generation, node }) if *generation == id.generation => Some(node),
200             _ => None,
201         }
202     }
203 
204     /// Convenience: shared access to a node's payload.
205     #[inline]
206     pub fn value(&self, id: NodeId) -> Option<&T> {
207         self.get(id).map(Node::value)
208     }
209 
210     /// Convenience: mutable access to a node's payload.
211     #[inline]
212     pub fn value_mut(&mut self, id: NodeId) -> Option<&mut T> {
213         self.get_mut(id).map(Node::value_mut)
214     }
215 
216     /// Mutable access to two distinct nodes at once. Returns `None` if the ids are equal or
217     /// either is stale. Needed by passes that move data between two nodes (e.g. reparenting or
218     /// parent→child layout) without cloning.
219     pub fn get_pair_mut(&mut self, a: NodeId, b: NodeId) -> Option<(&mut Node<T>, &mut Node<T>)> {
220         if a.index == b.index {
221             return None;
222         }
223         let (lo, hi, swapped) = if a.index < b.index { (a, b, false) } else { (b, a, true) };
224         let (left, right) = self.slots.split_at_mut(hi.index as usize);
225         let lo_node = match left.get_mut(lo.index as usize) {
226             Some(Slot::Occupied { generation, node }) if *generation == lo.generation => node,
227             _ => return None,
228         };
229         let hi_node = match right.get_mut(0) {
230             Some(Slot::Occupied { generation, node }) if *generation == hi.generation => node,
231             _ => return None,
232         };
233         Some(if swapped { (hi_node, lo_node) } else { (lo_node, hi_node) })
234     }
235 
236     /// This node's parent (or `None` for a root or a stale id).
237     #[inline]
238     pub fn parent(&self, id: NodeId) -> Option<NodeId> {
239         self.get(id).and_then(Node::parent)
240     }
241 
242     /// This node's direct children (empty for a leaf or a stale id).
243     #[inline]
244     pub fn children(&self, id: NodeId) -> &[NodeId] {
245         match self.get(id) {
246             Some(node) => node.children(),
247             None => &[],
248         }
249     }
250 
251     /// Make `child` the last child of `parent`, detaching it from any previous parent first.
252     ///
253     /// Panics if either id is stale, if `parent == child`, or if the link would create a cycle
254     /// (`parent` is `child` or one of its descendants). These are programmer errors — reads of a
255     /// stale id are still safe via [`get`](Arena::get); it is *mutating* through one that trips.
256     pub fn append_child(&mut self, parent: NodeId, child: NodeId) {
257         assert!(self.contains(parent), "append_child: parent is not a live node");
258         assert!(self.contains(child), "append_child: child is not a live node");
259         assert!(parent != child, "append_child: cannot make a node its own child");
260         assert!(
261             !self.is_ancestor(child, parent),
262             "append_child: would create a cycle (parent is a descendant of child)"
263         );
264 
265         self.detach(child);
266         self.get_mut(child).unwrap().parent = Some(parent);
267         self.get_mut(parent).unwrap().children.push(child);
268     }
269 
270     /// Unlink `id` from its parent, leaving it (and its subtree) in the arena as a detached root.
271     /// No-op if `id` is stale or already a root.
272     pub fn detach(&mut self, id: NodeId) {
273         let Some(parent) = self.parent(id) else { return };
274         if let Some(parent_node) = self.get_mut(parent) {
275             parent_node.children.retain(|&c| c != id);
276         }
277         if let Some(node) = self.get_mut(id) {
278             node.parent = None;
279         }
280     }
281 
282     /// Remove `id` and its entire subtree from the arena, freeing every slot. Detaches `id` from
283     /// its parent first. Every [`NodeId`] into the removed subtree is invalidated (later lookups
284     /// return `None`). Returns the number of nodes removed; no-op returning 0 for a stale id.
285     pub fn remove_subtree(&mut self, id: NodeId) -> usize {
286         if !self.contains(id) {
287             return 0;
288         }
289         self.detach(id);
290 
291         // Collect the subtree (pre-order) before mutating, so we don't invalidate mid-walk.
292         let mut to_free = Vec::new();
293         let mut stack = vec![id];
294         while let Some(current) = stack.pop() {
295             to_free.push(current);
296             // Children are pushed as-is; order within the free set does not matter.
297             stack.extend_from_slice(self.children(current));
298         }
299 
300         for node_id in &to_free {
301             let index = node_id.index as usize;
302             let old_generation = self.slots[index].generation();
303             self.slots[index] = Slot::Vacant {
304                 generation: bump(old_generation),
305                 next_free: self.free_head,
306             };
307             self.free_head = Some(node_id.index);
308         }
309         self.len -= to_free.len();
310         to_free.len()
311     }
312 
313     /// Whether `maybe_ancestor` is `id` itself or one of its ancestors.
314     pub fn is_ancestor(&self, maybe_ancestor: NodeId, id: NodeId) -> bool {
315         let mut current = Some(id);
316         while let Some(node) = current {
317             if node == maybe_ancestor {
318                 return true;
319             }
320             current = self.parent(node);
321         }
322         false
323     }
324 
325     /// Iterator over `id`'s ancestors, nearest first, excluding `id` itself. Empty for a stale id.
326     pub fn ancestors(&self, id: NodeId) -> Ancestors<'_, T> {
327         Ancestors { arena: self, next: self.parent(id) }
328     }
329 
330     /// Iterator over the subtree rooted at `id` in pre-order (`id` first, then each child's
331     /// subtree). This is the walk order for layout and paint passes. Empty for a stale id.
332     pub fn subtree(&self, id: NodeId) -> Subtree<'_, T> {
333         let stack = if self.contains(id) { vec![id] } else { Vec::new() };
334         Subtree { arena: self, stack }
335     }
336 
337     /// Remove every node.
338     pub fn clear(&mut self) {
339         self.slots.clear();
340         self.free_head = None;
341         self.len = 0;
342     }
343 }
344 
345 /// Iterator returned by [`Arena::ancestors`].
346 pub struct Ancestors<'a, T> {
347     arena: &'a Arena<T>,
348     next: Option<NodeId>,
349 }
350 
351 impl<'a, T> Iterator for Ancestors<'a, T> {
352     type Item = NodeId;
353     fn next(&mut self) -> Option<NodeId> {
354         let current = self.next?;
355         self.next = self.arena.parent(current);
356         Some(current)
357     }
358 }
359 
360 /// Iterator returned by [`Arena::subtree`] (pre-order DFS).
361 pub struct Subtree<'a, T> {
362     arena: &'a Arena<T>,
363     stack: Vec<NodeId>,
364 }
365 
366 impl<'a, T> Iterator for Subtree<'a, T> {
367     type Item = NodeId;
368     fn next(&mut self) -> Option<NodeId> {
369         let current = self.stack.pop()?;
370         // Push children in reverse so they are visited left-to-right.
371         let children = self.arena.children(current);
372         for &child in children.iter().rev() {
373             self.stack.push(child);
374         }
375         Some(current)
376     }
377 }
378 
379 #[cfg(test)]
380 mod tests {
381     use super::*;
382 
383     #[test]
384     fn insert_and_get() {
385         let mut arena: Arena<&str> = Arena::new();
386         let a = arena.insert("a");
387         assert_eq!(arena.len(), 1);
388         assert!(!arena.is_empty());
389         assert_eq!(arena.value(a), Some(&"a"));
390         assert_eq!(arena.parent(a), None); // detached root
391         assert!(arena.children(a).is_empty());
392     }
393 
394     #[test]
395     fn stale_id_reads_as_none_after_removal() {
396         // The core safety property: a handle to a removed node never dereferences freed data.
397         let mut arena: Arena<i32> = Arena::new();
398         let a = arena.insert(10);
399         assert!(arena.contains(a));
400         arena.remove_subtree(a);
401         assert!(!arena.contains(a));
402         assert!(arena.get(a).is_none());
403         assert_eq!(arena.value(a), None);
404         assert_eq!(arena.len(), 0);
405     }
406 
407     #[test]
408     fn slot_reuse_bumps_generation_and_invalidates_old_handle() {
409         let mut arena: Arena<i32> = Arena::new();
410         let first = arena.insert(1);
411         arena.remove_subtree(first);
412         let second = arena.insert(2); // reuses the freed slot
413 
414         assert_eq!(second.slot_index(), first.slot_index(), "slot should be reused");
415         assert_ne!(second, first, "generation must differ so the old handle is distinct");
416         assert_eq!(arena.value(second), Some(&2));
417         assert!(arena.get(first).is_none(), "stale handle to the reused slot is still None");
418     }
419 
420     #[test]
421     fn append_child_links_both_ends() {
422         let mut arena: Arena<&str> = Arena::new();
423         let parent = arena.insert("p");
424         let child = arena.insert("c");
425         arena.append_child(parent, child);
426 
427         assert_eq!(arena.parent(child), Some(parent));
428         assert_eq!(arena.children(parent), &[child]);
429     }
430 
431     #[test]
432     fn reparenting_removes_from_old_parent() {
433         let mut arena: Arena<&str> = Arena::new();
434         let a = arena.insert("a");
435         let b = arena.insert("b");
436         let child = arena.insert("c");
437 
438         arena.append_child(a, child);
439         assert_eq!(arena.children(a), &[child]);
440 
441         arena.append_child(b, child);
442         assert!(arena.children(a).is_empty(), "old parent must drop the child");
443         assert_eq!(arena.children(b), &[child]);
444         assert_eq!(arena.parent(child), Some(b));
445     }
446 
447     #[test]
448     fn detach_keeps_node_but_unlinks_parent() {
449         let mut arena: Arena<&str> = Arena::new();
450         let parent = arena.insert("p");
451         let child = arena.insert("c");
452         arena.append_child(parent, child);
453 
454         arena.detach(child);
455         assert_eq!(arena.parent(child), None);
456         assert!(arena.children(parent).is_empty());
457         assert!(arena.contains(child), "detach must not free the node");
458     }
459 
460     #[test]
461     fn remove_subtree_frees_all_descendants() {
462         let mut arena: Arena<i32> = Arena::new();
463         let root = arena.insert(0);
464         let a = arena.insert(1);
465         let b = arena.insert(2);
466         let a1 = arena.insert(11);
467         arena.append_child(root, a);
468         arena.append_child(root, b);
469         arena.append_child(a, a1);
470 
471         let freed = arena.remove_subtree(a);
472         assert_eq!(freed, 2, "a and a1");
473         assert!(!arena.contains(a));
474         assert!(!arena.contains(a1));
475         assert!(arena.contains(root));
476         assert!(arena.contains(b));
477         assert_eq!(arena.children(root), &[b], "a must be gone from root's children");
478     }
479 
480     #[test]
481     fn subtree_iterates_preorder_left_to_right() {
482         let mut arena: Arena<&str> = Arena::new();
483         let root = arena.insert("root");
484         let a = arena.insert("a");
485         let b = arena.insert("b");
486         let a1 = arena.insert("a1");
487         let a2 = arena.insert("a2");
488         arena.append_child(root, a);
489         arena.append_child(root, b);
490         arena.append_child(a, a1);
491         arena.append_child(a, a2);
492 
493         let order: Vec<&str> = arena.subtree(root).map(|id| *arena.value(id).unwrap()).collect();
494         assert_eq!(order, vec!["root", "a", "a1", "a2", "b"]);
495     }
496 
497     #[test]
498     fn ancestors_walk_nearest_first() {
499         let mut arena: Arena<&str> = Arena::new();
500         let root = arena.insert("root");
501         let a = arena.insert("a");
502         let a1 = arena.insert("a1");
503         arena.append_child(root, a);
504         arena.append_child(a, a1);
505 
506         let anc: Vec<NodeId> = arena.ancestors(a1).collect();
507         assert_eq!(anc, vec![a, root]);
508         assert!(arena.ancestors(root).next().is_none(), "root has no ancestors");
509     }
510 
511     #[test]
512     fn is_ancestor_reports_self_and_chain() {
513         let mut arena: Arena<i32> = Arena::new();
514         let root = arena.insert(0);
515         let a = arena.insert(1);
516         arena.append_child(root, a);
517 
518         assert!(arena.is_ancestor(root, a));
519         assert!(arena.is_ancestor(a, a), "a node is its own ancestor for cycle-check purposes");
520         assert!(!arena.is_ancestor(a, root));
521     }
522 
523     #[test]
524     #[should_panic(expected = "cycle")]
525     fn append_child_rejects_cycles() {
526         let mut arena: Arena<i32> = Arena::new();
527         let root = arena.insert(0);
528         let a = arena.insert(1);
529         arena.append_child(root, a);
530         // Trying to make root a child of a (its descendant) would form a cycle.
531         arena.append_child(a, root);
532     }
533 
534     #[test]
535     fn get_pair_mut_yields_distinct_nodes_in_argument_order() {
536         let mut arena: Arena<i32> = Arena::new();
537         let a = arena.insert(1);
538         let b = arena.insert(2);
539 
540         let (na, nb) = arena.get_pair_mut(a, b).unwrap();
541         *na.value_mut() += 100;
542         *nb.value_mut() += 200;
543         assert_eq!(arena.value(a), Some(&101));
544         assert_eq!(arena.value(b), Some(&202));
545 
546         // Order preserved when the higher-index id is passed first.
547         let (nb2, na2) = arena.get_pair_mut(b, a).unwrap();
548         assert_eq!(*nb2.value(), 202);
549         assert_eq!(*na2.value(), 101);
550 
551         assert!(arena.get_pair_mut(a, a).is_none(), "same id must be rejected");
552     }
553 
554     #[test]
555     fn clear_empties_the_arena() {
556         let mut arena: Arena<i32> = Arena::new();
557         let a = arena.insert(1);
558         arena.insert(2);
559         arena.clear();
560         assert!(arena.is_empty());
561         assert!(arena.get(a).is_none());
562     }
563 
564     // Validation against a real trait object: the arena must be able to *own* and tree actual
565     // `dyn WidgetHost` widgets (the payload type Phase 3 will use), not just Copy scalars.
566     #[test]
567     fn holds_and_trees_real_dyn_element_payloads() {
568         use crate::widget::WidgetHost;
569 
570         // A minimal real `WidgetHost` — `color` is the trait's only required method, everything
571         // else is defaulted, so this exercises the actual trait object without dragging in a
572         // heavyweight widget constructor.
573         struct Marker {
574             base: crate::widget::Widget,
575             tint: [f32; 4],
576             painted: std::cell::Cell<bool>,
577         }
578         impl WidgetHost for Marker {
579             crate::impl_widget_base!(Marker);
580             fn color(&self) -> [f32; 4] {
581                 self.painted.set(true);
582                 self.tint
583             }
584         }
585 
586         let mut arena: Arena<Box<dyn WidgetHost>> = Arena::new();
587         let root = arena.insert(Box::new(Marker { base: crate::widget::Widget::new(), tint: [1.0, 0.0, 0.0, 1.0], painted: false.into() }));
588         let child = arena.insert(Box::new(Marker { base: crate::widget::Widget::new(), tint: [0.0, 1.0, 0.0, 1.0], painted: false.into() }));
589         arena.append_child(root, child);
590 
591         // Walk the subtree the way a paint pass will, calling a real trait method on each node.
592         let tints: Vec<[f32; 4]> = arena
593             .subtree(root)
594             .map(|id| arena.value(id).unwrap().color())
595             .collect();
596         assert_eq!(tints, vec![[1.0, 0.0, 0.0, 1.0], [0.0, 1.0, 0.0, 1.0]]);
597 
598         // The trait object is genuinely stored (its base is reachable through the box).
599         let _ = arena.value(root).unwrap().base();
600 
601         // Removing the root frees the child too, proving ownership lives in the arena.
602         arena.remove_subtree(root);
603         assert!(arena.is_empty());
604     }
605 }