waremax_map/
traffic.rs

1//! Traffic management for edge and node capacity
2
3use crate::deadlock::{WaitForGraph, WaitingFor};
4use std::collections::{HashMap, HashSet};
5use waremax_core::{EdgeId, NodeId, RobotId};
6
7/// Manages traffic flow and capacity constraints in the warehouse
8pub struct TrafficManager {
9    edge_occupancy: HashMap<EdgeId, HashSet<RobotId>>,
10    node_occupancy: HashMap<NodeId, HashSet<RobotId>>,
11    edge_capacity: HashMap<EdgeId, u32>,
12    node_capacity: HashMap<NodeId, u32>,
13    default_edge_capacity: u32,
14    default_node_capacity: u32,
15    /// v2: Wait-for graph for deadlock detection
16    pub wait_graph: WaitForGraph,
17    /// v2: Whether deadlock detection is enabled
18    pub deadlock_detection_enabled: bool,
19}
20
21impl TrafficManager {
22    pub fn new(default_edge_capacity: u32, default_node_capacity: u32) -> Self {
23        Self {
24            edge_occupancy: HashMap::new(),
25            node_occupancy: HashMap::new(),
26            edge_capacity: HashMap::new(),
27            node_capacity: HashMap::new(),
28            default_edge_capacity,
29            default_node_capacity,
30            wait_graph: WaitForGraph::new(),
31            deadlock_detection_enabled: false,
32        }
33    }
34
35    /// Enable or disable deadlock detection
36    pub fn set_deadlock_detection(&mut self, enabled: bool) {
37        self.deadlock_detection_enabled = enabled;
38    }
39
40    pub fn set_edge_capacity(&mut self, edge: EdgeId, capacity: u32) {
41        self.edge_capacity.insert(edge, capacity);
42    }
43
44    pub fn set_node_capacity(&mut self, node: NodeId, capacity: u32) {
45        self.node_capacity.insert(node, capacity);
46    }
47
48    pub fn can_enter_edge(&self, edge: EdgeId, robot: RobotId) -> bool {
49        let capacity = self
50            .edge_capacity
51            .get(&edge)
52            .copied()
53            .unwrap_or(self.default_edge_capacity);
54        let occupants = self.edge_occupancy.get(&edge);
55
56        if let Some(set) = occupants {
57            if set.contains(&robot) {
58                return true;
59            }
60            (set.len() as u32) < capacity
61        } else {
62            capacity > 0
63        }
64    }
65
66    pub fn can_enter_node(&self, node: NodeId, robot: RobotId) -> bool {
67        let capacity = self
68            .node_capacity
69            .get(&node)
70            .copied()
71            .unwrap_or(self.default_node_capacity);
72        let occupants = self.node_occupancy.get(&node);
73
74        if let Some(set) = occupants {
75            if set.contains(&robot) {
76                return true;
77            }
78            (set.len() as u32) < capacity
79        } else {
80            capacity > 0
81        }
82    }
83
84    pub fn enter_edge(&mut self, edge: EdgeId, robot: RobotId) {
85        self.edge_occupancy.entry(edge).or_default().insert(robot);
86    }
87
88    pub fn leave_edge(&mut self, edge: EdgeId, robot: RobotId) {
89        if let Some(set) = self.edge_occupancy.get_mut(&edge) {
90            set.remove(&robot);
91        }
92    }
93
94    pub fn enter_node(&mut self, node: NodeId, robot: RobotId) {
95        self.node_occupancy.entry(node).or_default().insert(robot);
96    }
97
98    pub fn leave_node(&mut self, node: NodeId, robot: RobotId) {
99        if let Some(set) = self.node_occupancy.get_mut(&node) {
100            set.remove(&robot);
101        }
102    }
103
104    pub fn get_edge_occupancy(&self, edge: EdgeId) -> usize {
105        self.edge_occupancy.get(&edge).map_or(0, |s| s.len())
106    }
107
108    pub fn get_node_occupancy(&self, node: NodeId) -> usize {
109        self.node_occupancy.get(&node).map_or(0, |s| s.len())
110    }
111
112    pub fn robots_on_edge(&self, edge: EdgeId) -> impl Iterator<Item = RobotId> + '_ {
113        self.edge_occupancy
114            .get(&edge)
115            .into_iter()
116            .flat_map(|s| s.iter().copied())
117    }
118
119    pub fn robots_at_node(&self, node: NodeId) -> impl Iterator<Item = RobotId> + '_ {
120        self.node_occupancy
121            .get(&node)
122            .into_iter()
123            .flat_map(|s| s.iter().copied())
124    }
125
126    // === v2: Deadlock Detection Methods ===
127
128    /// Record that a robot is waiting for an edge
129    ///
130    /// Automatically determines which robots are blocking the edge.
131    pub fn record_edge_wait(&mut self, robot: RobotId, edge: EdgeId) {
132        if !self.deadlock_detection_enabled {
133            return;
134        }
135
136        let blockers: Vec<RobotId> = self.robots_on_edge(edge).filter(|&r| r != robot).collect();
137
138        self.wait_graph.add_wait(
139            robot,
140            WaitingFor::Edge {
141                edge_id: edge,
142                blocked_by: blockers,
143            },
144        );
145    }
146
147    /// Record that a robot is waiting for a node
148    ///
149    /// Automatically determines which robots are blocking the node.
150    pub fn record_node_wait(&mut self, robot: RobotId, node: NodeId) {
151        if !self.deadlock_detection_enabled {
152            return;
153        }
154
155        let blockers: Vec<RobotId> = self.robots_at_node(node).filter(|&r| r != robot).collect();
156
157        self.wait_graph.add_wait(
158            robot,
159            WaitingFor::Node {
160                node_id: node,
161                blocked_by: blockers,
162            },
163        );
164    }
165
166    /// Clear a robot's wait status (e.g., when it acquires the resource)
167    pub fn clear_wait(&mut self, robot: RobotId) {
168        self.wait_graph.remove_wait(robot);
169    }
170
171    /// Check if a robot is currently waiting
172    pub fn is_waiting(&self, robot: RobotId) -> bool {
173        self.wait_graph.is_waiting(robot)
174    }
175
176    /// Check for deadlocks in the current wait-for graph
177    ///
178    /// Returns Some(cycle) if a deadlock is detected, None otherwise.
179    pub fn check_deadlock(&self) -> Option<Vec<RobotId>> {
180        if !self.deadlock_detection_enabled {
181            return None;
182        }
183        self.wait_graph.detect_cycle()
184    }
185
186    /// Get all robots currently waiting
187    pub fn waiting_count(&self) -> usize {
188        self.wait_graph.waiting_count()
189    }
190}
waremax_map/traffic.rs

waremax_map/
traffic.rs
