src/tools/rust-analyzer/crates/ide/src/prime_caches/topologic_sort.rs

   1 //! helper data structure to schedule work for parallel prime caches.
   2 use std::{collections::VecDeque, hash::Hash};
   3
   4 use ide_db::FxHashMap;
   5
   6 pub(crate) struct TopologicSortIterBuilder<T> {
   7     nodes: FxHashMap<T, Entry<T>>,
   8 }
   9
  10 impl<T> TopologicSortIterBuilder<T>
  11 where
  12     T: Copy + Eq + PartialEq + Hash,
  13 {
  14     fn new() -> Self {
  15         Self { nodes: Default::default() }
  16     }
  17
  18     fn get_or_create_entry(&mut self, item: T) -> &mut Entry<T> {
  19         self.nodes.entry(item).or_default()
  20     }
  21
  22     pub(crate) fn add(&mut self, item: T, predecessors: impl IntoIterator<Item = T>) {
  23         let mut num_predecessors = 0;
  24
  25         for predecessor in predecessors.into_iter() {
  26             self.get_or_create_entry(predecessor).successors.push(item);
  27             num_predecessors += 1;
  28         }
  29
  30         let entry = self.get_or_create_entry(item);
  31         entry.num_predecessors += num_predecessors;
  32     }
  33
  34     pub(crate) fn build(self) -> TopologicalSortIter<T> {
  35         let ready = self
  36             .nodes
  37             .iter()
  38             .filter_map(
  39                 |(item, entry)| if entry.num_predecessors == 0 { Some(*item) } else { None },
  40             )
  41             .collect();
  42
  43         TopologicalSortIter { nodes: self.nodes, ready }
  44     }
  45 }
  46
  47 pub(crate) struct TopologicalSortIter<T> {
  48     ready: VecDeque<T>,
  49     nodes: FxHashMap<T, Entry<T>>,
  50 }
  51
  52 impl<T> TopologicalSortIter<T>
  53 where
  54     T: Copy + Eq + PartialEq + Hash,
  55 {
  56     pub(crate) fn builder() -> TopologicSortIterBuilder<T> {
  57         TopologicSortIterBuilder::new()
  58     }
  59
  60     pub(crate) fn pending(&self) -> usize {
  61         self.nodes.len()
  62     }
  63
  64     pub(crate) fn mark_done(&mut self, item: T) {
  65         let entry = self.nodes.remove(&item).expect("invariant: unknown item marked as done");
  66
  67         for successor in entry.successors {
  68             let succ_entry = self
  69                 .nodes
  70                 .get_mut(&successor)
  71                 .expect("invariant: unknown successor referenced by entry");
  72
  73             succ_entry.num_predecessors -= 1;
  74             if succ_entry.num_predecessors == 0 {
  75                 self.ready.push_back(successor);
  76             }
  77         }
  78     }
  79 }
  80
  81 impl<T> Iterator for TopologicalSortIter<T> {
  82     type Item = T;
  83
  84     fn next(&mut self) -> Option<Self::Item> {
  85         self.ready.pop_front()
  86     }
  87 }
  88
  89 struct Entry<T> {
  90     successors: Vec<T>,
  91     num_predecessors: usize,
  92 }
  93
  94 impl<T> Default for Entry<T> {
  95     fn default() -> Self {
  96         Self { successors: Default::default(), num_predecessors: 0 }
  97     }
  98 }