1use std::collections::{HashMap, HashSet};
13use std::hash::Hash;
14
15#[derive(Debug)]
16pub struct DependencyQueue<N: Hash + Eq, E: Hash + Eq, V> {
17 dep_map: HashMap<N, (HashSet<(N, E)>, V)>,
23
24 reverse_dep_map: HashMap<N, HashMap<E, HashSet<N>>>,
33
34 priority: HashMap<N, usize>,
36
37 cost: HashMap<N, usize>,
39}
40
41impl<N: Hash + Eq, E: Hash + Eq, V> Default for DependencyQueue<N, E, V> {
42 fn default() -> DependencyQueue<N, E, V> {
43 DependencyQueue::new()
44 }
45}
46
47impl<N: Hash + Eq, E: Hash + Eq, V> DependencyQueue<N, E, V> {
48 pub fn new() -> DependencyQueue<N, E, V> {
50 DependencyQueue {
51 dep_map: HashMap::new(),
52 reverse_dep_map: HashMap::new(),
53 priority: HashMap::new(),
54 cost: HashMap::new(),
55 }
56 }
57}
58
59impl<N: Hash + Eq + Clone, E: Eq + Hash + Clone, V> DependencyQueue<N, E, V> {
60 pub fn queue(
76 &mut self,
77 key: N,
78 value: V,
79 dependencies: impl IntoIterator<Item = (N, E)>,
80 cost: usize,
81 ) {
82 assert!(!self.dep_map.contains_key(&key));
83
84 let mut my_dependencies = HashSet::new();
85 for (dep, edge) in dependencies {
86 my_dependencies.insert((dep.clone(), edge.clone()));
87 self.reverse_dep_map
88 .entry(dep)
89 .or_insert_with(HashMap::new)
90 .entry(edge)
91 .or_insert_with(HashSet::new)
92 .insert(key.clone());
93 }
94 self.dep_map.insert(key.clone(), (my_dependencies, value));
95 self.cost.insert(key, cost);
96 }
97
98 pub fn queue_finished(&mut self) {
101 let mut out = HashMap::new();
102 for key in self.dep_map.keys() {
103 depth(key, &self.reverse_dep_map, &mut out);
104 }
105 self.priority = out
106 .into_iter()
107 .map(|(n, set)| {
108 let total_cost =
109 self.cost[&n] + set.iter().map(|key| self.cost[key]).sum::<usize>();
110 (n, total_cost)
111 })
112 .collect();
113
114 fn depth<'a, N: Hash + Eq + Clone, E: Hash + Eq + Clone>(
119 key: &N,
120 map: &HashMap<N, HashMap<E, HashSet<N>>>,
121 results: &'a mut HashMap<N, HashSet<N>>,
122 ) -> &'a HashSet<N> {
123 if results.contains_key(key) {
124 let depth = &results[key];
125 assert!(!depth.is_empty(), "cycle in DependencyQueue");
126 return depth;
127 }
128 results.insert(key.clone(), HashSet::new());
129
130 let mut set = HashSet::new();
131 set.insert(key.clone());
132
133 for dep in map
134 .get(key)
135 .into_iter()
136 .flat_map(|it| it.values())
137 .flatten()
138 {
139 set.extend(depth(dep, map, results).iter().cloned())
140 }
141
142 let slot = results.get_mut(key).unwrap();
143 *slot = set;
144 &*slot
145 }
146 }
147
148 pub fn dequeue(&mut self) -> Option<(N, V, usize)> {
153 let (key, priority) = self
154 .dep_map
155 .iter()
156 .filter(|(_, (deps, _))| deps.is_empty())
157 .map(|(key, _)| (key.clone(), self.priority[key]))
158 .max_by_key(|(_, priority)| *priority)?;
159 let (_, data) = self.dep_map.remove(&key).unwrap();
160 Some((key, data, priority))
161 }
162
163 pub fn is_empty(&self) -> bool {
165 self.dep_map.is_empty()
166 }
167
168 pub fn len(&self) -> usize {
170 self.dep_map.len()
171 }
172
173 pub fn finish(&mut self, node: &N, edge: &E) -> Vec<&N> {
182 let reverse_deps = self.reverse_dep_map.get(node).and_then(|map| map.get(edge));
184 let Some(reverse_deps) = reverse_deps else {
185 return Vec::new();
186 };
187 let key = (node.clone(), edge.clone());
188 let mut result = Vec::new();
189 for dep in reverse_deps.iter() {
190 let edges = &mut self.dep_map.get_mut(dep).unwrap().0;
191 assert!(edges.remove(&key));
192 if edges.is_empty() {
193 result.push(dep);
194 }
195 }
196 result
197 }
198}
199
200#[cfg(test)]
201mod test {
202 use super::DependencyQueue;
203
204 #[test]
205 fn deep_first_equal_cost() {
206 let mut q = DependencyQueue::new();
207
208 q.queue(1, (), vec![], 1);
209 q.queue(2, (), vec![(1, ())], 1);
210 q.queue(3, (), vec![], 1);
211 q.queue(4, (), vec![(2, ()), (3, ())], 1);
212 q.queue(5, (), vec![(4, ()), (3, ())], 1);
213 q.queue_finished();
214
215 assert_eq!(q.dequeue(), Some((1, (), 5)));
216 assert_eq!(q.dequeue(), Some((3, (), 4)));
217 assert_eq!(q.dequeue(), None);
218 q.finish(&3, &());
219 assert_eq!(q.dequeue(), None);
220 q.finish(&1, &());
221 assert_eq!(q.dequeue(), Some((2, (), 4)));
222 assert_eq!(q.dequeue(), None);
223 q.finish(&2, &());
224 assert_eq!(q.dequeue(), Some((4, (), 3)));
225 assert_eq!(q.dequeue(), None);
226 q.finish(&4, &());
227 assert_eq!(q.dequeue(), Some((5, (), 2)));
228 }
229
230 #[test]
231 fn sort_by_highest_cost() {
232 let mut q = DependencyQueue::new();
233
234 q.queue(1, (), vec![], 1);
235 q.queue(2, (), vec![(1, ())], 1);
236 q.queue(3, (), vec![], 4);
237 q.queue(4, (), vec![(2, ()), (3, ())], 1);
238 q.queue_finished();
239
240 assert_eq!(q.dequeue(), Some((3, (), 9)));
241 assert_eq!(q.dequeue(), Some((1, (), 4)));
242 assert_eq!(q.dequeue(), None);
243 q.finish(&3, &());
244 assert_eq!(q.dequeue(), None);
245 q.finish(&1, &());
246 assert_eq!(q.dequeue(), Some((2, (), 3)));
247 assert_eq!(q.dequeue(), None);
248 q.finish(&2, &());
249 assert_eq!(q.dequeue(), Some((4, (), 2)));
250 assert_eq!(q.dequeue(), None);
251 q.finish(&4, &());
252 assert_eq!(q.dequeue(), None);
253 }
254}