src/tools/rust-analyzer/crates/profile/src/tree.rs

   1 //! A simple tree implementation which tries to not allocate all over the place.
   2 use std::ops;
   3
   4 use la_arena::Arena;
   5
   6 #[derive(Default)]
   7 pub(crate) struct Tree<T> {
   8     nodes: Arena<Node<T>>,
   9     current_path: Vec<(Idx<T>, Option<Idx<T>>)>,
  10 }
  11
  12 pub(crate) type Idx<T> = la_arena::Idx<Node<T>>;
  13
  14 impl<T> Tree<T> {
  15     pub(crate) fn start(&mut self)
  16     where
  17         T: Default,
  18     {
  19         let me = self.nodes.alloc(Node::new(T::default()));
  20         if let Some((parent, last_child)) = self.current_path.last_mut() {
  21             let slot = match *last_child {
  22                 Some(last_child) => &mut self.nodes[last_child].next_sibling,
  23                 None => &mut self.nodes[*parent].first_child,
  24             };
  25             let prev = slot.replace(me);
  26             assert!(prev.is_none());
  27             *last_child = Some(me);
  28         }
  29
  30         self.current_path.push((me, None));
  31     }
  32
  33     pub(crate) fn finish(&mut self, data: T) {
  34         let (me, _last_child) = self.current_path.pop().unwrap();
  35         self.nodes[me].data = data;
  36     }
  37
  38     pub(crate) fn root(&self) -> Option<Idx<T>> {
  39         self.nodes.iter().next().map(|(idx, _)| idx)
  40     }
  41
  42     pub(crate) fn children(&self, idx: Idx<T>) -> impl Iterator<Item = Idx<T>> + '_ {
  43         NodeIter { nodes: &self.nodes, next: self.nodes[idx].first_child }
  44     }
  45     pub(crate) fn clear(&mut self) {
  46         self.nodes.clear();
  47         self.current_path.clear();
  48     }
  49 }
  50
  51 impl<T> ops::Index<Idx<T>> for Tree<T> {
  52     type Output = T;
  53     fn index(&self, index: Idx<T>) -> &T {
  54         &self.nodes[index].data
  55     }
  56 }
  57
  58 pub(crate) struct Node<T> {
  59     data: T,
  60     first_child: Option<Idx<T>>,
  61     next_sibling: Option<Idx<T>>,
  62 }
  63
  64 impl<T> Node<T> {
  65     fn new(data: T) -> Node<T> {
  66         Node { data, first_child: None, next_sibling: None }
  67     }
  68 }
  69
  70 struct NodeIter<'a, T> {
  71     nodes: &'a Arena<Node<T>>,
  72     next: Option<Idx<T>>,
  73 }
  74
  75 impl<'a, T> Iterator for NodeIter<'a, T> {
  76     type Item = Idx<T>;
  77
  78     fn next(&mut self) -> Option<Idx<T>> {
  79         self.next.map(|next| {
  80             self.next = self.nodes[next].next_sibling;
  81             next
  82         })
  83     }
  84 }