Skip to main content

jpx_core/extensions/
geo.rs

1//! Geographic/geospatial functions.
2
3use std::collections::HashSet;
4
5use geoutils::Location;
6use serde_json::Value;
7
8use crate::functions::{Function, number_value};
9use crate::interpreter::SearchResult;
10use crate::registry::register_if_enabled;
11use crate::{Context, Runtime, arg, defn};
12
13// =============================================================================
14// geo_distance(lat1, lon1, lat2, lon2) -> number (meters)
15// =============================================================================
16
17defn!(
18    GeoDistanceFn,
19    vec![arg!(number), arg!(number), arg!(number), arg!(number)],
20    None
21);
22
23impl Function for GeoDistanceFn {
24    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
25        self.signature.validate(args, ctx)?;
26        let lat1 = args[0].as_f64().unwrap();
27        let lon1 = args[1].as_f64().unwrap();
28        let lat2 = args[2].as_f64().unwrap();
29        let lon2 = args[3].as_f64().unwrap();
30
31        let loc1 = Location::new(lat1, lon1);
32        let loc2 = Location::new(lat2, lon2);
33
34        let distance = loc1.haversine_distance_to(&loc2);
35        Ok(number_value(distance.meters()))
36    }
37}
38
39// =============================================================================
40// geo_distance_km(lat1, lon1, lat2, lon2) -> number (kilometers)
41// =============================================================================
42
43defn!(
44    GeoDistanceKmFn,
45    vec![arg!(number), arg!(number), arg!(number), arg!(number)],
46    None
47);
48
49impl Function for GeoDistanceKmFn {
50    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
51        self.signature.validate(args, ctx)?;
52        let lat1 = args[0].as_f64().unwrap();
53        let lon1 = args[1].as_f64().unwrap();
54        let lat2 = args[2].as_f64().unwrap();
55        let lon2 = args[3].as_f64().unwrap();
56
57        let loc1 = Location::new(lat1, lon1);
58        let loc2 = Location::new(lat2, lon2);
59
60        let distance = loc1.haversine_distance_to(&loc2);
61        Ok(number_value(distance.meters() / 1000.0))
62    }
63}
64
65// =============================================================================
66// geo_distance_miles(lat1, lon1, lat2, lon2) -> number (miles)
67// =============================================================================
68
69defn!(
70    GeoDistanceMilesFn,
71    vec![arg!(number), arg!(number), arg!(number), arg!(number)],
72    None
73);
74
75impl Function for GeoDistanceMilesFn {
76    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
77        self.signature.validate(args, ctx)?;
78        let lat1 = args[0].as_f64().unwrap();
79        let lon1 = args[1].as_f64().unwrap();
80        let lat2 = args[2].as_f64().unwrap();
81        let lon2 = args[3].as_f64().unwrap();
82
83        let loc1 = Location::new(lat1, lon1);
84        let loc2 = Location::new(lat2, lon2);
85
86        // 1 meter = 0.000621371 miles
87        const METERS_TO_MILES: f64 = 0.000621371;
88
89        let distance = loc1.haversine_distance_to(&loc2);
90        Ok(number_value(distance.meters() * METERS_TO_MILES))
91    }
92}
93
94// =============================================================================
95// geo_bearing(lat1, lon1, lat2, lon2) -> number (degrees 0-360)
96// =============================================================================
97
98defn!(
99    GeoBearingFn,
100    vec![arg!(number), arg!(number), arg!(number), arg!(number)],
101    None
102);
103
104impl Function for GeoBearingFn {
105    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
106        self.signature.validate(args, ctx)?;
107        let lat1 = args[0].as_f64().unwrap();
108        let lon1 = args[1].as_f64().unwrap();
109        let lat2 = args[2].as_f64().unwrap();
110        let lon2 = args[3].as_f64().unwrap();
111
112        // Calculate initial bearing using the forward azimuth formula
113        let lat1_rad = lat1.to_radians();
114        let lat2_rad = lat2.to_radians();
115        let delta_lon = (lon2 - lon1).to_radians();
116
117        let x = delta_lon.sin() * lat2_rad.cos();
118        let y = lat1_rad.cos() * lat2_rad.sin() - lat1_rad.sin() * lat2_rad.cos() * delta_lon.cos();
119
120        let bearing_rad = x.atan2(y);
121        let mut bearing = bearing_rad.to_degrees();
122
123        // Normalize to 0-360
124        if bearing < 0.0 {
125            bearing += 360.0;
126        }
127
128        Ok(number_value(bearing))
129    }
130}
131
132/// Register geo functions filtered by the enabled set.
133pub fn register_filtered(runtime: &mut Runtime, enabled: &HashSet<&str>) {
134    register_if_enabled(
135        runtime,
136        "geo_distance",
137        enabled,
138        Box::new(GeoDistanceFn::new()),
139    );
140    register_if_enabled(
141        runtime,
142        "geo_distance_km",
143        enabled,
144        Box::new(GeoDistanceKmFn::new()),
145    );
146    register_if_enabled(
147        runtime,
148        "geo_distance_miles",
149        enabled,
150        Box::new(GeoDistanceMilesFn::new()),
151    );
152    register_if_enabled(
153        runtime,
154        "geo_bearing",
155        enabled,
156        Box::new(GeoBearingFn::new()),
157    );
158}