1use std::collections::BTreeMap;
2
3use crate::compiler::prelude::*;
4use crate::value;
5
6use super::util::round_to_precision;
7
8const EARTH_RADIUS_IN_METERS: f64 = 6_371_008.8;
9const EARTH_RADIUS_IN_KILOMETERS: f64 = EARTH_RADIUS_IN_METERS / 1000.0;
10const EARTH_RADIUS_IN_MILES: f64 = EARTH_RADIUS_IN_KILOMETERS * 0.621_371_2;
11
12fn haversine_distance(
13 latitude1: Value,
14 longitude1: Value,
15 latitude2: Value,
16 longitude2: Value,
17 measurement_unit: &MeasurementUnit,
18) -> Resolved {
19 let latitude1 = latitude1.try_float()?.to_radians();
20 let longitude1 = longitude1.try_float()?.to_radians();
21 let latitude2 = latitude2.try_float()?.to_radians();
22 let longitude2 = longitude2.try_float()?.to_radians();
23
24 let mut result = ObjectMap::new();
25
26 let dlon = longitude2 - longitude1;
28 let dlat = latitude2 - latitude1;
29 let a =
30 (dlat / 2.0).sin().powi(2) + latitude1.cos() * latitude2.cos() * (dlon / 2.0).sin().powi(2);
31 let distance = 2.0 * a.sqrt().asin();
32
33 result.insert(
34 "distance".into(),
35 match measurement_unit {
36 MeasurementUnit::Kilometers => Value::from_f64_or_zero(round_to_precision(
37 distance * EARTH_RADIUS_IN_KILOMETERS,
38 7,
39 f64::round,
40 )),
41 MeasurementUnit::Miles => Value::from_f64_or_zero(round_to_precision(
42 distance * EARTH_RADIUS_IN_MILES,
43 7,
44 f64::round,
45 )),
46 },
47 );
48
49 let y = dlon.sin() * latitude2.cos();
51 let x = latitude1.cos() * latitude2.sin() - latitude1.sin() * latitude2.cos() * dlon.cos();
52 let bearing = (y.atan2(x).to_degrees() + 360.0) % 360.0;
53
54 result.insert(
55 "bearing".into(),
56 Value::from_f64_or_zero(round_to_precision(bearing, 3, f64::round)),
57 );
58
59 Ok(result.into())
60}
61
62fn measurement_systems() -> Vec<Value> {
63 vec![value!("kilometers"), value!("miles")]
64}
65
66#[derive(Clone, Debug)]
67enum MeasurementUnit {
68 Kilometers,
69 Miles,
70}
71
72#[derive(Clone, Copy, Debug)]
73pub struct Haversine;
74
75impl Function for Haversine {
76 fn identifier(&self) -> &'static str {
77 "haversine"
78 }
79
80 fn usage(&self) -> &'static str {
81 "Calculates [haversine](https://en.wikipedia.org/wiki/Haversine_formula) distance and bearing between two points."
82 }
83
84 fn parameters(&self) -> &'static [Parameter] {
85 &[
86 Parameter {
87 keyword: "latitude1",
88 kind: kind::FLOAT,
89 required: true,
90 },
91 Parameter {
92 keyword: "longitude1",
93 kind: kind::FLOAT,
94 required: true,
95 },
96 Parameter {
97 keyword: "latitude2",
98 kind: kind::FLOAT,
99 required: true,
100 },
101 Parameter {
102 keyword: "longitude2",
103 kind: kind::FLOAT,
104 required: true,
105 },
106 Parameter {
107 keyword: "measurement_unit",
108 kind: kind::BYTES,
109 required: false,
110 },
111 ]
112 }
113
114 fn compile(
115 &self,
116 state: &state::TypeState,
117 _ctx: &mut FunctionCompileContext,
118 arguments: ArgumentList,
119 ) -> Compiled {
120 let latitude1 = arguments.required("latitude1");
121 let longitude1 = arguments.required("longitude1");
122 let latitude2 = arguments.required("latitude2");
123 let longitude2 = arguments.required("longitude2");
124
125 let measurement_unit = match arguments
126 .optional_enum("measurement_unit", &measurement_systems(), state)?
127 .unwrap_or_else(|| value!("kilometers"))
128 .try_bytes()
129 .ok()
130 .as_deref()
131 {
132 Some(b"kilometers") => MeasurementUnit::Kilometers,
133 Some(b"miles") => MeasurementUnit::Miles,
134 _ => return Err(Box::new(ExpressionError::from("invalid measurement unit"))),
135 };
136
137 Ok(HaversineFn {
138 latitude1,
139 longitude1,
140 latitude2,
141 longitude2,
142 measurement_unit,
143 }
144 .as_expr())
145 }
146
147 fn examples(&self) -> &'static [Example] {
148 &[
149 example! {
150 title: "Haversine in kilometers",
151 source: "haversine(0.0, 0.0, 10.0, 10.0)",
152 result: Ok(indoc!(
153 r#"{
154 "distance": 1568.5227233,
155 "bearing": 44.561
156 }"#
157 )),
158 },
159 example! {
160 title: "Haversine in miles",
161 source: r#"haversine(0.0, 0.0, 10.0, 10.0, measurement_unit: "miles")"#,
162 result: Ok(indoc!(
163 r#"{
164 "distance": 974.6348468,
165 "bearing": 44.561
166 }"#
167 )),
168 },
169 ]
170 }
171}
172
173#[derive(Clone, Debug)]
174struct HaversineFn {
175 latitude1: Box<dyn Expression>,
176 longitude1: Box<dyn Expression>,
177 latitude2: Box<dyn Expression>,
178 longitude2: Box<dyn Expression>,
179 measurement_unit: MeasurementUnit,
180}
181
182impl FunctionExpression for HaversineFn {
183 fn resolve(&self, ctx: &mut Context) -> Resolved {
184 let latitude1 = self.latitude1.resolve(ctx)?;
185 let longitude1 = self.longitude1.resolve(ctx)?;
186 let latitude2 = self.latitude2.resolve(ctx)?;
187 let longitude2 = self.longitude2.resolve(ctx)?;
188
189 haversine_distance(
190 latitude1,
191 longitude1,
192 latitude2,
193 longitude2,
194 &self.measurement_unit,
195 )
196 }
197
198 fn type_def(&self, _state: &state::TypeState) -> TypeDef {
199 TypeDef::object(inner_kind()).infallible()
200 }
201}
202
203fn inner_kind() -> BTreeMap<Field, Kind> {
204 BTreeMap::from([
205 (Field::from("distance"), Kind::float()),
206 (Field::from("bearing"), Kind::float()),
207 ])
208}
209
210#[cfg(test)]
211mod tests {
212 use super::*;
213 use crate::value;
214
215 test_function![
216 haversine => Haversine;
217
218 basic_kilometers {
219 args: func_args![latitude1: value!(0.0), longitude1: value!(0.0), latitude2: value!(10.0), longitude2: value!(10.0)],
220 want: Ok(value!({ "distance": 1_568.522_723_3, "bearing": 44.561 })),
221 tdef: TypeDef::object(inner_kind()).infallible(),
222 }
223
224 basic_miles {
225 args: func_args![latitude1: value!(0.0), longitude1: value!(0.0), latitude2: value!(10.0), longitude2: value!(10.0), measurement_unit: value!("miles")],
226 want: Ok(value!({ "distance": 974.634_846_8, "bearing": 44.561 })),
227 tdef: TypeDef::object(inner_kind()).infallible(),
228 }
229 ];
230}