1use yo_common::num::DIGITS_MAX;
43use yo_common::{Code, Error, Result};
44
45use crate::db::Db;
46use crate::elem::Elements;
47use crate::geo::{self, Kind, Shape, Unit};
48use crate::keyspace::Keyspace;
49use crate::strings;
50use crate::zset::{Bound, Zset};
51use crate::zsets::{ZAdd, member_bytes};
52
53const NO_MEMBER: &str = "could not decode requested zset member";
56
57#[must_use]
63pub fn out_of_range(lon: f64, lat: f64) -> Error {
64 yo_alloc::allow(|| {
65 Error::fmt(
66 Code::Invalid,
67 format_args!("invalid longitude,latitude pair {lon:.6},{lat:.6}"),
68 )
69 })
70}
71
72#[must_use]
74pub fn no_member() -> Error {
75 Error::new(Code::Invalid, NO_MEMBER)
76}
77
78#[derive(Debug, Clone, Copy, PartialEq, Eq)]
80pub enum Sort {
81 Near,
83 Far,
85}
86
87#[derive(Debug, Clone, Copy, Default)]
89pub struct Limit {
90 pub sort: Option<Sort>,
93 pub count: Option<usize>,
95 pub any: bool,
98}
99
100impl Limit {
101 #[must_use]
107 fn ordering(&self) -> Option<Sort> {
108 match self.sort {
109 Some(s) => Some(s),
110 None if self.count.is_some() && !self.any => Some(Sort::Near),
111 None => None,
112 }
113 }
114
115 #[must_use]
117 fn cap(&self) -> Option<usize> {
118 self.any.then_some(self.count).flatten()
119 }
120}
121
122#[derive(Debug, Clone, Copy)]
124pub struct Hit {
125 at: usize,
127 len: usize,
129 pub score: u64,
131 pub lon: f64,
133 pub lat: f64,
135 pub metres: f64,
137}
138
139#[derive(Debug, Default)]
146pub struct Scratch {
147 hits: Vec<Hit>,
150 names: Vec<u8>,
152}
153
154impl Scratch {
155 pub fn iter(&self) -> impl Iterator<Item = (&[u8], &Hit)> {
157 self.hits
158 .iter()
159 .map(|h| (&self.names[h.at..h.at + h.len], h))
160 }
161
162 #[must_use]
164 pub fn len(&self) -> usize {
165 self.hits.len()
166 }
167
168 #[must_use]
170 pub fn is_empty(&self) -> bool {
171 self.hits.is_empty()
172 }
173
174 fn clear(&mut self) {
176 self.hits.clear();
177 self.names.clear();
178 }
179
180 fn push(&mut self, name: &[u8], score: u64, lon: f64, lat: f64, metres: f64) {
182 let at = self.names.len();
183 self.names.extend_from_slice(name);
184 self.hits.push(Hit {
185 at,
186 len: name.len(),
187 score,
188 lon,
189 lat,
190 metres,
191 });
192 }
193
194 fn order(&mut self, limit: Limit) {
202 let Some(sort) = limit.ordering() else {
203 self.hits.truncate(limit.count.unwrap_or(usize::MAX));
204 return;
205 };
206 let near = |a: &Hit, b: &Hit| a.metres.total_cmp(&b.metres);
207 let far = |a: &Hit, b: &Hit| b.metres.total_cmp(&a.metres);
208 let want = limit.count.unwrap_or(self.hits.len()).min(self.hits.len());
209 if want < self.hits.len() {
210 match sort {
211 Sort::Near => self.hits.select_nth_unstable_by(want, near),
212 Sort::Far => self.hits.select_nth_unstable_by(want, far),
213 };
214 self.hits.truncate(want);
215 }
216 match sort {
217 Sort::Near => self.hits.sort_unstable_by(near),
218 Sort::Far => self.hits.sort_unstable_by(far),
219 }
220 }
221}
222
223impl Keyspace {
224 pub fn geoadd<'m, I>(&mut self, key: &[u8], points: I, opts: ZAdd) -> Result<usize>
235 where
236 I: Iterator<Item = (f64, f64, &'m [u8])> + Clone,
237 {
238 let (added, changed) = self.geoadd_counts(key, points, opts)?;
239 Ok(if opts.changed { added + changed } else { added })
240 }
241
242 pub fn geoadd_counts<'m, I>(
249 &mut self,
250 key: &[u8],
251 points: I,
252 opts: ZAdd,
253 ) -> Result<(usize, usize)>
254 where
255 I: Iterator<Item = (f64, f64, &'m [u8])> + Clone,
256 {
257 for (lon, lat, member) in points.clone() {
258 strings::check_len(key, member.len())?;
259 if geo::score(lon, lat).is_none() {
260 return Err(out_of_range(lon, lat));
261 }
262 }
263 self.zadd_counts(
264 key,
265 points.map(|(lon, lat, m)| {
266 (geo::score(lon, lat).expect("checked") as f64, m)
269 }),
270 opts,
271 )
272 }
273
274 pub fn geopos<'m, F>(
281 &mut self,
282 key: &[u8],
283 members: impl Iterator<Item = &'m [u8]>,
284 mut f: F,
285 ) -> Result<()>
286 where
287 F: FnMut(Option<(f64, f64)>),
288 {
289 let Some(at) = self.zset_slot(key)? else {
290 members.for_each(|_| f(None));
291 return Ok(());
292 };
293 let z = self.zset_at(at);
294 for m in members {
295 f(z.score(m).and_then(geo::decode));
296 }
297 Ok(())
298 }
299
300 pub fn geohash<'m, F>(
302 &mut self,
303 key: &[u8],
304 members: impl Iterator<Item = &'m [u8]>,
305 mut f: F,
306 ) -> Result<()>
307 where
308 F: FnMut(Option<&[u8]>),
309 {
310 let Some(at) = self.zset_slot(key)? else {
311 members.for_each(|_| f(None));
312 return Ok(());
313 };
314 let z = self.zset_at(at);
315 for m in members {
316 let text = z
317 .score(m)
318 .and_then(geo::decode)
319 .and_then(|(lon, lat)| geo::geohash(lon, lat));
320 match text {
321 Some(bytes) => f(Some(&bytes)),
322 None => f(None),
323 }
324 }
325 Ok(())
326 }
327
328 pub fn geodist(&mut self, key: &[u8], a: &[u8], b: &[u8]) -> Result<Option<f64>> {
336 let Some(at) = self.zset_slot(key)? else {
337 return Ok(None);
338 };
339 let z = self.zset_at(at);
340 let (Some(sa), Some(sb)) = (z.score(a), z.score(b)) else {
341 return Ok(None);
342 };
343 let (Some(pa), Some(pb)) = (geo::decode(sa), geo::decode(sb)) else {
344 return Ok(None);
345 };
346 Ok(Some(geo::distance(pa.0, pa.1, pb.0, pb.1)))
347 }
348
349 pub fn geocentre(&mut self, key: &[u8], member: &[u8]) -> Result<Option<(f64, f64)>> {
355 let Some(at) = self.zset_slot(key)? else {
356 return Ok(None);
357 };
358 match self.zset_at(at).score(member).and_then(geo::decode) {
359 Some(xy) => Ok(Some(xy)),
360 None => Err(no_member()),
361 }
362 }
363
364 pub fn geosearch(&mut self, key: &[u8], shape: &Shape, limit: Limit) -> Result<usize> {
376 let mut found = std::mem::take(&mut self.geo);
377 found.clear();
378 let outcome = match self.zset_slot(key) {
379 Err(e) => Err(e),
380 Ok(None) => Ok(()),
381 Ok(Some(at)) => {
382 collect(self.zset_at(at), shape, limit, &mut found);
383 Ok(())
384 }
385 };
386 found.order(limit);
387 let n = found.hits.len();
388 self.geo = found;
389 outcome.map(|()| n)
390 }
391
392 #[must_use]
394 pub fn geohits(&self) -> &Scratch {
395 &self.geo
396 }
397
398 pub fn geosearchstore(
410 &mut self,
411 dest: &[u8],
412 src: &[u8],
413 shape: &Shape,
414 limit: Limit,
415 dist: bool,
416 ) -> Result<usize> {
417 let n = self.geosearch(src, shape, limit)?;
418 let found = std::mem::take(&mut self.geo);
419 let mut got = Elements::with_capacity(n.max(16));
420 for (name, hit) in found.iter() {
421 let score = if dist {
422 hit.metres / shape.unit.metres()
423 } else {
424 hit.score as f64
425 };
426 let _ = got.insert(name, score);
427 }
428 self.geo = found;
429 let limits = self.zset_limits;
430 let built = Zset::from_elements(got, &limits);
431 Ok(self.put_zset(dest, built))
432 }
433}
434
435impl Db {
436 pub fn geosearchstore(
458 &self,
459 dest: &[u8],
460 src: &[u8],
461 from: Option<&[u8]>,
462 shape: &Shape,
463 limit: Limit,
464 dist: bool,
465 ) -> Result<usize> {
466 let (home, onto) = (self.stripe_of(src), self.stripe_of(dest));
467 if home == onto {
468 let mut stripe = self.hold_stripe(home);
469 let shape = recentre(&mut stripe, src, from, shape)?;
470 return stripe.geosearchstore(dest, src, &shape, limit, dist);
471 }
472 let mut held = self.hold_many([home, onto].into_iter());
478 let shape = recentre(held.stripe_mut(home), src, from, shape)?;
479 let n = held.stripe_mut(home).geosearch(src, &shape, limit)?;
480 let mut got = Elements::with_capacity(n.max(16));
481 for (name, hit) in held.stripe(home).geohits().iter() {
482 let score = if dist {
483 hit.metres / shape.unit.metres()
484 } else {
485 hit.score as f64
486 };
487 let _ = got.insert(name, score);
488 }
489 let limits = held.stripe(onto).zset_limits;
490 let built = Zset::from_elements(got, &limits);
491 Ok(held.stripe_mut(onto).put_zset(dest, built))
492 }
493}
494
495fn recentre(
502 stripe: &mut Keyspace,
503 src: &[u8],
504 from: Option<&[u8]>,
505 shape: &Shape,
506) -> Result<Shape> {
507 let mut shape = *shape;
508 if let Some(member) = from
509 && let Some(centre) = stripe.geocentre(src, member)?
510 {
511 (shape.lon, shape.lat) = centre;
512 }
513 Ok(shape)
514}
515
516fn collect(z: &Zset, shape: &Shape, limit: Limit, out: &mut Scratch) {
522 let search = geo::areas(shape);
523 let cap = limit.cap();
524 let mut digits = [0u8; DIGITS_MAX];
525 let mut last = 0usize;
530 for i in 0..search.boxes.len() {
531 let hash = search.boxes[i];
532 if hash.bits == 0 && hash.step == 0 {
533 continue;
534 }
535 if last != 0 && hash == search.boxes[last] {
536 continue;
537 }
538 if cap.is_some_and(|n| out.hits.len() >= n) {
539 break;
540 }
541 let (low, high) = geo::range(hash);
542 let window = z.window_by_score(Bound::closed(low as f64), Bound::open(high as f64));
543 z.walk(window.start, window.len(), false, |m, raw| {
544 if cap.is_some_and(|n| out.hits.len() >= n) {
545 return;
546 }
547 let Some((lon, lat)) = geo::decode(raw) else {
548 return;
549 };
550 let Some(metres) = shape.covers(lon, lat) else {
551 return;
552 };
553 out.push(member_bytes(m, &mut digits), raw as u64, lon, lat, metres);
554 });
555 last = i;
556 }
557}
558
559#[must_use]
561pub fn circle(lon: f64, lat: f64, radius: f64, unit: Unit) -> Shape {
562 Shape {
563 lon,
564 lat,
565 kind: Kind::Circle { radius },
566 unit,
567 }
568}
569
570#[cfg(test)]
571mod tests {
572 use super::*;
573 use crate::Clock;
574
575 const PLACES: [(f64, f64, &[u8]); 4] = [
578 (13.361389, 38.115556, b"Palermo"),
579 (15.087269, 37.502669, b"Catania"),
580 (12.758489, 38.788135, b"edge"),
581 (2.352222, 48.856613, b"Paris"),
582 ];
583
584 fn ks() -> Keyspace {
585 let mut db = Keyspace::new();
586 let opts = ZAdd::default();
587 db.geoadd(b"g", PLACES.iter().copied(), opts)
588 .expect("the places are all in range");
589 db
590 }
591
592 fn names(db: &Keyspace) -> Vec<Vec<u8>> {
593 db.geohits().iter().map(|(n, _)| n.to_vec()).collect()
594 }
595
596 #[test]
597 fn a_geo_key_is_a_sorted_set_of_hashes() {
598 let mut db = ks();
599 assert_eq!(db.zcard(b"g").expect("a zset"), 4);
600 assert_eq!(
602 db.zscore(b"g", b"Palermo").expect("a zset"),
603 Some(3_479_099_956_230_698.0)
604 );
605 assert_eq!(
606 db.zscore(b"g", b"Catania").expect("a zset"),
607 Some(3_479_447_370_796_909.0)
608 );
609 }
610
611 #[test]
612 fn nothing_is_stored_when_one_pair_is_out_of_range() {
613 let mut db = Keyspace::new();
614 let bad: [(f64, f64, &[u8]); 2] = [(13.0, 38.0, b"good"), (13.0, 86.0, b"bad")];
615 let err = db
616 .geoadd(b"g", bad.iter().copied(), ZAdd::default())
617 .expect_err("86 is past the projection");
618 assert_eq!(
619 err.message(),
620 "invalid longitude,latitude pair 13.000000,86.000000"
621 );
622 assert_eq!(db.zcard(b"g").expect("a zset"), 0);
623 }
624
625 #[test]
626 fn a_position_comes_back_where_it_went_in_give_or_take_two_metres() {
627 let mut db = ks();
628 let mut got = Vec::new();
629 db.geopos(b"g", [&b"Palermo"[..], b"nope"].into_iter(), |p| {
630 got.push(p)
631 })
632 .expect("a zset");
633 let (lon, lat) = got[0].expect("Palermo is there");
634 assert_eq!(format!("{lon}"), "13.361389338970184");
635 assert_eq!(format!("{lat}"), "38.1155563954963");
636 assert_eq!(got[1], None);
637 }
638
639 #[test]
640 fn the_distance_between_two_members_is_the_one_a_real_server_answers() {
641 let mut db = ks();
642 let d = db
643 .geodist(b"g", b"Palermo", b"Catania")
644 .expect("a zset")
645 .expect("both are there");
646 assert_eq!(format!("{d:.4}"), "166274.1516");
647 assert_eq!(format!("{:.4}", d / Unit::Km.metres()), "166.2742");
648 assert_eq!(db.geodist(b"g", b"Palermo", b"nope").expect("a zset"), None);
650 assert_eq!(db.geodist(b"nope", b"a", b"b").expect("no key"), None);
651 }
652
653 #[test]
654 fn a_hash_string_is_eleven_characters_and_ends_in_a_zero() {
655 let mut db = ks();
656 let mut got: Vec<Option<Vec<u8>>> = Vec::new();
657 db.geohash(
658 b"g",
659 [&b"Palermo"[..], b"Catania", b"nope"].into_iter(),
660 |h| {
661 got.push(h.map(<[u8]>::to_vec));
662 },
663 )
664 .expect("a zset");
665 assert_eq!(got[0].as_deref(), Some(&b"sqc8b49rny0"[..]));
666 assert_eq!(got[1].as_deref(), Some(&b"sqdtr74hyu0"[..]));
667 assert_eq!(got[2], None);
668 }
669
670 #[test]
671 fn a_radius_search_finds_what_is_inside_it_nearest_first() {
672 let mut db = ks();
673 let shape = circle(15.0, 37.0, 200.0, Unit::Km);
674 let limit = Limit {
675 sort: Some(Sort::Near),
676 ..Limit::default()
677 };
678 assert_eq!(db.geosearch(b"g", &shape, limit).expect("a zset"), 2);
679 assert_eq!(names(&db), [b"Catania".to_vec(), b"Palermo".to_vec()]);
680 let hits: Vec<f64> = db.geohits().iter().map(|(_, h)| h.metres).collect();
682 assert_eq!(format!("{:.4}", hits[0] / 1000.0), "56.4413");
683 assert_eq!(format!("{:.4}", hits[1] / 1000.0), "190.4424");
684 }
685
686 #[test]
687 fn a_box_search_reaches_the_corners_a_circle_does_not() {
688 let mut db = ks();
689 let shape = Shape {
690 lon: 13.361389,
691 lat: 38.115556,
692 kind: Kind::Rect {
693 width: 400.0,
694 height: 400.0,
695 },
696 unit: Unit::Km,
697 };
698 let limit = Limit {
699 sort: Some(Sort::Near),
700 ..Limit::default()
701 };
702 assert_eq!(db.geosearch(b"g", &shape, limit).expect("a zset"), 3);
703 assert_eq!(
704 names(&db),
705 [b"Palermo".to_vec(), b"edge".to_vec(), b"Catania".to_vec()]
706 );
707 }
708
709 #[test]
710 fn a_count_takes_the_nearest_and_desc_takes_the_furthest() {
711 let mut db = ks();
712 let shape = circle(15.0, 37.0, 200.0, Unit::Km);
713 let limit = Limit {
715 count: Some(1),
716 ..Limit::default()
717 };
718 assert_eq!(db.geosearch(b"g", &shape, limit).expect("a zset"), 1);
719 assert_eq!(names(&db), [b"Catania".to_vec()]);
720 let limit = Limit {
721 sort: Some(Sort::Far),
722 count: Some(1),
723 ..Limit::default()
724 };
725 assert_eq!(db.geosearch(b"g", &shape, limit).expect("a zset"), 1);
726 assert_eq!(names(&db), [b"Palermo".to_vec()]);
727 }
728
729 #[test]
730 fn any_stops_at_the_count_rather_than_finding_the_nearest() {
731 let mut db = ks();
732 let shape = circle(15.0, 37.0, 200.0, Unit::Km);
733 let limit = Limit {
734 count: Some(1),
735 any: true,
736 ..Limit::default()
737 };
738 assert_eq!(db.geosearch(b"g", &shape, limit).expect("a zset"), 1);
739 assert_eq!(db.geohits().len(), 1);
742 }
743
744 #[test]
745 fn a_search_that_finds_nothing_is_not_an_error() {
746 let mut db = ks();
747 let shape = circle(0.0, 0.0, 1.0, Unit::M);
748 assert_eq!(
749 db.geosearch(b"g", &shape, Limit::default())
750 .expect("a zset"),
751 0
752 );
753 assert!(db.geohits().is_empty());
754 assert_eq!(
756 db.geosearch(b"nope", &shape, Limit::default())
757 .expect("no key"),
758 0
759 );
760 }
761
762 #[test]
763 fn a_search_around_a_member_is_a_search_around_where_it_is() {
764 let mut db = ks();
765 let (lon, lat) = db
766 .geocentre(b"g", b"Palermo")
767 .expect("a zset")
768 .expect("Palermo is there");
769 let shape = circle(lon, lat, 200.0, Unit::Km);
770 let limit = Limit {
771 sort: Some(Sort::Near),
772 ..Limit::default()
773 };
774 assert_eq!(db.geosearch(b"g", &shape, limit).expect("a zset"), 3);
775 assert_eq!(names(&db)[0], b"Palermo".to_vec());
776 assert!(db.geocentre(b"g", b"nope").is_err());
779 assert_eq!(db.geocentre(b"nope", b"nope").expect("no key"), None);
780 }
781
782 #[test]
783 fn a_store_keeps_the_hashes_and_a_storedist_keeps_the_distances() {
784 let mut db = ks();
785 let shape = circle(15.0, 37.0, 200.0, Unit::Km);
786 let limit = Limit {
787 sort: Some(Sort::Near),
788 ..Limit::default()
789 };
790 assert_eq!(
791 db.geosearchstore(b"d", b"g", &shape, limit, false)
792 .expect("a zset"),
793 2
794 );
795 assert_eq!(
796 db.zscore(b"d", b"Catania").expect("a zset"),
797 Some(3_479_447_370_796_909.0)
798 );
799 assert_eq!(
801 db.geodist(b"d", b"Palermo", b"Catania")
802 .expect("a zset")
803 .map(|d| format!("{d:.4}")),
804 Some("166274.1516".to_string())
805 );
806 assert_eq!(
807 db.geosearchstore(b"e", b"g", &shape, limit, true)
808 .expect("a zset"),
809 2
810 );
811 let d = db
812 .zscore(b"e", b"Catania")
813 .expect("a zset")
814 .expect("stored");
815 assert_eq!(format!("{d:.4}"), "56.4413");
816 }
817
818 #[test]
819 fn a_store_that_finds_nothing_deletes_what_was_there() {
820 let mut db = ks();
821 let shape = circle(15.0, 37.0, 200.0, Unit::Km);
822 assert_eq!(
823 db.geosearchstore(b"d", b"g", &shape, Limit::default(), false)
824 .expect("a zset"),
825 2
826 );
827 let empty = circle(0.0, 0.0, 1.0, Unit::M);
828 assert_eq!(
829 db.geosearchstore(b"d", b"g", &empty, Limit::default(), false)
830 .expect("a zset"),
831 0
832 );
833 assert_eq!(db.zcard(b"d").expect("no key"), 0);
834 }
835
836 #[test]
837 fn a_store_from_a_member_takes_its_centre_off_the_held_source() {
838 let db = Db::with_clock(Clock::fixed(1_000_000), 8);
839 db.hold(b"g")
840 .geoadd(b"g", PLACES.iter().copied(), ZAdd::default())
841 .expect("the places are all in range");
842 let dest = (0u32..)
845 .map(|i| format!("d{i}").into_bytes())
846 .find(|d| db.stripe_of(d) != db.stripe_of(b"g"))
847 .expect("some name lands on another stripe");
848 let shape = circle(0.0, 0.0, 200.0, Unit::Km);
851 let limit = Limit::default();
852 assert_eq!(
853 db.geosearchstore(&dest, b"g", Some(b"Catania"), &shape, limit, false)
854 .expect("a zset"),
855 2
856 );
857 assert!(
858 db.hold(&dest)
859 .zscore(&dest, b"Palermo")
860 .expect("a zset")
861 .is_some()
862 );
863 assert!(
866 db.geosearchstore(&dest, b"g", Some(b"nope"), &shape, limit, false)
867 .is_err()
868 );
869 assert_eq!(
872 db.geosearchstore(&dest, b"gone", Some(b"nope"), &shape, limit, false)
873 .expect("no key"),
874 0
875 );
876 }
877
878 #[test]
879 fn a_search_across_the_date_line_finds_both_sides_of_it() {
880 let mut db = Keyspace::new();
881 let pair: [(f64, f64, &[u8]); 2] = [(179.9, 0.0, b"west"), (-179.9, 0.0, b"east")];
882 db.geoadd(b"d", pair.iter().copied(), ZAdd::default())
883 .expect("both are in range");
884 let d = db
887 .geodist(b"d", b"west", b"east")
888 .expect("a zset")
889 .expect("both are there");
890 assert_eq!(format!("{:.4}", d / Unit::Km.metres()), "22.2454");
891 let shape = circle(179.95, 0.0, 50.0, Unit::Km);
892 assert_eq!(
893 db.geosearch(b"d", &shape, Limit::default())
894 .expect("a zset"),
895 2
896 );
897 let shape = circle(180.0, 0.0, 50.0, Unit::Km);
902 assert_eq!(
903 db.geosearch(b"d", &shape, Limit::default())
904 .expect("a zset"),
905 1
906 );
907 assert_eq!(names(&db), [b"west".to_vec()]);
908 }
909
910 #[test]
911 fn a_key_holding_something_else_is_refused_everywhere() {
912 let mut db = Keyspace::new();
913 db.set(b"s", b"v", strings::SetOptions::default())
914 .expect("a fresh key");
915 let shape = circle(0.0, 0.0, 1.0, Unit::Km);
916 assert_eq!(
917 db.geoadd(b"s", PLACES.iter().copied(), ZAdd::default())
918 .expect_err("a string")
919 .code(),
920 Code::WrongType
921 );
922 assert!(db.geopos(b"s", [&b"x"[..]].into_iter(), |_| {}).is_err());
923 assert!(db.geohash(b"s", [&b"x"[..]].into_iter(), |_| {}).is_err());
924 assert!(db.geodist(b"s", b"a", b"b").is_err());
925 assert!(db.geosearch(b"s", &shape, Limit::default()).is_err());
926 assert!(
927 db.geosearchstore(b"d", b"s", &shape, Limit::default(), false)
928 .is_err()
929 );
930 }
931}