1use yo_common::num::DIGITS_MAX;
43use yo_common::{Code, Error, Result};
44
45use crate::elem::Elements;
46use crate::geo::{self, Kind, Shape, Unit};
47use crate::keyspace::Keyspace;
48use crate::strings;
49use crate::zset::{Bound, Zset};
50use crate::zsets::{ZAdd, member_bytes};
51
52const NO_MEMBER: &str = "could not decode requested zset member";
55
56#[must_use]
62pub fn out_of_range(lon: f64, lat: f64) -> Error {
63 yo_alloc::allow(|| {
64 Error::fmt(
65 Code::Invalid,
66 format_args!("invalid longitude,latitude pair {lon:.6},{lat:.6}"),
67 )
68 })
69}
70
71#[must_use]
73pub fn no_member() -> Error {
74 Error::new(Code::Invalid, NO_MEMBER)
75}
76
77#[derive(Debug, Clone, Copy, PartialEq, Eq)]
79pub enum Sort {
80 Near,
82 Far,
84}
85
86#[derive(Debug, Clone, Copy, Default)]
88pub struct Limit {
89 pub sort: Option<Sort>,
92 pub count: Option<usize>,
94 pub any: bool,
97}
98
99impl Limit {
100 #[must_use]
106 fn ordering(&self) -> Option<Sort> {
107 match self.sort {
108 Some(s) => Some(s),
109 None if self.count.is_some() && !self.any => Some(Sort::Near),
110 None => None,
111 }
112 }
113
114 #[must_use]
116 fn cap(&self) -> Option<usize> {
117 self.any.then_some(self.count).flatten()
118 }
119}
120
121#[derive(Debug, Clone, Copy)]
123pub struct Hit {
124 at: usize,
126 len: usize,
128 pub score: u64,
130 pub lon: f64,
132 pub lat: f64,
134 pub metres: f64,
136}
137
138#[derive(Debug, Default)]
145pub struct Scratch {
146 hits: Vec<Hit>,
149 names: Vec<u8>,
151}
152
153impl Scratch {
154 pub fn iter(&self) -> impl Iterator<Item = (&[u8], &Hit)> {
156 self.hits
157 .iter()
158 .map(|h| (&self.names[h.at..h.at + h.len], h))
159 }
160
161 #[must_use]
163 pub fn len(&self) -> usize {
164 self.hits.len()
165 }
166
167 #[must_use]
169 pub fn is_empty(&self) -> bool {
170 self.hits.is_empty()
171 }
172
173 fn clear(&mut self) {
175 self.hits.clear();
176 self.names.clear();
177 }
178
179 fn push(&mut self, name: &[u8], score: u64, lon: f64, lat: f64, metres: f64) {
181 let at = self.names.len();
182 self.names.extend_from_slice(name);
183 self.hits.push(Hit {
184 at,
185 len: name.len(),
186 score,
187 lon,
188 lat,
189 metres,
190 });
191 }
192
193 fn order(&mut self, limit: Limit) {
201 let Some(sort) = limit.ordering() else {
202 self.hits.truncate(limit.count.unwrap_or(usize::MAX));
203 return;
204 };
205 let near = |a: &Hit, b: &Hit| a.metres.total_cmp(&b.metres);
206 let far = |a: &Hit, b: &Hit| b.metres.total_cmp(&a.metres);
207 let want = limit.count.unwrap_or(self.hits.len()).min(self.hits.len());
208 if want < self.hits.len() {
209 match sort {
210 Sort::Near => self.hits.select_nth_unstable_by(want, near),
211 Sort::Far => self.hits.select_nth_unstable_by(want, far),
212 };
213 self.hits.truncate(want);
214 }
215 match sort {
216 Sort::Near => self.hits.sort_unstable_by(near),
217 Sort::Far => self.hits.sort_unstable_by(far),
218 }
219 }
220}
221
222impl Keyspace {
223 pub fn geoadd<'m, I>(&mut self, key: &[u8], points: I, opts: ZAdd) -> Result<usize>
234 where
235 I: Iterator<Item = (f64, f64, &'m [u8])> + Clone,
236 {
237 for (lon, lat, member) in points.clone() {
238 strings::check_len(key, member.len())?;
239 if geo::score(lon, lat).is_none() {
240 return Err(out_of_range(lon, lat));
241 }
242 }
243 self.zadd(
244 key,
245 points.map(|(lon, lat, m)| {
246 (geo::score(lon, lat).expect("checked") as f64, m)
249 }),
250 opts,
251 )
252 }
253
254 pub fn geopos<'m, F>(
261 &mut self,
262 key: &[u8],
263 members: impl Iterator<Item = &'m [u8]>,
264 mut f: F,
265 ) -> Result<()>
266 where
267 F: FnMut(Option<(f64, f64)>),
268 {
269 let Some(at) = self.zset_slot(key)? else {
270 members.for_each(|_| f(None));
271 return Ok(());
272 };
273 let z = self.zset_at(at);
274 for m in members {
275 f(z.score(m).and_then(geo::decode));
276 }
277 Ok(())
278 }
279
280 pub fn geohash<'m, F>(
282 &mut self,
283 key: &[u8],
284 members: impl Iterator<Item = &'m [u8]>,
285 mut f: F,
286 ) -> Result<()>
287 where
288 F: FnMut(Option<&[u8]>),
289 {
290 let Some(at) = self.zset_slot(key)? else {
291 members.for_each(|_| f(None));
292 return Ok(());
293 };
294 let z = self.zset_at(at);
295 for m in members {
296 let text = z
297 .score(m)
298 .and_then(geo::decode)
299 .and_then(|(lon, lat)| geo::geohash(lon, lat));
300 match text {
301 Some(bytes) => f(Some(&bytes)),
302 None => f(None),
303 }
304 }
305 Ok(())
306 }
307
308 pub fn geodist(&mut self, key: &[u8], a: &[u8], b: &[u8]) -> Result<Option<f64>> {
316 let Some(at) = self.zset_slot(key)? else {
317 return Ok(None);
318 };
319 let z = self.zset_at(at);
320 let (Some(sa), Some(sb)) = (z.score(a), z.score(b)) else {
321 return Ok(None);
322 };
323 let (Some(pa), Some(pb)) = (geo::decode(sa), geo::decode(sb)) else {
324 return Ok(None);
325 };
326 Ok(Some(geo::distance(pa.0, pa.1, pb.0, pb.1)))
327 }
328
329 pub fn geocentre(&mut self, key: &[u8], member: &[u8]) -> Result<Option<(f64, f64)>> {
335 let Some(at) = self.zset_slot(key)? else {
336 return Ok(None);
337 };
338 match self.zset_at(at).score(member).and_then(geo::decode) {
339 Some(xy) => Ok(Some(xy)),
340 None => Err(no_member()),
341 }
342 }
343
344 pub fn geosearch(&mut self, key: &[u8], shape: &Shape, limit: Limit) -> Result<usize> {
356 let mut found = std::mem::take(&mut self.geo);
357 found.clear();
358 let outcome = match self.zset_slot(key) {
359 Err(e) => Err(e),
360 Ok(None) => Ok(()),
361 Ok(Some(at)) => {
362 collect(self.zset_at(at), shape, limit, &mut found);
363 Ok(())
364 }
365 };
366 found.order(limit);
367 let n = found.hits.len();
368 self.geo = found;
369 outcome.map(|()| n)
370 }
371
372 #[must_use]
374 pub fn geohits(&self) -> &Scratch {
375 &self.geo
376 }
377
378 pub fn geosearchstore(
390 &mut self,
391 dest: &[u8],
392 src: &[u8],
393 shape: &Shape,
394 limit: Limit,
395 dist: bool,
396 ) -> Result<usize> {
397 let n = self.geosearch(src, shape, limit)?;
398 let found = std::mem::take(&mut self.geo);
399 let mut got = Elements::with_capacity(n.max(16));
400 for (name, hit) in found.iter() {
401 let score = if dist {
402 hit.metres / shape.unit.metres()
403 } else {
404 hit.score as f64
405 };
406 let _ = got.insert(name, score);
407 }
408 self.geo = found;
409 let limits = self.zset_limits;
410 let built = Zset::from_elements(got, &limits);
411 Ok(self.put_zset(dest, built))
412 }
413}
414
415fn collect(z: &Zset, shape: &Shape, limit: Limit, out: &mut Scratch) {
421 let search = geo::areas(shape);
422 let cap = limit.cap();
423 let mut digits = [0u8; DIGITS_MAX];
424 let mut last = 0usize;
429 for i in 0..search.boxes.len() {
430 let hash = search.boxes[i];
431 if hash.bits == 0 && hash.step == 0 {
432 continue;
433 }
434 if last != 0 && hash == search.boxes[last] {
435 continue;
436 }
437 if cap.is_some_and(|n| out.hits.len() >= n) {
438 break;
439 }
440 let (low, high) = geo::range(hash);
441 let window = z.window_by_score(Bound::closed(low as f64), Bound::open(high as f64));
442 z.walk(window.start, window.len(), false, |m, raw| {
443 if cap.is_some_and(|n| out.hits.len() >= n) {
444 return;
445 }
446 let Some((lon, lat)) = geo::decode(raw) else {
447 return;
448 };
449 let Some(metres) = shape.covers(lon, lat) else {
450 return;
451 };
452 out.push(member_bytes(m, &mut digits), raw as u64, lon, lat, metres);
453 });
454 last = i;
455 }
456}
457
458#[must_use]
460pub fn circle(lon: f64, lat: f64, radius: f64, unit: Unit) -> Shape {
461 Shape {
462 lon,
463 lat,
464 kind: Kind::Circle { radius },
465 unit,
466 }
467}
468
469#[cfg(test)]
470mod tests {
471 use super::*;
472
473 const PLACES: [(f64, f64, &[u8]); 4] = [
476 (13.361389, 38.115556, b"Palermo"),
477 (15.087269, 37.502669, b"Catania"),
478 (12.758489, 38.788135, b"edge"),
479 (2.352222, 48.856613, b"Paris"),
480 ];
481
482 fn ks() -> Keyspace {
483 let mut db = Keyspace::new();
484 let opts = ZAdd::default();
485 db.geoadd(b"g", PLACES.iter().copied(), opts)
486 .expect("the places are all in range");
487 db
488 }
489
490 fn names(db: &Keyspace) -> Vec<Vec<u8>> {
491 db.geohits().iter().map(|(n, _)| n.to_vec()).collect()
492 }
493
494 #[test]
495 fn a_geo_key_is_a_sorted_set_of_hashes() {
496 let mut db = ks();
497 assert_eq!(db.zcard(b"g").expect("a zset"), 4);
498 assert_eq!(
500 db.zscore(b"g", b"Palermo").expect("a zset"),
501 Some(3_479_099_956_230_698.0)
502 );
503 assert_eq!(
504 db.zscore(b"g", b"Catania").expect("a zset"),
505 Some(3_479_447_370_796_909.0)
506 );
507 }
508
509 #[test]
510 fn nothing_is_stored_when_one_pair_is_out_of_range() {
511 let mut db = Keyspace::new();
512 let bad: [(f64, f64, &[u8]); 2] = [(13.0, 38.0, b"good"), (13.0, 86.0, b"bad")];
513 let err = db
514 .geoadd(b"g", bad.iter().copied(), ZAdd::default())
515 .expect_err("86 is past the projection");
516 assert_eq!(
517 err.message(),
518 "invalid longitude,latitude pair 13.000000,86.000000"
519 );
520 assert_eq!(db.zcard(b"g").expect("a zset"), 0);
521 }
522
523 #[test]
524 fn a_position_comes_back_where_it_went_in_give_or_take_two_metres() {
525 let mut db = ks();
526 let mut got = Vec::new();
527 db.geopos(b"g", [&b"Palermo"[..], b"nope"].into_iter(), |p| {
528 got.push(p)
529 })
530 .expect("a zset");
531 let (lon, lat) = got[0].expect("Palermo is there");
532 assert_eq!(format!("{lon}"), "13.361389338970184");
533 assert_eq!(format!("{lat}"), "38.1155563954963");
534 assert_eq!(got[1], None);
535 }
536
537 #[test]
538 fn the_distance_between_two_members_is_the_one_a_real_server_answers() {
539 let mut db = ks();
540 let d = db
541 .geodist(b"g", b"Palermo", b"Catania")
542 .expect("a zset")
543 .expect("both are there");
544 assert_eq!(format!("{d:.4}"), "166274.1516");
545 assert_eq!(format!("{:.4}", d / Unit::Km.metres()), "166.2742");
546 assert_eq!(db.geodist(b"g", b"Palermo", b"nope").expect("a zset"), None);
548 assert_eq!(db.geodist(b"nope", b"a", b"b").expect("no key"), None);
549 }
550
551 #[test]
552 fn a_hash_string_is_eleven_characters_and_ends_in_a_zero() {
553 let mut db = ks();
554 let mut got: Vec<Option<Vec<u8>>> = Vec::new();
555 db.geohash(
556 b"g",
557 [&b"Palermo"[..], b"Catania", b"nope"].into_iter(),
558 |h| {
559 got.push(h.map(<[u8]>::to_vec));
560 },
561 )
562 .expect("a zset");
563 assert_eq!(got[0].as_deref(), Some(&b"sqc8b49rny0"[..]));
564 assert_eq!(got[1].as_deref(), Some(&b"sqdtr74hyu0"[..]));
565 assert_eq!(got[2], None);
566 }
567
568 #[test]
569 fn a_radius_search_finds_what_is_inside_it_nearest_first() {
570 let mut db = ks();
571 let shape = circle(15.0, 37.0, 200.0, Unit::Km);
572 let limit = Limit {
573 sort: Some(Sort::Near),
574 ..Limit::default()
575 };
576 assert_eq!(db.geosearch(b"g", &shape, limit).expect("a zset"), 2);
577 assert_eq!(names(&db), [b"Catania".to_vec(), b"Palermo".to_vec()]);
578 let hits: Vec<f64> = db.geohits().iter().map(|(_, h)| h.metres).collect();
580 assert_eq!(format!("{:.4}", hits[0] / 1000.0), "56.4413");
581 assert_eq!(format!("{:.4}", hits[1] / 1000.0), "190.4424");
582 }
583
584 #[test]
585 fn a_box_search_reaches_the_corners_a_circle_does_not() {
586 let mut db = ks();
587 let shape = Shape {
588 lon: 13.361389,
589 lat: 38.115556,
590 kind: Kind::Rect {
591 width: 400.0,
592 height: 400.0,
593 },
594 unit: Unit::Km,
595 };
596 let limit = Limit {
597 sort: Some(Sort::Near),
598 ..Limit::default()
599 };
600 assert_eq!(db.geosearch(b"g", &shape, limit).expect("a zset"), 3);
601 assert_eq!(
602 names(&db),
603 [b"Palermo".to_vec(), b"edge".to_vec(), b"Catania".to_vec()]
604 );
605 }
606
607 #[test]
608 fn a_count_takes_the_nearest_and_desc_takes_the_furthest() {
609 let mut db = ks();
610 let shape = circle(15.0, 37.0, 200.0, Unit::Km);
611 let limit = Limit {
613 count: Some(1),
614 ..Limit::default()
615 };
616 assert_eq!(db.geosearch(b"g", &shape, limit).expect("a zset"), 1);
617 assert_eq!(names(&db), [b"Catania".to_vec()]);
618 let limit = Limit {
619 sort: Some(Sort::Far),
620 count: Some(1),
621 ..Limit::default()
622 };
623 assert_eq!(db.geosearch(b"g", &shape, limit).expect("a zset"), 1);
624 assert_eq!(names(&db), [b"Palermo".to_vec()]);
625 }
626
627 #[test]
628 fn any_stops_at_the_count_rather_than_finding_the_nearest() {
629 let mut db = ks();
630 let shape = circle(15.0, 37.0, 200.0, Unit::Km);
631 let limit = Limit {
632 count: Some(1),
633 any: true,
634 ..Limit::default()
635 };
636 assert_eq!(db.geosearch(b"g", &shape, limit).expect("a zset"), 1);
637 assert_eq!(db.geohits().len(), 1);
640 }
641
642 #[test]
643 fn a_search_that_finds_nothing_is_not_an_error() {
644 let mut db = ks();
645 let shape = circle(0.0, 0.0, 1.0, Unit::M);
646 assert_eq!(
647 db.geosearch(b"g", &shape, Limit::default())
648 .expect("a zset"),
649 0
650 );
651 assert!(db.geohits().is_empty());
652 assert_eq!(
654 db.geosearch(b"nope", &shape, Limit::default())
655 .expect("no key"),
656 0
657 );
658 }
659
660 #[test]
661 fn a_search_around_a_member_is_a_search_around_where_it_is() {
662 let mut db = ks();
663 let (lon, lat) = db
664 .geocentre(b"g", b"Palermo")
665 .expect("a zset")
666 .expect("Palermo is there");
667 let shape = circle(lon, lat, 200.0, Unit::Km);
668 let limit = Limit {
669 sort: Some(Sort::Near),
670 ..Limit::default()
671 };
672 assert_eq!(db.geosearch(b"g", &shape, limit).expect("a zset"), 3);
673 assert_eq!(names(&db)[0], b"Palermo".to_vec());
674 assert!(db.geocentre(b"g", b"nope").is_err());
677 assert_eq!(db.geocentre(b"nope", b"nope").expect("no key"), None);
678 }
679
680 #[test]
681 fn a_store_keeps_the_hashes_and_a_storedist_keeps_the_distances() {
682 let mut db = ks();
683 let shape = circle(15.0, 37.0, 200.0, Unit::Km);
684 let limit = Limit {
685 sort: Some(Sort::Near),
686 ..Limit::default()
687 };
688 assert_eq!(
689 db.geosearchstore(b"d", b"g", &shape, limit, false)
690 .expect("a zset"),
691 2
692 );
693 assert_eq!(
694 db.zscore(b"d", b"Catania").expect("a zset"),
695 Some(3_479_447_370_796_909.0)
696 );
697 assert_eq!(
699 db.geodist(b"d", b"Palermo", b"Catania")
700 .expect("a zset")
701 .map(|d| format!("{d:.4}")),
702 Some("166274.1516".to_string())
703 );
704 assert_eq!(
705 db.geosearchstore(b"e", b"g", &shape, limit, true)
706 .expect("a zset"),
707 2
708 );
709 let d = db
710 .zscore(b"e", b"Catania")
711 .expect("a zset")
712 .expect("stored");
713 assert_eq!(format!("{d:.4}"), "56.4413");
714 }
715
716 #[test]
717 fn a_store_that_finds_nothing_deletes_what_was_there() {
718 let mut db = ks();
719 let shape = circle(15.0, 37.0, 200.0, Unit::Km);
720 assert_eq!(
721 db.geosearchstore(b"d", b"g", &shape, Limit::default(), false)
722 .expect("a zset"),
723 2
724 );
725 let empty = circle(0.0, 0.0, 1.0, Unit::M);
726 assert_eq!(
727 db.geosearchstore(b"d", b"g", &empty, Limit::default(), false)
728 .expect("a zset"),
729 0
730 );
731 assert_eq!(db.zcard(b"d").expect("no key"), 0);
732 }
733
734 #[test]
735 fn a_search_across_the_date_line_finds_both_sides_of_it() {
736 let mut db = Keyspace::new();
737 let pair: [(f64, f64, &[u8]); 2] = [(179.9, 0.0, b"west"), (-179.9, 0.0, b"east")];
738 db.geoadd(b"d", pair.iter().copied(), ZAdd::default())
739 .expect("both are in range");
740 let d = db
743 .geodist(b"d", b"west", b"east")
744 .expect("a zset")
745 .expect("both are there");
746 assert_eq!(format!("{:.4}", d / Unit::Km.metres()), "22.2454");
747 let shape = circle(179.95, 0.0, 50.0, Unit::Km);
748 assert_eq!(
749 db.geosearch(b"d", &shape, Limit::default())
750 .expect("a zset"),
751 2
752 );
753 let shape = circle(180.0, 0.0, 50.0, Unit::Km);
758 assert_eq!(
759 db.geosearch(b"d", &shape, Limit::default())
760 .expect("a zset"),
761 1
762 );
763 assert_eq!(names(&db), [b"west".to_vec()]);
764 }
765
766 #[test]
767 fn a_key_holding_something_else_is_refused_everywhere() {
768 let mut db = Keyspace::new();
769 db.set(b"s", b"v", strings::SetOptions::default())
770 .expect("a fresh key");
771 let shape = circle(0.0, 0.0, 1.0, Unit::Km);
772 assert_eq!(
773 db.geoadd(b"s", PLACES.iter().copied(), ZAdd::default())
774 .expect_err("a string")
775 .code(),
776 Code::WrongType
777 );
778 assert!(db.geopos(b"s", [&b"x"[..]].into_iter(), |_| {}).is_err());
779 assert!(db.geohash(b"s", [&b"x"[..]].into_iter(), |_| {}).is_err());
780 assert!(db.geodist(b"s", b"a", b"b").is_err());
781 assert!(db.geosearch(b"s", &shape, Limit::default()).is_err());
782 assert!(
783 db.geosearchstore(b"d", b"s", &shape, Limit::default(), false)
784 .is_err()
785 );
786 }
787}