Skip to main content

nms_graph/
edges.rs

1//! Graph edge construction strategies.
2//!
3//! Edges connect systems in the petgraph. Edge weights are distances in
4//! light-years. Two strategies are provided:
5//!
6//! - **KNN**: connect each system to its k nearest neighbors
7//! - **Warp range**: connect systems within a maximum warp distance
8
9use nms_core::address::LY_PER_VOXEL;
10use rstar::PointDistance;
11
12use crate::model::GalaxyModel;
13use crate::spatial::SystemId;
14
15/// Strategy for generating graph edges.
16#[derive(Debug, Clone, Copy)]
17pub enum EdgeStrategy {
18    /// Connect each system to its k nearest neighbors.
19    Knn { k: usize },
20    /// Connect all systems within warp range (light-years).
21    WarpRange { max_ly: f64 },
22}
23
24impl Default for EdgeStrategy {
25    fn default() -> Self {
26        EdgeStrategy::Knn { k: 10 }
27    }
28}
29
30impl GalaxyModel {
31    /// Build edges using the given strategy. Clears existing edges first.
32    pub fn build_edges(&mut self, strategy: EdgeStrategy) {
33        // Remove all existing edges
34        let edge_ids: Vec<_> = self.graph.edge_indices().collect();
35        for edge_id in edge_ids {
36            self.graph.remove_edge(edge_id);
37        }
38
39        match strategy {
40            EdgeStrategy::Knn { k } => self.build_knn_edges(k),
41            EdgeStrategy::WarpRange { max_ly } => self.build_warp_range_edges(max_ly),
42        }
43    }
44
45    /// Connect each system to its k nearest neighbors.
46    ///
47    /// For each system, queries the R-tree for the k+1 nearest points
48    /// (including itself), skips self, and adds undirected edges.
49    /// Duplicate edges (A-B when B-A already exists) are avoided.
50    /// Each system is matched against its own galaxy's spatial index.
51    fn build_knn_edges(&mut self, k: usize) {
52        let system_ids: Vec<SystemId> = self.systems.keys().copied().collect();
53
54        for &sys_id in &system_ids {
55            let system = match self.systems.get(&sys_id) {
56                Some(s) => s,
57                None => continue,
58            };
59
60            let sys_addr = system.address;
61            let galaxy = sys_addr.reality_index;
62            let spatial = match self.spatial.get(&galaxy) {
63                Some(s) => s,
64                None => continue,
65            };
66
67            let query_point = [
68                sys_addr.voxel_x() as f64,
69                sys_addr.voxel_y() as f64,
70                sys_addr.voxel_z() as f64,
71            ];
72
73            // Collect R-tree neighbor IDs first, then look up distances.
74            let neighbor_ids: Vec<SystemId> = spatial
75                .nearest_neighbor_iter(&query_point)
76                .filter(|sp| sp.id != sys_id)
77                .take(k)
78                .map(|sp| sp.id)
79                .collect();
80
81            let neighbors: Vec<_> = neighbor_ids
82                .into_iter()
83                .filter_map(|nid| {
84                    let neighbor_sys = self.systems.get(&nid)?;
85                    let dist_ly = sys_addr.distance_ly(&neighbor_sys.address);
86                    Some((nid, dist_ly))
87                })
88                .collect();
89
90            let from_node = match self.node_map.get(&sys_id) {
91                Some(&n) => n,
92                None => continue,
93            };
94
95            for (neighbor_id, dist_ly) in neighbors {
96                let to_node = match self.node_map.get(&neighbor_id) {
97                    Some(&n) => n,
98                    None => continue,
99                };
100
101                // Avoid duplicate edges (check both directions)
102                if self.graph.find_edge(from_node, to_node).is_none() {
103                    self.graph.add_edge(from_node, to_node, dist_ly);
104                }
105            }
106        }
107    }
108
109    /// Connect all pairs of systems within a warp range (in light-years).
110    ///
111    /// For each system, queries its galaxy's R-tree for all neighbors within
112    /// `max_ly / LY_PER_VOXEL` voxel units and adds edges.
113    fn build_warp_range_edges(&mut self, max_ly: f64) {
114        let system_ids: Vec<SystemId> = self.systems.keys().copied().collect();
115
116        for &sys_id in &system_ids {
117            let system = match self.systems.get(&sys_id) {
118                Some(s) => s,
119                None => continue,
120            };
121
122            let sys_addr = system.address;
123            let galaxy = sys_addr.reality_index;
124            let spatial = match self.spatial.get(&galaxy) {
125                Some(s) => s,
126                None => continue,
127            };
128
129            let query_point = [
130                sys_addr.voxel_x() as f64,
131                sys_addr.voxel_y() as f64,
132                sys_addr.voxel_z() as f64,
133            ];
134
135            let voxel_radius = max_ly / LY_PER_VOXEL + 1.0;
136            let voxel_radius_sq = voxel_radius * voxel_radius;
137
138            // Collect R-tree neighbor IDs first, then look up distances.
139            let neighbor_ids: Vec<SystemId> = spatial
140                .nearest_neighbor_iter(&query_point)
141                .take_while(|sp| sp.distance_2(&query_point) <= voxel_radius_sq)
142                .filter(|sp| sp.id != sys_id)
143                .map(|sp| sp.id)
144                .collect();
145
146            let neighbors: Vec<_> = neighbor_ids
147                .into_iter()
148                .filter_map(|nid| {
149                    let neighbor_sys = self.systems.get(&nid)?;
150                    let dist_ly = sys_addr.distance_ly(&neighbor_sys.address);
151                    if dist_ly <= max_ly {
152                        Some((nid, dist_ly))
153                    } else {
154                        None
155                    }
156                })
157                .collect();
158
159            let from_node = match self.node_map.get(&sys_id) {
160                Some(&n) => n,
161                None => continue,
162            };
163
164            for (neighbor_id, dist_ly) in neighbors {
165                let to_node = match self.node_map.get(&neighbor_id) {
166                    Some(&n) => n,
167                    None => continue,
168                };
169
170                if self.graph.find_edge(from_node, to_node).is_none() {
171                    self.graph.add_edge(from_node, to_node, dist_ly);
172                }
173            }
174        }
175    }
176
177    /// Add KNN edges for a single newly inserted system.
178    ///
179    /// Call this after `insert_system()` to connect the new node to
180    /// its neighbors without rebuilding the entire graph.
181    /// Looks up the system's galaxy to use the correct spatial index.
182    pub fn connect_new_system(&mut self, sys_id: SystemId, k: usize) {
183        let system = match self.systems.get(&sys_id) {
184            Some(s) => s,
185            None => return,
186        };
187
188        let sys_addr = system.address;
189        let galaxy = sys_addr.reality_index;
190        let spatial = match self.spatial.get(&galaxy) {
191            Some(s) => s,
192            None => return,
193        };
194
195        let query_point = [
196            sys_addr.voxel_x() as f64,
197            sys_addr.voxel_y() as f64,
198            sys_addr.voxel_z() as f64,
199        ];
200
201        // Collect R-tree neighbor IDs first, then look up distances.
202        let neighbor_ids: Vec<SystemId> = spatial
203            .nearest_neighbor_iter(&query_point)
204            .filter(|sp| sp.id != sys_id)
205            .take(k)
206            .map(|sp| sp.id)
207            .collect();
208
209        let neighbors: Vec<_> = neighbor_ids
210            .into_iter()
211            .filter_map(|nid| {
212                let neighbor_sys = self.systems.get(&nid)?;
213                let dist_ly = sys_addr.distance_ly(&neighbor_sys.address);
214                Some((nid, dist_ly))
215            })
216            .collect();
217
218        let from_node = match self.node_map.get(&sys_id) {
219            Some(&n) => n,
220            None => return,
221        };
222
223        for (neighbor_id, dist_ly) in neighbors {
224            let to_node = match self.node_map.get(&neighbor_id) {
225                Some(&n) => n,
226                None => continue,
227            };
228
229            if self.graph.find_edge(from_node, to_node).is_none() {
230                self.graph.add_edge(from_node, to_node, dist_ly);
231            }
232        }
233    }
234}
235
236#[cfg(test)]
237mod tests {
238    use super::*;
239    use nms_core::address::GalacticAddress;
240    use nms_core::biome::Biome;
241    use nms_core::system::{Planet, System};
242
243    /// Build a model with N systems in a line along the X axis.
244    fn line_model(n: usize, spacing: i16) -> GalaxyModel {
245        let json = r#"{
246            "Version": 4720, "Platform": "Mac|Final", "ActiveContext": "Main",
247            "CommonStateData": {"SaveName": "Test", "TotalPlayTime": 100},
248            "BaseContext": {"GameMode": 1, "PlayerStateData": {"UniverseAddress": {"RealityIndex": 0, "GalacticAddress": {"VoxelX": 0, "VoxelY": 0, "VoxelZ": 0, "SolarSystemIndex": 0, "PlanetIndex": 0}}, "Units": 0, "Nanites": 0, "Specials": 0, "PersistentPlayerBases": []}},
249            "ExpeditionContext": {"GameMode": 6, "PlayerStateData": {"UniverseAddress": {"RealityIndex": 0, "GalacticAddress": {"VoxelX": 0, "VoxelY": 0, "VoxelZ": 0, "SolarSystemIndex": 0, "PlanetIndex": 0}}, "Units": 0, "Nanites": 0, "Specials": 0, "PersistentPlayerBases": []}},
250            "DiscoveryManagerData": {"DiscoveryData-v1": {"ReserveStore": 0, "ReserveManaged": 0, "Store": {"Record": []}}}
251        }"#;
252        let save = nms_save::parse_save(json.as_bytes()).unwrap();
253        let mut model = GalaxyModel::from_save(&save);
254
255        for i in 0..n {
256            let x = (i as i16) * spacing;
257            let ssi = (i + 1) as u16;
258            let addr = GalacticAddress::new(x, 0, 0, ssi, 0, 0);
259            let planet = Planet::new(0, Some(Biome::Barren), None, false, None, None);
260            let system = System::new(addr, None, None, None, vec![planet]);
261            model.insert_system(system);
262        }
263
264        model
265    }
266
267    #[test]
268    fn test_knn_edges_created() {
269        let mut model = line_model(5, 10);
270        model.build_edges(EdgeStrategy::Knn { k: 2 });
271        assert!(model.graph.edge_count() > 0);
272        assert!(model.graph.edge_count() <= 10);
273    }
274
275    #[test]
276    fn test_knn_edges_all_nodes_connected() {
277        let mut model = line_model(5, 10);
278        model.build_edges(EdgeStrategy::Knn { k: 2 });
279        for &node_idx in model.node_map.values() {
280            let degree = model.graph.edges(node_idx).count();
281            assert!(degree >= 1, "Node with 0 edges found");
282        }
283    }
284
285    #[test]
286    fn test_warp_range_edges_respects_distance() {
287        let mut model = line_model(5, 10);
288        // 10 voxels * 400 = 4000 ly. Range 5000 should connect adjacent only.
289        model.build_edges(EdgeStrategy::WarpRange { max_ly: 5000.0 });
290        for edge in model.graph.edge_indices() {
291            let weight = model.graph[edge];
292            assert!(weight <= 5000.0, "Edge weight {weight} exceeds warp range");
293        }
294    }
295
296    #[test]
297    fn test_warp_range_zero_no_edges() {
298        let mut model = line_model(5, 10);
299        model.build_edges(EdgeStrategy::WarpRange { max_ly: 0.0 });
300        assert_eq!(model.graph.edge_count(), 0);
301    }
302
303    #[test]
304    fn test_build_edges_clears_previous() {
305        let mut model = line_model(5, 10);
306        model.build_edges(EdgeStrategy::Knn { k: 4 });
307        let count1 = model.graph.edge_count();
308        model.build_edges(EdgeStrategy::Knn { k: 1 });
309        let count2 = model.graph.edge_count();
310        assert!(count2 < count1);
311    }
312
313    #[test]
314    fn test_connect_new_system_adds_edges() {
315        let mut model = line_model(3, 10);
316        model.build_edges(EdgeStrategy::Knn { k: 2 });
317        let edges_before = model.graph.edge_count();
318
319        // Insert a new system
320        let addr = GalacticAddress::new(5, 0, 0, 0xFFF, 0, 0);
321        let system = System::new(addr, None, None, None, vec![]);
322        let sys_id = crate::spatial::SystemId::from_address(&addr);
323        model.insert_system(system);
324        model.connect_new_system(sys_id, 2);
325
326        assert!(model.graph.edge_count() > edges_before);
327    }
328
329    #[test]
330    fn test_no_self_loops() {
331        let mut model = line_model(5, 10);
332        model.build_edges(EdgeStrategy::Knn { k: 4 });
333        for edge in model.graph.edge_indices() {
334            let (a, b) = model.graph.edge_endpoints(edge).unwrap();
335            assert_ne!(a, b, "Self-loop detected");
336        }
337    }
338
339    #[test]
340    fn test_no_duplicate_edges() {
341        let mut model = line_model(5, 10);
342        model.build_edges(EdgeStrategy::Knn { k: 4 });
343        let mut seen = std::collections::HashSet::new();
344        for edge in model.graph.edge_indices() {
345            let (a, b) = model.graph.edge_endpoints(edge).unwrap();
346            let key = if a < b { (a, b) } else { (b, a) };
347            assert!(seen.insert(key), "Duplicate edge found: {key:?}");
348        }
349    }
350
351    #[test]
352    fn test_edge_strategy_default_is_knn_10() {
353        let strategy = EdgeStrategy::default();
354        match strategy {
355            EdgeStrategy::Knn { k } => assert_eq!(k, 10),
356            _ => panic!("Default should be Knn"),
357        }
358    }
359
360    #[test]
361    fn test_warp_range_large_connects_all() {
362        let mut model = line_model(3, 10);
363        // 3 systems, spacing 10 voxels = 4000 ly each. Range 100000 ly covers all.
364        model.build_edges(EdgeStrategy::WarpRange { max_ly: 100_000.0 });
365        // 3 systems fully connected = 3 edges
366        assert_eq!(model.graph.edge_count(), 3);
367    }
368
369    #[test]
370    fn test_connect_new_system_nonexistent_is_noop() {
371        let mut model = line_model(3, 10);
372        model.build_edges(EdgeStrategy::Knn { k: 2 });
373        let edges_before = model.graph.edge_count();
374
375        // Try connecting a system that doesn't exist
376        model.connect_new_system(SystemId(0xDEADBEEF), 2);
377
378        assert_eq!(model.graph.edge_count(), edges_before);
379    }
380
381    #[test]
382    fn test_edge_weights_are_positive() {
383        let mut model = line_model(5, 10);
384        model.build_edges(EdgeStrategy::Knn { k: 2 });
385        for edge in model.graph.edge_indices() {
386            let weight = model.graph[edge];
387            assert!(weight > 0.0, "Edge weight should be positive, got {weight}");
388        }
389    }
390}