1use std::collections::{HashMap, HashSet};
7use waremax_core::{EdgeId, NodeId, RobotId};
8
9#[derive(Clone, Debug)]
11pub enum WaitingFor {
12 Edge {
14 edge_id: EdgeId,
15 blocked_by: Vec<RobotId>,
16 },
17 Node {
19 node_id: NodeId,
20 blocked_by: Vec<RobotId>,
21 },
22}
23
24impl WaitingFor {
25 pub fn blockers(&self) -> &[RobotId] {
27 match self {
28 WaitingFor::Edge { blocked_by, .. } => blocked_by,
29 WaitingFor::Node { blocked_by, .. } => blocked_by,
30 }
31 }
32}
33
34#[derive(Clone, Debug, Default)]
42pub struct WaitForGraph {
43 waiting_for: HashMap<RobotId, WaitingFor>,
45}
46
47impl WaitForGraph {
48 pub fn new() -> Self {
50 Self {
51 waiting_for: HashMap::new(),
52 }
53 }
54
55 pub fn add_wait(&mut self, robot: RobotId, waiting_for: WaitingFor) {
57 self.waiting_for.insert(robot, waiting_for);
58 }
59
60 pub fn remove_wait(&mut self, robot: RobotId) {
62 self.waiting_for.remove(&robot);
63 }
64
65 pub fn is_waiting(&self, robot: RobotId) -> bool {
67 self.waiting_for.contains_key(&robot)
68 }
69
70 pub fn get_wait(&self, robot: RobotId) -> Option<&WaitingFor> {
72 self.waiting_for.get(&robot)
73 }
74
75 pub fn waiting_count(&self) -> usize {
77 self.waiting_for.len()
78 }
79
80 pub fn detect_cycle(&self) -> Option<Vec<RobotId>> {
87 for &start_robot in self.waiting_for.keys() {
89 if let Some(cycle) = self.find_cycle_from(start_robot) {
90 return Some(cycle);
91 }
92 }
93 None
94 }
95
96 fn find_cycle_from(&self, start: RobotId) -> Option<Vec<RobotId>> {
98 let mut visited = HashSet::new();
99 let mut path = Vec::new();
100 let mut path_set = HashSet::new();
101
102 self.dfs_find_cycle(start, &mut visited, &mut path, &mut path_set, start)
103 }
104
105 fn dfs_find_cycle(
107 &self,
108 current: RobotId,
109 visited: &mut HashSet<RobotId>,
110 path: &mut Vec<RobotId>,
111 path_set: &mut HashSet<RobotId>,
112 start: RobotId,
113 ) -> Option<Vec<RobotId>> {
114 if visited.contains(¤t) {
116 return None;
117 }
118
119 if path_set.contains(¤t) {
121 let cycle_start_idx = path.iter().position(|&r| r == current)?;
123 return Some(path[cycle_start_idx..].to_vec());
124 }
125
126 path.push(current);
128 path_set.insert(current);
129
130 if let Some(waiting) = self.waiting_for.get(¤t) {
132 for &blocker in waiting.blockers() {
134 if blocker == start && path.len() > 1 {
136 path.push(start);
137 return Some(path.clone());
138 }
139
140 if self.waiting_for.contains_key(&blocker) {
142 if let Some(cycle) =
143 self.dfs_find_cycle(blocker, visited, path, path_set, start)
144 {
145 return Some(cycle);
146 }
147 }
148 }
149 }
150
151 path.pop();
153 path_set.remove(¤t);
154 visited.insert(current);
155
156 None
157 }
158
159 pub fn detect_all_cycles(&self) -> Vec<Vec<RobotId>> {
161 let mut cycles = Vec::new();
162 let mut found_robots = HashSet::new();
163
164 for &start_robot in self.waiting_for.keys() {
165 if found_robots.contains(&start_robot) {
166 continue;
167 }
168
169 if let Some(cycle) = self.find_cycle_from(start_robot) {
170 for &robot in &cycle {
171 found_robots.insert(robot);
172 }
173 cycles.push(cycle);
174 }
175 }
176
177 cycles
178 }
179
180 pub fn clear(&mut self) {
182 self.waiting_for.clear();
183 }
184}
185
186#[cfg(test)]
187mod tests {
188 use super::*;
189
190 #[test]
191 fn test_empty_graph_no_cycle() {
192 let graph = WaitForGraph::new();
193 assert!(graph.detect_cycle().is_none());
194 }
195
196 #[test]
197 fn test_single_wait_no_cycle() {
198 let mut graph = WaitForGraph::new();
199 graph.add_wait(
200 RobotId(1),
201 WaitingFor::Edge {
202 edge_id: EdgeId(100),
203 blocked_by: vec![RobotId(2)],
204 },
205 );
206 assert!(graph.detect_cycle().is_none());
208 }
209
210 #[test]
211 fn test_two_robot_cycle() {
212 let mut graph = WaitForGraph::new();
213 graph.add_wait(
215 RobotId(1),
216 WaitingFor::Edge {
217 edge_id: EdgeId(100),
218 blocked_by: vec![RobotId(2)],
219 },
220 );
221 graph.add_wait(
223 RobotId(2),
224 WaitingFor::Node {
225 node_id: NodeId(50),
226 blocked_by: vec![RobotId(1)],
227 },
228 );
229
230 let cycle = graph.detect_cycle();
231 assert!(cycle.is_some());
232 let cycle = cycle.unwrap();
233 assert!(cycle.contains(&RobotId(1)));
234 assert!(cycle.contains(&RobotId(2)));
235 }
236
237 #[test]
238 fn test_three_robot_cycle() {
239 let mut graph = WaitForGraph::new();
240 graph.add_wait(
242 RobotId(1),
243 WaitingFor::Edge {
244 edge_id: EdgeId(100),
245 blocked_by: vec![RobotId(2)],
246 },
247 );
248 graph.add_wait(
249 RobotId(2),
250 WaitingFor::Edge {
251 edge_id: EdgeId(101),
252 blocked_by: vec![RobotId(3)],
253 },
254 );
255 graph.add_wait(
256 RobotId(3),
257 WaitingFor::Edge {
258 edge_id: EdgeId(102),
259 blocked_by: vec![RobotId(1)],
260 },
261 );
262
263 let cycle = graph.detect_cycle();
264 assert!(cycle.is_some());
265 let cycle = cycle.unwrap();
266 assert_eq!(cycle.len(), 4); }
268
269 #[test]
270 fn test_chain_no_cycle() {
271 let mut graph = WaitForGraph::new();
272 graph.add_wait(
274 RobotId(1),
275 WaitingFor::Edge {
276 edge_id: EdgeId(100),
277 blocked_by: vec![RobotId(2)],
278 },
279 );
280 graph.add_wait(
281 RobotId(2),
282 WaitingFor::Edge {
283 edge_id: EdgeId(101),
284 blocked_by: vec![RobotId(3)],
285 },
286 );
287 assert!(graph.detect_cycle().is_none());
290 }
291
292 #[test]
293 fn test_remove_wait_breaks_cycle() {
294 let mut graph = WaitForGraph::new();
295 graph.add_wait(
296 RobotId(1),
297 WaitingFor::Edge {
298 edge_id: EdgeId(100),
299 blocked_by: vec![RobotId(2)],
300 },
301 );
302 graph.add_wait(
303 RobotId(2),
304 WaitingFor::Node {
305 node_id: NodeId(50),
306 blocked_by: vec![RobotId(1)],
307 },
308 );
309
310 assert!(graph.detect_cycle().is_some());
312
313 graph.remove_wait(RobotId(1));
315
316 assert!(graph.detect_cycle().is_none());
318 }
319
320 #[test]
321 fn test_multiple_blockers() {
322 let mut graph = WaitForGraph::new();
323 graph.add_wait(
325 RobotId(1),
326 WaitingFor::Edge {
327 edge_id: EdgeId(100),
328 blocked_by: vec![RobotId(2), RobotId(3)],
329 },
330 );
331 graph.add_wait(
333 RobotId(2),
334 WaitingFor::Node {
335 node_id: NodeId(50),
336 blocked_by: vec![RobotId(1)],
337 },
338 );
339 let cycle = graph.detect_cycle();
342 assert!(cycle.is_some());
343 }
344
345 #[test]
346 fn test_is_waiting() {
347 let mut graph = WaitForGraph::new();
348 assert!(!graph.is_waiting(RobotId(1)));
349
350 graph.add_wait(
351 RobotId(1),
352 WaitingFor::Edge {
353 edge_id: EdgeId(100),
354 blocked_by: vec![RobotId(2)],
355 },
356 );
357
358 assert!(graph.is_waiting(RobotId(1)));
359 assert!(!graph.is_waiting(RobotId(2)));
360 }
361
362 #[test]
363 fn test_waiting_count() {
364 let mut graph = WaitForGraph::new();
365 assert_eq!(graph.waiting_count(), 0);
366
367 graph.add_wait(
368 RobotId(1),
369 WaitingFor::Edge {
370 edge_id: EdgeId(100),
371 blocked_by: vec![RobotId(2)],
372 },
373 );
374 assert_eq!(graph.waiting_count(), 1);
375
376 graph.add_wait(
377 RobotId(2),
378 WaitingFor::Node {
379 node_id: NodeId(50),
380 blocked_by: vec![RobotId(1)],
381 },
382 );
383 assert_eq!(graph.waiting_count(), 2);
384
385 graph.remove_wait(RobotId(1));
386 assert_eq!(graph.waiting_count(), 1);
387 }
388}