1use rkyv::{Archive, Deserialize, Serialize};
4use serde::{Deserialize as SerdeDeserialize, Serialize as SerdeSerialize};
5use std::collections::HashMap;
6use waremax_core::{EdgeId, NodeId};
7
8#[derive(
10 Archive,
11 Deserialize,
12 Serialize,
13 SerdeDeserialize,
14 SerdeSerialize,
15 Clone,
16 Debug,
17 PartialEq,
18 Default,
19)]
20pub enum EdgeDirection {
21 #[default]
23 Bidirectional,
24 OneWay,
26}
27
28#[derive(
30 Archive,
31 Deserialize,
32 Serialize,
33 SerdeDeserialize,
34 SerdeSerialize,
35 Clone,
36 Debug,
37 PartialEq,
38 Default,
39)]
40pub enum NodeType {
41 #[default]
42 Aisle,
43 StationPick,
44 StationDrop,
45 StationInbound,
46 StationOutbound,
47 Charging,
48 Staging,
49 Rack,
50}
51
52#[derive(Archive, Deserialize, Serialize, Clone, Debug)]
54pub struct Node {
55 pub id: NodeId,
56 pub string_id: String,
57 pub x: f64,
58 pub y: f64,
59 pub node_type: NodeType,
60 pub capacity: u32,
61}
62
63impl Node {
64 pub fn new(id: NodeId, string_id: String, x: f64, y: f64, node_type: NodeType) -> Self {
65 Self {
66 id,
67 string_id,
68 x,
69 y,
70 node_type,
71 capacity: 1,
72 }
73 }
74}
75
76#[derive(Archive, Deserialize, Serialize, Clone, Debug)]
78pub struct Edge {
79 pub id: EdgeId,
80 pub from: NodeId,
81 pub to: NodeId,
82 pub length_m: f64,
83 pub capacity: u32,
84 pub direction: EdgeDirection,
86 pub speed_multiplier: f64,
89}
90
91impl Edge {
92 pub fn new(id: EdgeId, from: NodeId, to: NodeId, length_m: f64) -> Self {
93 Self {
94 id,
95 from,
96 to,
97 length_m,
98 capacity: 1,
99 direction: EdgeDirection::Bidirectional,
100 speed_multiplier: 1.0,
101 }
102 }
103
104 pub fn with_direction(mut self, direction: EdgeDirection) -> Self {
106 self.direction = direction;
107 self
108 }
109
110 pub fn with_capacity(mut self, capacity: u32) -> Self {
112 self.capacity = capacity;
113 self
114 }
115
116 pub fn with_speed_multiplier(mut self, multiplier: f64) -> Self {
120 self.speed_multiplier = multiplier;
121 self
122 }
123}
124
125#[derive(Clone, Default)]
127pub struct WarehouseMap {
128 pub nodes: HashMap<NodeId, Node>,
129 pub edges: HashMap<EdgeId, Edge>,
130 pub adjacency: HashMap<NodeId, Vec<(NodeId, EdgeId, f64)>>,
131 pub string_to_node: HashMap<String, NodeId>,
132 pub blocked_nodes: Vec<NodeId>,
133 pub blocked_edges: Vec<EdgeId>,
134}
135
136impl WarehouseMap {
137 pub fn new() -> Self {
138 Self::default()
139 }
140
141 pub fn add_node(&mut self, node: Node) {
142 let id = node.id;
143 let string_id = node.string_id.clone();
144 self.adjacency.entry(id).or_default();
145 self.string_to_node.insert(string_id, id);
146 self.nodes.insert(id, node);
147 }
148
149 pub fn add_edge(&mut self, edge: Edge) {
154 let from = edge.from;
155 let to = edge.to;
156 let length = edge.length_m;
157 let edge_id = edge.id;
158 let direction = edge.direction.clone();
159
160 self.adjacency
161 .entry(from)
162 .or_default()
163 .push((to, edge_id, length));
164
165 if direction == EdgeDirection::Bidirectional {
166 let reverse_id = EdgeId(edge_id.0 + 100000);
167 self.adjacency
168 .entry(to)
169 .or_default()
170 .push((from, reverse_id, length));
171 self.edges.insert(
172 reverse_id,
173 Edge {
174 id: reverse_id,
175 from: to,
176 to: from,
177 length_m: length,
178 capacity: edge.capacity,
179 direction: EdgeDirection::OneWay, speed_multiplier: edge.speed_multiplier, },
182 );
183 }
184
185 self.edges.insert(edge_id, edge);
186 }
187
188 pub fn get_node(&self, id: NodeId) -> Option<&Node> {
189 self.nodes.get(&id)
190 }
191
192 pub fn get_node_by_string(&self, s: &str) -> Option<&Node> {
193 self.string_to_node.get(s).and_then(|id| self.nodes.get(id))
194 }
195
196 pub fn get_edge(&self, id: EdgeId) -> Option<&Edge> {
197 self.edges.get(&id)
198 }
199
200 pub fn neighbors(&self, node: NodeId) -> impl Iterator<Item = (NodeId, EdgeId, f64)> + '_ {
201 self.adjacency
202 .get(&node)
203 .into_iter()
204 .flat_map(|v| v.iter().copied())
205 .filter(move |(neighbor, edge_id, _)| {
206 !self.blocked_nodes.contains(neighbor) && !self.blocked_edges.contains(edge_id)
207 })
208 }
209
210 pub fn euclidean_distance(&self, from: NodeId, to: NodeId) -> f64 {
211 let n1 = self.nodes.get(&from);
212 let n2 = self.nodes.get(&to);
213 match (n1, n2) {
214 (Some(a), Some(b)) => ((a.x - b.x).powi(2) + (a.y - b.y).powi(2)).sqrt(),
215 _ => f64::INFINITY,
216 }
217 }
218
219 pub fn node_count(&self) -> usize {
220 self.nodes.len()
221 }
222
223 pub fn edge_count(&self) -> usize {
224 self.edges.len()
225 }
226}