//! A simple tree implementation which tries to not allocate all over the place. use std::ops; use ra_arena::Arena; #[derive(Default)] pub struct Tree<T> { nodes: Arena<Node<T>>, current_path: Vec<(Idx<T>, Option<Idx<T>>)>, } pub type Idx<T> = ra_arena::Idx<Node<T>>; impl<T> Tree<T> { pub fn start(&mut self) where T: Default, { let me = self.nodes.alloc(Node::new(T::default())); if let Some((parent, last_child)) = self.current_path.last_mut() { let slot = match *last_child { Some(last_child) => &mut self.nodes[last_child].next_sibling, None => &mut self.nodes[*parent].first_child, }; let prev = slot.replace(me); assert!(prev.is_none()); *last_child = Some(me); } self.current_path.push((me, None)); } pub fn finish(&mut self, data: T) { let (me, _last_child) = self.current_path.pop().unwrap(); self.nodes[me].data = data; } pub fn root(&self) -> Option<Idx<T>> { self.nodes.iter().next().map(|(idx, _)| idx) } pub fn children(&self, idx: Idx<T>) -> impl Iterator<Item = Idx<T>> + '_ { NodeIter { nodes: &self.nodes, next: self.nodes[idx].first_child } } pub fn clear(&mut self) { self.nodes.clear(); self.current_path.clear(); } } impl<T> ops::Index<Idx<T>> for Tree<T> { type Output = T; fn index(&self, index: Idx<T>) -> &T { &self.nodes[index].data } } pub struct Node<T> { data: T, first_child: Option<Idx<T>>, next_sibling: Option<Idx<T>>, } impl<T> Node<T> { fn new(data: T) -> Node<T> { Node { data, first_child: None, next_sibling: None } } } struct NodeIter<'a, T> { nodes: &'a Arena<Node<T>>, next: Option<Idx<T>>, } impl<'a, T> Iterator for NodeIter<'a, T> { type Item = Idx<T>; fn next(&mut self) -> Option<Idx<T>> { self.next.map(|next| { self.next = self.nodes[next].next_sibling; next }) } }