Skip to main content

yo_kv/
geos.rs

1//! The geospatial commands.
2//!
3//! Ten of them, and every one is a sorted set command underneath, because that
4//! is what a geo key is. `GEOADD` is `ZADD` with the coordinates turned into a
5//! score, `GEOPOS` and `GEOHASH` and `GEODIST` are `ZSCORE` with arithmetic on
6//! the answer, and the six search forms are nine `ZRANGEBYSCORE` calls with a
7//! distance filter over the results. Nothing here holds any state a sorted set
8//! does not already hold, which is why `TYPE` on a geo key says `zset` and why a
9//! client can `ZREM` a place out of one.
10//!
11//! # The six search commands are one command
12//!
13//! `GEOSEARCH`, `GEOSEARCHSTORE`, `GEORADIUS`, `GEORADIUS_RO`,
14//! `GEORADIUSBYMEMBER` and `GEORADIUSBYMEMBER_RO` differ in how they spell where
15//! the centre is and whether they are allowed to write. Once the centre, the
16//! shape and the options are parsed there is one search, and it lives in
17//! [`Keyspace::geosearch`]. The store forms run the same search and put the
18//! results in a key instead of handing them back.
19//!
20//! # Why the results are held rather than streamed
21//!
22//! A search cannot answer in order as it goes. The nine boxes are walked in hash
23//! order and the reply is in distance order, so every candidate has to be in
24//! hand before the first one can be written. The wire needs the count before the
25//! members for the same reason every other range command does.
26//!
27//! So there is a [`Scratch`] on the keyspace holding the hits and one byte
28//! buffer with every member's name in it, cleared and refilled per search rather
29//! than allocated per search. A search that found a million points holds a
30//! million until the next one, which is the same trade [`crate::setops`] makes
31//! and for the same reason: the buffer had to exist for the length of the
32//! command anyway.
33//!
34//! # Errors
35//!
36//! `WRONGTYPE` for a key holding something that is not a sorted set, and a
37//! missing key is an empty one everywhere except `GEODIST`, which answers nil
38//! for a key that is not there rather than for a member that is not there, and
39//! `GEORADIUSBYMEMBER`, which needs a member to take its centre from and says so
40//! when it cannot find one.
41
42use 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
53/// What Redis says when a search is asked to start from a member it cannot read
54/// a position out of.
55const NO_MEMBER: &str = "could not decode requested zset member";
56
57/// The error for a point the projection does not reach.
58///
59/// The coordinates are printed back with six decimal places, which is `%f` and
60/// is what Redis formats them with, so `GEOADD k 181 38 x` complains about
61/// `181.000000,38.000000` rather than about `181,38`.
62#[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/// The error for a member a search cannot take its centre from.
73#[must_use]
74pub fn no_member() -> Error {
75    Error::new(Code::Invalid, NO_MEMBER)
76}
77
78/// Which way a search orders what it found.
79#[derive(Debug, Clone, Copy, PartialEq, Eq)]
80pub enum Sort {
81    /// Nearest first. `ASC`.
82    Near,
83    /// Furthest first. `DESC`.
84    Far,
85}
86
87/// What a search was asked for beyond its shape.
88#[derive(Debug, Clone, Copy, Default)]
89pub struct Limit {
90    /// `ASC` or `DESC`, or nothing, which leaves the results in the order the
91    /// boxes were walked in.
92    pub sort: Option<Sort>,
93    /// `COUNT`, or nothing for all of them.
94    pub count: Option<usize>,
95    /// `ANY`, which stops the walk as soon as `count` are in hand rather than
96    /// finding them all and keeping the nearest.
97    pub any: bool,
98}
99
100impl Limit {
101    /// The ordering the search actually runs with.
102    ///
103    /// A `COUNT` with no `ASC` or `DESC` means the nearest ones, so it implies
104    /// `ASC`. `ANY` is the exception: it says the caller does not care which
105    /// ones, only how many, and sorting would undo the whole point of it.
106    #[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    /// How many the walk may stop at, or nothing if it has to find them all.
116    #[must_use]
117    fn cap(&self) -> Option<usize> {
118        self.any.then_some(self.count).flatten()
119    }
120}
121
122/// One member a search found.
123#[derive(Debug, Clone, Copy)]
124pub struct Hit {
125    /// Where the member's name starts in the scratch buffer.
126    at: usize,
127    /// How long the name is.
128    len: usize,
129    /// The raw score, which is what `WITHHASH` answers.
130    pub score: u64,
131    /// The stored longitude.
132    pub lon: f64,
133    /// The stored latitude.
134    pub lat: f64,
135    /// How far from the centre of the search, in metres.
136    pub metres: f64,
137}
138
139/// The buffers a search fills, kept rather than built per call.
140///
141/// Two of them: the hits, and one run of bytes with every member's name in it.
142/// A name is a slice of the second, which is why [`Hit`] carries an offset and a
143/// length rather than a `Vec<u8>` each. A search over ten thousand points is one
144/// buffer that grows once instead of ten thousand small allocations.
145#[derive(Debug, Default)]
146pub struct Scratch {
147    /// What the search found, in the order the boxes were walked unless it was
148    /// sorted afterwards.
149    hits: Vec<Hit>,
150    /// The names, end to end.
151    names: Vec<u8>,
152}
153
154impl Scratch {
155    /// Everything the last search found, longest to say and cheapest to read.
156    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    /// How many hits the last search left here.
163    #[must_use]
164    pub fn len(&self) -> usize {
165        self.hits.len()
166    }
167
168    /// Whether the last search found nothing.
169    #[must_use]
170    pub fn is_empty(&self) -> bool {
171        self.hits.is_empty()
172    }
173
174    /// Forget everything, keeping the space.
175    fn clear(&mut self) {
176        self.hits.clear();
177        self.names.clear();
178    }
179
180    /// Remember one point.
181    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    /// Put the hits in distance order and cut them down to `count`.
195    ///
196    /// The cut comes first when there is one, through a selection rather than a
197    /// full sort, because a search over a city with `COUNT 10` should not pay to
198    /// order the other side of the city. Ties between equal distances land in
199    /// whatever order the selection leaves them, which is what a real server's
200    /// `qsort` does too.
201    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    /// `GEOADD key [NX|XX] [CH] longitude latitude member [...]`.
225    ///
226    /// Answers what the `ZADD` underneath answers, which is how many members
227    /// were added, or how many were added or moved with `CH`.
228    ///
229    /// Every coordinate is checked before anything is stored, so a call with one
230    /// bad pair in the middle of it leaves the key exactly as it was. Redis does
231    /// the same, and it matters more here than it looks: `GEOADD` is how a whole
232    /// dataset gets loaded, and a partial load with no way to tell where it
233    /// stopped is worse than a refusal.
234    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    /// The same write, with both halves of the count kept apart.
243    ///
244    /// For the reason [`Keyspace::zadd_counts`] exists: the reply wants one
245    /// number and the notification wants to know whether either half was more
246    /// than nothing, since a `GEOADD` that moved a member without adding one has
247    /// written to the key and says so.
248    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                // Checked in the pass above, and nothing between the two passes
267                // can change what a pair of coordinates hashes to.
268                (geo::score(lon, lat).expect("checked") as f64, m)
269            }),
270            opts,
271        )
272    }
273
274    /// `GEOPOS key member [member ...]`.
275    ///
276    /// Hands each position over as it is found rather than collecting them,
277    /// because the reply is as long as the argument list and the wire already
278    /// knows that number. A member that is not there, or whose score is not a
279    /// position anything wrote, gets nothing.
280    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    /// `GEOHASH key member [member ...]`, the same shape one step further on.
301    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    /// `GEODIST key member1 member2 [unit]`, in metres whatever the unit was.
329    ///
330    /// The caller divides, because the unit is a wire concern and the number
331    /// this hands back is the one every other distance in the crate is in.
332    ///
333    /// Nothing at all if either member is missing, and nothing for a key that is
334    /// not there, which are the same nil on the wire.
335    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    /// Where a member is, for a search that takes its centre from one.
350    ///
351    /// `FROMMEMBER`, and the two `GEORADIUSBYMEMBER` forms. A key that is not
352    /// there answers nothing, because those commands have their own reply for
353    /// that and it is not the error a missing member gets.
354    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    /// Run a search and leave what it found on the keyspace.
365    ///
366    /// Answers how many hits there are, which is what the wire needs before it
367    /// can write the array header. The hits themselves come from
368    /// [`Keyspace::geohits`], which borrows rather than copies.
369    ///
370    /// The nine boxes are walked in Redis's order and a box that a previous one
371    /// already covered is skipped, which is not an optimisation: at a radius of
372    /// a few thousand kilometres the step is small enough that neighbouring
373    /// boxes come out identical, and walking one twice would report every member
374    /// in it twice.
375    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    /// What the last [`Keyspace::geosearch`] found.
393    #[must_use]
394    pub fn geohits(&self) -> &Scratch {
395        &self.geo
396    }
397
398    /// `GEOSEARCHSTORE`, and the `STORE` and `STOREDIST` forms of `GEORADIUS`.
399    ///
400    /// Answers how many members went into the destination. A search that found
401    /// nothing deletes the destination rather than leaving an empty sorted set
402    /// or leaving the old contents, which is the rule every store form follows.
403    ///
404    /// `dist` is `STOREDIST`, which stores the distance in the shape's unit as
405    /// the score instead of the geohash. The two are not interchangeable: a key
406    /// written with `STOREDIST` is a sorted set of distances and is not a geo
407    /// key any more, and `GEOPOS` on it answers positions somewhere off the
408    /// coast of Africa rather than an error.
409    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    /// `GEOSEARCHSTORE` and the two `GEORADIUS` store forms, when the source and
437    /// the destination are not on the same stripe.
438    ///
439    /// The search runs on the source's stripe and leaves its hits in that
440    /// stripe's scratch, which is where they are read from while the result is
441    /// built. The sorted set that comes out is put on the destination's stripe
442    /// under that stripe's promotion thresholds, the same way every other store
443    /// form works.
444    ///
445    /// `from` is the member the centre comes from, for the forms that take it
446    /// from one rather than from a pair of coordinates. It is read here, with
447    /// the stripes already held, so that the centre and the members it is
448    /// measured against are the same key at the same moment. A caller that read
449    /// it beforehand would be handing in a point the member may have moved away
450    /// from since.
451    ///
452    /// # Errors
453    ///
454    /// `WRONGTYPE` from the source, which is checked before the destination is
455    /// touched, and the missing member complaint when `from` names a member the
456    /// source does not have.
457    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        // Both at once and in stripe order, since the search leaves its hits in
473        // the source stripe's scratch, the destination's limits decide what is
474        // built from them, and the destination is written from that. The search
475        // runs with both already held, because a scratch that was filled and
476        // then let go of is a scratch the next search on that stripe overwrites.
477        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
495/// The shape with its centre read out of the source, for the search forms whose
496/// centre is a member rather than a pair of coordinates.
497///
498/// Called with the stripe held, which is the whole point of it. A source that is
499/// not there leaves the shape as it came in, since a search on a key that is not
500/// there finds nothing wherever it is pointed.
501fn 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
516/// Walk the nine boxes and keep every member the shape covers.
517///
518/// Split out of [`Keyspace::geosearch`] so that the sorted set borrow and the
519/// scratch buffer are two arguments rather than two borrows of the same
520/// keyspace, which they cannot both be.
521fn 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    // Redis compares each box against the last one it actually walked, and it
526    // starts that at index zero, which means the centre box is never the thing a
527    // neighbour is compared against. Keeping the quirk keeps the duplicate
528    // behaviour identical at radii large enough for the boxes to collide.
529    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/// A circle around a point, which is what five of the six search forms want.
560#[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    /// The three places every Redis geo example uses, and one more far enough
576    /// away to be outside every search here.
577    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        // The scores a real 8.10.1 has after the same `GEOADD`.
601        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        // One member missing is the same nil as the whole key missing.
649        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        // The distances a real server prints for the same search.
681        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        // A `COUNT` on its own means the nearest, with no `ASC` written.
714        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        // Whichever box came first, which is not necessarily the nearest, and
740        // that is the whole point of the option.
741        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        // Nor is a key that is not there.
755        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        // A member that is not there is the error and not a nil, which is what
777        // separates it from a key that is not there.
778        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        // A stored search is still a geo key, so the round trip works.
800        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        // A destination somewhere other than where the source is, so that the
843        // two stripe path is the one under test.
844        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        // The shape points at nowhere near Sicily, so a result with the two
849        // Sicilian places in it is a centre that came from the member.
850        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        // A member that is not there is the error, whatever the shape it came
864        // in with says.
865        assert!(
866            db.geosearchstore(&dest, b"g", Some(b"nope"), &shape, limit, false)
867                .is_err()
868        );
869        // And a source that is not there is not, because a search on a key that
870        // is not there finds nothing wherever it is pointed.
871        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        // The two are 22.2454 kilometres apart going the short way, and a search
885        // that treated longitude as a plain number would find one of them.
886        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        // Centred exactly on 180 it finds only the western one, because 180 is
898        // the last longitude the projection has and there is no box east of it
899        // to be a neighbour. A real server answers the same one member, so this
900        // is pinned rather than fixed.
901        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}