nms_graph/
query.rs

1//! Spatial query methods on the GalaxyModel.
2
3use nms_core::address::GalacticAddress;
4use nms_core::biome::Biome;
5use nms_core::system::Planet;
6use rstar::PointDistance;
7
8use crate::error::GraphError;
9use crate::model::{GalaxyModel, PlanetKey};
10use crate::spatial::SystemId;
11
12/// Filter criteria for planet queries.
13#[derive(Debug, Clone, Default)]
14pub struct BiomeFilter {
15    pub biome: Option<Biome>,
16    pub infested: Option<bool>,
17    pub named_only: bool,
18}
19
20impl GalaxyModel {
21    /// Resolve a reference point to a `GalacticAddress`.
22    ///
23    /// Accepts:
24    /// - A direct address (returned as-is)
25    /// - A base name (looked up in the base index)
26    /// - None for both (uses player position)
27    pub fn resolve_position(
28        &self,
29        address: Option<&GalacticAddress>,
30        base_name: Option<&str>,
31    ) -> Result<GalacticAddress, GraphError> {
32        if let Some(addr) = address {
33            return Ok(*addr);
34        }
35        if let Some(name) = base_name {
36            return self
37                .base(name)
38                .map(|b| b.address)
39                .ok_or_else(|| GraphError::BaseNotFound(name.to_string()));
40        }
41        self.player_position()
42            .copied()
43            .ok_or(GraphError::NoPlayerPosition)
44    }
45
46    /// Find the N nearest systems to a reference point.
47    ///
48    /// Returns `(SystemId, distance_in_ly)` pairs sorted by distance ascending.
49    pub fn nearest_systems(&self, from: &GalacticAddress, n: usize) -> Vec<(SystemId, f64)> {
50        let query_point = [
51            from.voxel_x() as f64,
52            from.voxel_y() as f64,
53            from.voxel_z() as f64,
54        ];
55
56        self.spatial
57            .nearest_neighbor_iter(&query_point)
58            .take(n)
59            .map(|sp| {
60                let voxel_dist_sq = sp.distance_2(&query_point);
61                let ly = voxel_dist_sq.sqrt() * 400.0;
62                (sp.id, ly)
63            })
64            .collect()
65    }
66
67    /// Find all systems within a radius (in light-years) of a reference point.
68    ///
69    /// Returns `(SystemId, distance_in_ly)` pairs sorted by distance ascending.
70    pub fn systems_within_radius(
71        &self,
72        from: &GalacticAddress,
73        radius_ly: f64,
74    ) -> Vec<(SystemId, f64)> {
75        let query_point = [
76            from.voxel_x() as f64,
77            from.voxel_y() as f64,
78            from.voxel_z() as f64,
79        ];
80        let voxel_radius = radius_ly / 400.0;
81        let voxel_radius_sq = voxel_radius * voxel_radius;
82
83        let mut results: Vec<(SystemId, f64)> = self
84            .spatial
85            .nearest_neighbor_iter(&query_point)
86            .map(|sp| {
87                let dist_sq = sp.distance_2(&query_point);
88                (sp.id, dist_sq)
89            })
90            .take_while(|&(_, dist_sq)| dist_sq <= voxel_radius_sq)
91            .map(|(id, dist_sq)| (id, dist_sq.sqrt() * 400.0))
92            .collect();
93
94        results.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());
95        results
96    }
97
98    /// Find the N nearest planets to a reference point, with optional filtering.
99    ///
100    /// Iterates systems by proximity, then checks their planets against the filter.
101    /// Returns `(PlanetKey, &Planet, system_distance_ly)` tuples.
102    pub fn nearest_planets<'a>(
103        &'a self,
104        from: &GalacticAddress,
105        n: usize,
106        filter: &BiomeFilter,
107    ) -> Vec<(PlanetKey, &'a Planet, f64)> {
108        let query_point = [
109            from.voxel_x() as f64,
110            from.voxel_y() as f64,
111            from.voxel_z() as f64,
112        ];
113
114        let mut results = Vec::with_capacity(n);
115
116        for sp in self.spatial.nearest_neighbor_iter(&query_point) {
117            if results.len() >= n {
118                break;
119            }
120
121            let dist_ly = sp.distance_2(&query_point).sqrt() * 400.0;
122
123            // Get all planets in this system
124            if let Some(system) = self.systems.get(&sp.id) {
125                for planet in &system.planets {
126                    if results.len() >= n {
127                        break;
128                    }
129                    if matches_filter(planet, filter) {
130                        let key = (sp.id, planet.index);
131                        results.push((key, planet, dist_ly));
132                    }
133                }
134            }
135        }
136
137        results
138    }
139
140    /// Find all planets within a radius that match a filter.
141    pub fn planets_within_radius<'a>(
142        &'a self,
143        from: &GalacticAddress,
144        radius_ly: f64,
145        filter: &BiomeFilter,
146    ) -> Vec<(PlanetKey, &'a Planet, f64)> {
147        let systems = self.systems_within_radius(from, radius_ly);
148        let mut results = Vec::new();
149
150        for (sys_id, dist_ly) in systems {
151            if let Some(system) = self.systems.get(&sys_id) {
152                for planet in &system.planets {
153                    if matches_filter(planet, filter) {
154                        let key = (sys_id, planet.index);
155                        results.push((key, planet, dist_ly));
156                    }
157                }
158            }
159        }
160
161        results
162    }
163}
164
165/// Check if a planet matches the given filter criteria.
166fn matches_filter(planet: &Planet, filter: &BiomeFilter) -> bool {
167    if let Some(biome) = filter.biome {
168        if planet.biome != Some(biome) {
169            return false;
170        }
171    }
172    if let Some(infested) = filter.infested {
173        if planet.infested != infested {
174            return false;
175        }
176    }
177    if filter.named_only && planet.name.is_none() {
178        return false;
179    }
180    true
181}
182
183#[cfg(test)]
184mod tests {
185    use super::*;
186    use nms_core::address::GalacticAddress;
187    use nms_core::biome::Biome;
188    use nms_core::system::{Planet, System};
189
190    /// Build a model with systems at known positions for spatial testing.
191    fn spatial_test_model() -> GalaxyModel {
192        let json = r#"{
193            "Version": 4720, "Platform": "Mac|Final", "ActiveContext": "Main",
194            "CommonStateData": {"SaveName": "Test", "TotalPlayTime": 100},
195            "BaseContext": {
196                "GameMode": 1,
197                "PlayerStateData": {
198                    "UniverseAddress": {"RealityIndex": 0, "GalacticAddress": {"VoxelX": 0, "VoxelY": 0, "VoxelZ": 0, "SolarSystemIndex": 1, "PlanetIndex": 0}},
199                    "Units": 0, "Nanites": 0, "Specials": 0,
200                    "PersistentPlayerBases": [
201                        {"BaseVersion": 8, "GalacticAddress": "0x001000000064", "Position": [0.0,0.0,0.0], "Forward": [1.0,0.0,0.0], "LastUpdateTimestamp": 0, "Objects": [], "RID": "", "Owner": {"LID":"","UID":"1","USN":"","PTK":"ST","TS":0}, "Name": "Test Base", "BaseType": {"PersistentBaseTypes": "HomePlanetBase"}, "LastEditedById": "", "LastEditedByUsername": ""}
202                    ]
203                }
204            },
205            "ExpeditionContext": {
206                "GameMode": 6,
207                "PlayerStateData": {
208                    "UniverseAddress": {"RealityIndex": 0, "GalacticAddress": {"VoxelX": 0, "VoxelY": 0, "VoxelZ": 0, "SolarSystemIndex": 0, "PlanetIndex": 0}},
209                    "Units": 0, "Nanites": 0, "Specials": 0, "PersistentPlayerBases": []
210                }
211            },
212            "DiscoveryManagerData": {"DiscoveryData-v1": {"ReserveStore": 0, "ReserveManaged": 0, "Store": {"Record": []}}}
213        }"#;
214        let save = nms_save::parse_save(json.as_bytes()).unwrap();
215        let mut model = GalaxyModel::from_save(&save);
216
217        // Insert systems at known positions along the X axis
218        let positions = [
219            (10, 0, 0, 0x100, "Near"), // 10 voxels = 4000 ly
220            (50, 0, 0, 0x200, "Mid"),  // 50 voxels = 20000 ly
221            (200, 0, 0, 0x300, "Far"), // 200 voxels = 80000 ly
222        ];
223
224        for (x, y, z, ssi, name) in positions {
225            let addr = GalacticAddress::new(x, y, z, ssi, 0, 0);
226            let planet = Planet::new(0, Some(Biome::Lush), None, false, None, None);
227            let system = System::new(addr, Some(name.into()), None, None, vec![planet]);
228            model.insert_system(system);
229        }
230
231        // Add a Scorched infested planet to "Near"
232        let near_addr = GalacticAddress::new(10, 0, 0, 0x100, 1, 0);
233        let near_id = crate::spatial::SystemId::from_address(&near_addr);
234        let scorched = Planet::new(1, Some(Biome::Scorched), None, true, None, None);
235        let key = (near_id, 1);
236        model.planets.insert(key, scorched.clone());
237        model
238            .biome_index
239            .entry(Biome::Scorched)
240            .or_default()
241            .push(key);
242        // Also add to the system's planet list
243        if let Some(sys) = model.systems.get_mut(&near_id) {
244            sys.planets.push(scorched);
245        }
246
247        model
248    }
249
250    #[test]
251    fn test_nearest_systems_returns_sorted() {
252        let model = spatial_test_model();
253        let origin = GalacticAddress::new(0, 0, 0, 1, 0, 0);
254        let results = model.nearest_systems(&origin, 10);
255        for i in 1..results.len() {
256            assert!(results[i].1 >= results[i - 1].1);
257        }
258    }
259
260    #[test]
261    fn test_nearest_systems_limit() {
262        let model = spatial_test_model();
263        let origin = GalacticAddress::new(0, 0, 0, 1, 0, 0);
264        let results = model.nearest_systems(&origin, 2);
265        assert_eq!(results.len(), 2);
266    }
267
268    #[test]
269    fn test_nearest_systems_distances_are_in_ly() {
270        let model = spatial_test_model();
271        let origin = GalacticAddress::new(0, 0, 0, 1, 0, 0);
272        let results = model.nearest_systems(&origin, 10);
273        // "Near" is at voxel 10 = 4000 ly; find the result with ~4000 ly
274        let near_result = results.iter().find(|(_, d)| (*d - 4000.0).abs() < 1.0);
275        assert!(near_result.is_some(), "Expected a system at ~4000 ly");
276    }
277
278    #[test]
279    fn test_systems_within_radius_filters() {
280        let model = spatial_test_model();
281        let origin = GalacticAddress::new(0, 0, 0, 1, 0, 0);
282        // 10 voxels = 4000 ly, so radius of 5000 should include "Near" but not "Mid"
283        let results = model.systems_within_radius(&origin, 5000.0);
284        assert!(!results.is_empty());
285        for (_, dist) in &results {
286            assert!(*dist <= 5000.0);
287        }
288    }
289
290    #[test]
291    fn test_systems_within_radius_excludes_far() {
292        let model = spatial_test_model();
293        let origin = GalacticAddress::new(0, 0, 0, 1, 0, 0);
294        // Radius 5000 ly should not include "Mid" (20000 ly) or "Far" (80000 ly)
295        let results = model.systems_within_radius(&origin, 5000.0);
296        for (_, dist) in &results {
297            assert!(*dist < 20000.0, "Far system should be excluded");
298        }
299    }
300
301    #[test]
302    fn test_systems_within_radius_zero() {
303        let model = spatial_test_model();
304        let origin = GalacticAddress::new(999, 127, 999, 1, 0, 0);
305        let results = model.systems_within_radius(&origin, 0.0);
306        assert!(results.is_empty());
307    }
308
309    #[test]
310    fn test_nearest_planets_with_biome_filter() {
311        let model = spatial_test_model();
312        let origin = GalacticAddress::new(0, 0, 0, 1, 0, 0);
313        let filter = BiomeFilter {
314            biome: Some(Biome::Scorched),
315            ..Default::default()
316        };
317        let results = model.nearest_planets(&origin, 10, &filter);
318        assert!(!results.is_empty());
319        for (_, planet, _) in &results {
320            assert_eq!(planet.biome, Some(Biome::Scorched));
321        }
322    }
323
324    #[test]
325    fn test_nearest_planets_with_infested_filter() {
326        let model = spatial_test_model();
327        let origin = GalacticAddress::new(0, 0, 0, 1, 0, 0);
328        let filter = BiomeFilter {
329            infested: Some(true),
330            ..Default::default()
331        };
332        let results = model.nearest_planets(&origin, 10, &filter);
333        assert!(!results.is_empty());
334        for (_, planet, _) in &results {
335            assert!(planet.infested);
336        }
337    }
338
339    #[test]
340    fn test_nearest_planets_no_filter() {
341        let model = spatial_test_model();
342        let origin = GalacticAddress::new(0, 0, 0, 1, 0, 0);
343        let filter = BiomeFilter::default();
344        let results = model.nearest_planets(&origin, 5, &filter);
345        assert!(!results.is_empty());
346        assert!(results.len() <= 5);
347    }
348
349    #[test]
350    fn test_nearest_planets_limit_respected() {
351        let model = spatial_test_model();
352        let origin = GalacticAddress::new(0, 0, 0, 1, 0, 0);
353        let filter = BiomeFilter::default();
354        let results = model.nearest_planets(&origin, 1, &filter);
355        assert_eq!(results.len(), 1);
356    }
357
358    #[test]
359    fn test_planets_within_radius() {
360        let model = spatial_test_model();
361        let origin = GalacticAddress::new(0, 0, 0, 1, 0, 0);
362        let filter = BiomeFilter {
363            biome: Some(Biome::Lush),
364            ..Default::default()
365        };
366        // 5000 ly should reach "Near" (4000 ly) but not "Mid" (20000 ly)
367        let results = model.planets_within_radius(&origin, 5000.0, &filter);
368        for (_, planet, dist) in &results {
369            assert_eq!(planet.biome, Some(Biome::Lush));
370            assert!(*dist <= 5000.0);
371        }
372    }
373
374    #[test]
375    fn test_planets_within_radius_no_match() {
376        let model = spatial_test_model();
377        let origin = GalacticAddress::new(0, 0, 0, 1, 0, 0);
378        let filter = BiomeFilter {
379            biome: Some(Biome::Lava),
380            ..Default::default()
381        };
382        let results = model.planets_within_radius(&origin, 100000.0, &filter);
383        assert!(results.is_empty());
384    }
385
386    #[test]
387    fn test_resolve_position_direct_address() {
388        let model = spatial_test_model();
389        let addr = GalacticAddress::new(42, 0, 0, 1, 0, 0);
390        let resolved = model.resolve_position(Some(&addr), None).unwrap();
391        assert_eq!(resolved, addr);
392    }
393
394    #[test]
395    fn test_resolve_position_from_base() {
396        let model = spatial_test_model();
397        let resolved = model.resolve_position(None, Some("Test Base")).unwrap();
398        // Base was inserted with address 0x001000000064
399        assert_eq!(resolved.packed(), 0x001000000064);
400    }
401
402    #[test]
403    fn test_resolve_position_base_not_found() {
404        let model = spatial_test_model();
405        let result = model.resolve_position(None, Some("No Such Base"));
406        assert!(result.is_err());
407    }
408
409    #[test]
410    fn test_resolve_position_player_position() {
411        let model = spatial_test_model();
412        let resolved = model.resolve_position(None, None).unwrap();
413        // Player is at origin (0,0,0)
414        assert_eq!(resolved.voxel_x(), 0);
415        assert_eq!(resolved.voxel_y(), 0);
416        assert_eq!(resolved.voxel_z(), 0);
417    }
418
419    #[test]
420    fn test_resolve_position_no_player_state_errors() {
421        let json = r#"{
422            "Version": 4720, "Platform": "Mac|Final", "ActiveContext": "Main",
423            "CommonStateData": {"SaveName": "Test", "TotalPlayTime": 0},
424            "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": []}},
425            "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": []}},
426            "DiscoveryManagerData": {"DiscoveryData-v1": {"ReserveStore": 0, "ReserveManaged": 0, "Store": {"Record": []}}}
427        }"#;
428        let save = nms_save::parse_save(json.as_bytes()).unwrap();
429        let mut model = GalaxyModel::from_save(&save);
430        model.player_state = None;
431        assert!(model.resolve_position(None, None).is_err());
432    }
433
434    #[test]
435    fn test_resolve_position_address_takes_priority() {
436        let model = spatial_test_model();
437        let addr = GalacticAddress::new(42, 10, -5, 0x999, 0, 0);
438        // Even with a base name, direct address wins
439        let resolved = model
440            .resolve_position(Some(&addr), Some("Test Base"))
441            .unwrap();
442        assert_eq!(resolved, addr);
443    }
444
445    #[test]
446    fn test_matches_filter_all_pass() {
447        let planet = Planet::new(0, Some(Biome::Lush), None, false, Some("Eden".into()), None);
448        let filter = BiomeFilter::default();
449        assert!(matches_filter(&planet, &filter));
450    }
451
452    #[test]
453    fn test_matches_filter_biome_mismatch() {
454        let planet = Planet::new(0, Some(Biome::Lush), None, false, None, None);
455        let filter = BiomeFilter {
456            biome: Some(Biome::Toxic),
457            ..Default::default()
458        };
459        assert!(!matches_filter(&planet, &filter));
460    }
461
462    #[test]
463    fn test_matches_filter_biome_match() {
464        let planet = Planet::new(0, Some(Biome::Toxic), None, false, None, None);
465        let filter = BiomeFilter {
466            biome: Some(Biome::Toxic),
467            ..Default::default()
468        };
469        assert!(matches_filter(&planet, &filter));
470    }
471
472    #[test]
473    fn test_matches_filter_infested_mismatch() {
474        let planet = Planet::new(0, Some(Biome::Lush), None, false, None, None);
475        let filter = BiomeFilter {
476            infested: Some(true),
477            ..Default::default()
478        };
479        assert!(!matches_filter(&planet, &filter));
480    }
481
482    #[test]
483    fn test_matches_filter_named_only() {
484        let unnamed = Planet::new(0, Some(Biome::Lush), None, false, None, None);
485        let named = Planet::new(0, Some(Biome::Lush), None, false, Some("X".into()), None);
486        let filter = BiomeFilter {
487            named_only: true,
488            ..Default::default()
489        };
490        assert!(!matches_filter(&unnamed, &filter));
491        assert!(matches_filter(&named, &filter));
492    }
493
494    #[test]
495    fn test_matches_filter_combined() {
496        let planet = Planet::new(
497            0,
498            Some(Biome::Scorched),
499            None,
500            true,
501            Some("Inferno".into()),
502            None,
503        );
504        let filter = BiomeFilter {
505            biome: Some(Biome::Scorched),
506            infested: Some(true),
507            named_only: true,
508        };
509        assert!(matches_filter(&planet, &filter));
510    }
511}
nms_graph/query.rs

nms_graph/
query.rs
