Skip to main content

yo_kv/
zsets.rs

1//! The sorted set commands.
2//!
3//! One method per Redis command on [`Keyspace`], the same arrangement the string,
4//! set, hash and list commands use and for the same reason: a key belongs to the
5//! database and not to a type, so `ZADD` against a string has to be able to see
6//! that it is a string. The sorted set itself, and the choice between its two
7//! representations, is [`crate::zset`]. This file is what the wire and the
8//! embedded API both call.
9//!
10//! # Every range command is the same command
11//!
12//! There are nine ways to ask a sorted set for a run of members. `ZRANGE` alone
13//! has three, once `BYSCORE` and `BYLEX` are counted, and then `REV` doubles
14//! them and `ZREVRANGE`, `ZREVRANGEBYSCORE` and `ZREVRANGEBYLEX` exist as older
15//! spellings of the same thing. `ZRANGESTORE` is a tenth, and the three
16//! `ZREMRANGE` forms are three more, and `ZCOUNT` and `ZLEXCOUNT` are two of
17//! those with the walk left off.
18//!
19//! Writing fourteen of those separately is fourteen chances to get an exclusive
20//! bound or a negative index wrong in one of them. So there is one [`Query`],
21//! and every one of those commands is a `Query` turned into a [`Window`], which
22//! is a rank, a count and a direction. What a command does with the window is
23//! all that separates it from the others: walk it, count it, remove it, or walk
24//! it into another key.
25//!
26//! Two calls and not one, because the wire needs the count before it needs the
27//! members: a RESP array writes its length first, and a reply that collected the
28//! members into a `Vec` in order to count them would allocate on the read path,
29//! which is the thing `Y1` is about. So [`Keyspace::zwindow`] answers how many
30//! there are and [`Keyspace::zwalk`] hands them over one at a time. The memo in
31//! the keyspace means the second call does not resolve the key again.
32//!
33//! # The commands over more than one key
34//!
35//! `ZUNION`, `ZINTER`, `ZDIFF` and their three store forms all come through
36//! [`Keyspace::zsetop`] or [`Keyspace::zsetop_store`], because the only thing
37//! that separates them is which members survive and that is [`crate::zsetops`]'s
38//! decision, not this file's. What is decided here is which keys they are
39//! allowed to name. A plain set is a sorted set where every score is one, so
40//! `ZUNIONSTORE d 2 zs plain` is legal and every input resolves through
41//! `live_slot_either`. A key that is not there stays in place as
42//! [`Operand::Missing`] rather than being dropped, because `WEIGHTS` is
43//! positional and closing the gap would hand every later input the wrong
44//! weight.
45//!
46//! `ZINTERCARD` is not `ZINTER` with the members thrown away, because it can
47//! stop at its limit and never has to work out a single score.
48//!
49//! # Errors
50//!
51//! Every command here answers `WRONGTYPE` for a key holding something that is
52//! not a sorted set, and treats a missing key as an empty one, which between
53//! them cover every case because a key is a sorted set, or another type, or
54//! absent. The commands over several keys resolve every key before they build
55//! anything, so `ZUNIONSTORE d 2 z not-a-zset` leaves `d` alone rather than
56//! finding out too late.
57
58use yo_common::num::DIGITS_MAX;
59use yo_common::{Code, Error, Result};
60
61use crate::db::{Db, Holds};
62use crate::elem::Elements;
63use crate::keyspace::Keyspace;
64use crate::scan::Cursor;
65use crate::strings;
66use crate::value::{self, Kind};
67use crate::zset::{Added, Bound, Lex, Member, Zset};
68use crate::zsetops::{self, Aggregate, Op, Operand};
69
70/// Which members a `ZADD` is allowed to touch.
71#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
72pub enum Gate {
73    /// Add and update both. Plain `ZADD`.
74    #[default]
75    Always,
76    /// Only add members that are not there. `ZADD NX`.
77    IfMissing,
78    /// Only update members that are there. `ZADD XX`.
79    IfPresent,
80}
81
82/// Which way a `ZADD` is allowed to move a score it already has.
83#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
84pub enum Move {
85    /// Either way. Plain `ZADD`.
86    #[default]
87    Any,
88    /// Only up. `ZADD GT`.
89    Up,
90    /// Only down. `ZADD LT`.
91    Down,
92}
93
94/// What a `ZADD` was asked to do.
95///
96/// `NX` with `GT` or `LT` is refused by the parser rather than here, because a
97/// gate that only lets new members through and a rule about which way an
98/// existing score may move cannot both apply to the same member and Redis calls
99/// that a syntax error.
100#[derive(Debug, Clone, Copy, Default)]
101pub struct ZAdd {
102    /// `NX` or `XX`.
103    pub gate: Gate,
104    /// `GT` or `LT`.
105    pub only: Move,
106    /// `CH`, which counts changed scores as well as new members.
107    pub changed: bool,
108}
109
110/// Which end of a sorted set a command works from.
111#[derive(Debug, Clone, Copy, PartialEq, Eq)]
112pub enum From {
113    /// The lowest score. `ZPOPMIN`, and `ZMPOP MIN`.
114    Min,
115    /// The highest score. `ZPOPMAX`, and `ZMPOP MAX`.
116    Max,
117}
118
119/// What a range is measured in.
120#[derive(Debug, Clone, Copy)]
121pub enum By<'a> {
122    /// Positions, where a negative one counts from the end. `ZRANGE k 0 -1`.
123    Rank {
124        /// The first position, inclusive.
125        start: i64,
126        /// The last position, inclusive.
127        stop: i64,
128    },
129    /// Scores. `ZRANGEBYSCORE`, and `ZRANGE ... BYSCORE`.
130    Score {
131        /// The lowest score wanted.
132        min: Bound,
133        /// The highest score wanted.
134        max: Bound,
135    },
136    /// Members, which is only meaningful when every score is the same.
137    /// `ZRANGEBYLEX`, and `ZRANGE ... BYLEX`.
138    Lex {
139        /// The first member wanted.
140        min: Lex<'a>,
141        /// The last member wanted.
142        max: Lex<'a>,
143    },
144}
145
146/// A run of members, however it was asked for.
147///
148/// `min` and `max` are always the low end and the high end of the range itself,
149/// whichever order the command wrote them in. `REV` reverses the walk, it does
150/// not reverse the range, which is why `ZRANGEBYSCORE k 1 5` and
151/// `ZREVRANGEBYSCORE k 5 1` cover the same members.
152#[derive(Debug, Clone, Copy)]
153pub struct Query<'a> {
154    /// What the range is measured in.
155    pub by: By<'a>,
156    /// Walk from the high end down. `REV`, and the `ZREV` spellings.
157    pub rev: bool,
158    /// How many to skip once the range is found. `LIMIT`'s first number.
159    pub offset: usize,
160    /// How many to take, or all of them. `LIMIT`'s second number, where Redis's
161    /// negative count means all.
162    pub count: Option<usize>,
163}
164
165impl<'a> Query<'a> {
166    /// A plain `ZRANGE key start stop`.
167    #[must_use]
168    pub const fn rank(start: i64, stop: i64) -> Query<'a> {
169        Query {
170            by: By::Rank { start, stop },
171            rev: false,
172            offset: 0,
173            count: None,
174        }
175    }
176
177    /// A plain `ZRANGEBYSCORE key min max`.
178    #[must_use]
179    pub const fn score(min: Bound, max: Bound) -> Query<'a> {
180        Query {
181            by: By::Score { min, max },
182            rev: false,
183            offset: 0,
184            count: None,
185        }
186    }
187
188    /// A plain `ZRANGEBYLEX key min max`.
189    #[must_use]
190    pub const fn lex(min: Lex<'a>, max: Lex<'a>) -> Query<'a> {
191        Query {
192            by: By::Lex { min, max },
193            rev: false,
194            offset: 0,
195            count: None,
196        }
197    }
198
199    /// The same query walked from the other end.
200    #[must_use]
201    pub const fn rev(mut self, rev: bool) -> Query<'a> {
202        self.rev = rev;
203        self
204    }
205
206    /// The same query with a `LIMIT` on it.
207    #[must_use]
208    pub const fn limit(mut self, offset: usize, count: Option<usize>) -> Query<'a> {
209        self.offset = offset;
210        self.count = count;
211        self
212    }
213}
214
215/// Where a run of members starts, how long it is, and which way it goes.
216///
217/// This is what every range command reduces to, and it is a rank rather than a
218/// pair of members because the tree answers in ranks. `from` is the first member
219/// the walk hands over, so on a reverse walk it is the high end and the walk
220/// counts down.
221#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
222pub struct Window {
223    /// The rank the walk starts at.
224    pub from: usize,
225    /// How many members the walk hands over.
226    pub count: usize,
227    /// Whether the walk counts down.
228    pub rev: bool,
229}
230
231/// A score that is not a number, which no command may store.
232fn nan() -> Error {
233    Error::new(Code::Invalid, "resulting score is not a number (NaN)")
234}
235
236impl Keyspace {
237    /// `ZADD key [NX|XX] [GT|LT] [CH] score member [score member ...]`.
238    ///
239    /// Answers how many members were added, or how many were added or changed if
240    /// `CH` was given, which is the only thing `CH` does.
241    ///
242    /// The pairs arrive as an iterator and not a slice, the way `SADD`'s members
243    /// do, because the wire layer has them as positions in the connection's read
244    /// buffer and a slice would mean collecting them first. It is walked more
245    /// than once, which is why it has to be `Clone`.
246    pub fn zadd<'m, I>(&mut self, key: &[u8], pairs: I, opts: ZAdd) -> Result<usize>
247    where
248        I: Iterator<Item = (f64, &'m [u8])> + Clone,
249    {
250        for (score, m) in pairs.clone() {
251            strings::check_len(key, m.len())?;
252            if score.is_nan() {
253                return Err(nan());
254            }
255        }
256        let at = match self.zset_slot(key)? {
257            Some(at) => at,
258            None => {
259                // A key is not created for a `ZADD XX` that cannot add anything,
260                // and it is not created for an empty pair list either. Redis's
261                // parser rejects `ZADD k` before it gets this far, but the
262                // embedded API has no parser in front of it and an empty sorted
263                // set left behind would be a key that exists and holds nothing.
264                if opts.gate == Gate::IfPresent || pairs.clone().next().is_none() {
265                    return Ok(0);
266                }
267                self.new_zset(key)
268            }
269        };
270        let limits = self.zset_limits;
271        let z = self
272            .zsets
273            .get_mut(at)
274            .expect("the record points at its body");
275        let mut added = 0usize;
276        let mut changed = 0usize;
277        for (score, m) in pairs {
278            let Some(want) = gated(z, m, score, opts) else {
279                continue;
280            };
281            // The element table being full means twenty four million members in
282            // one key. Nothing was stored, so nothing is counted.
283            if z.add(m, score, &limits) == Added::Full {
284                continue;
285            }
286            match want {
287                Added::New => added += 1,
288                Added::Changed => changed += 1,
289                _ => {}
290            }
291        }
292        // A `ZADD XX` on a key that did not exist never got here, so the only
293        // way to be holding an empty sorted set is a pair list that turned out
294        // to be empty after the gate, and that key was ours to make.
295        if z.is_empty() {
296            self.drop_key(key);
297        }
298        Ok(if opts.changed { added + changed } else { added })
299    }
300
301    /// `ZADD key ... INCR score member`, and `ZINCRBY key increment member`.
302    ///
303    /// Answers the member's new score, or nothing at all if a gate refused it,
304    /// which is the nil `ZADD INCR` replies with and is why this is one method
305    /// rather than an `INCR` flag on [`Keyspace::zadd`] that would have to
306    /// return two different shapes.
307    ///
308    /// `ZINCRBY` is this with no gate, where the answer is never nil.
309    pub fn zincrby(
310        &mut self,
311        key: &[u8],
312        member: &[u8],
313        by: f64,
314        opts: ZAdd,
315    ) -> Result<Option<f64>> {
316        strings::check_len(key, member.len())?;
317        if by.is_nan() {
318            return Err(nan());
319        }
320        let at = match self.zset_slot(key)? {
321            Some(at) => at,
322            None => {
323                if opts.gate == Gate::IfPresent {
324                    return Ok(None);
325                }
326                self.new_zset(key)
327            }
328        };
329        let limits = self.zset_limits;
330        let z = self
331            .zsets
332            .get_mut(at)
333            .expect("the record points at its body");
334        let now = z.score(member);
335        let want = now.unwrap_or(0.0) + by;
336        // Infinity plus its opposite. Redis refuses this and leaves the score
337        // alone rather than storing a NaN nothing could ever compare against.
338        if want.is_nan() {
339            if z.is_empty() {
340                self.drop_key(key);
341            }
342            return Err(nan());
343        }
344        let allowed = match (now, opts.gate, opts.only) {
345            (Some(_), Gate::IfMissing, _) | (None, Gate::IfPresent, _) => false,
346            (Some(was), _, Move::Up) => want > was,
347            (Some(was), _, Move::Down) => want < was,
348            _ => true,
349        };
350        if !allowed || z.add(member, want, &limits) == Added::Full {
351            if z.is_empty() {
352                self.drop_key(key);
353            }
354            return Ok(None);
355        }
356        Ok(Some(want))
357    }
358
359    /// `ZCARD key`.
360    pub fn zcard(&mut self, key: &[u8]) -> Result<usize> {
361        Ok(match self.zset_slot(key)? {
362            Some(at) => self.zset_at(at).len(),
363            None => 0,
364        })
365    }
366
367    /// `ZSCORE key member`.
368    pub fn zscore(&mut self, key: &[u8], member: &[u8]) -> Result<Option<f64>> {
369        Ok(match self.zset_slot(key)? {
370            Some(at) => self.zset_at(at).score(member),
371            None => None,
372        })
373    }
374
375    /// `ZMSCORE key member [member ...]`, which is `ZSCORE` in bulk.
376    ///
377    /// One key lookup for the whole call rather than one per member, which is
378    /// the only reason the command exists.
379    pub fn zmscore<'m>(
380        &mut self,
381        key: &[u8],
382        members: impl Iterator<Item = &'m [u8]>,
383        out: &mut Vec<Option<f64>>,
384    ) -> Result<()> {
385        out.clear();
386        let Some(at) = self.zset_slot(key)? else {
387            out.extend(members.map(|_| None));
388            return Ok(());
389        };
390        let z = self.zset_at(at);
391        out.extend(members.map(|m| z.score(m)));
392        Ok(())
393    }
394
395    /// `ZREM key member [member ...]`. Answers how many were there.
396    ///
397    /// A sorted set that loses its last member loses its key too, because an
398    /// empty sorted set does not exist in Redis.
399    pub fn zrem<'m>(
400        &mut self,
401        key: &[u8],
402        members: impl Iterator<Item = &'m [u8]>,
403    ) -> Result<usize> {
404        let Some(at) = self.zset_slot(key)? else {
405            return Ok(0);
406        };
407        let z = self
408            .zsets
409            .get_mut(at)
410            .expect("the record points at its body");
411        let mut gone = 0;
412        for m in members {
413            if z.remove(m) {
414                gone += 1;
415            }
416        }
417        if z.is_empty() {
418            self.drop_key(key);
419        }
420        Ok(gone)
421    }
422
423    /// `ZRANK key member [WITHSCORE]`, and `ZREVRANK` with `rev` set.
424    ///
425    /// The score comes back whether it was asked for or not, because finding the
426    /// rank already read it and handing it over costs nothing.
427    pub fn zrank(&mut self, key: &[u8], member: &[u8], rev: bool) -> Result<Option<(usize, f64)>> {
428        let Some(at) = self.zset_slot(key)? else {
429            return Ok(None);
430        };
431        let z = self.zset_at(at);
432        let Some(rank) = z.rank(member) else {
433            return Ok(None);
434        };
435        let score = z.score(member).unwrap_or(0.0);
436        Ok(Some((if rev { z.len() - 1 - rank } else { rank }, score)))
437    }
438
439    /// How many members a query covers, and where they start.
440    ///
441    /// Every range command starts here. It is separate from [`Keyspace::zwalk`]
442    /// because a RESP array writes its length before its members, and a reply
443    /// that collected the members in order to count them would allocate on the
444    /// read path.
445    pub fn zwindow(&mut self, key: &[u8], q: &Query<'_>) -> Result<Window> {
446        let Some(at) = self.zset_slot(key)? else {
447            return Ok(Window::default());
448        };
449        Ok(window(self.zset_at(at), q))
450    }
451
452    /// Hand over the members a window covers, in order, without collecting them.
453    ///
454    /// The window is the caller's rather than the query's, so that a caller that
455    /// has already asked [`Keyspace::zwindow`] does not compute it twice, and so
456    /// that `ZRANGESTORE` can walk a window it has already narrowed.
457    pub fn zwalk<F>(&mut self, key: &[u8], w: Window, f: F) -> Result<()>
458    where
459        F: FnMut(Member<'_>, f64),
460    {
461        let Some(at) = self.zset_slot(key)? else {
462            return Ok(());
463        };
464        self.zset_at(at).walk(w.from, w.count, w.rev, f);
465        Ok(())
466    }
467
468    /// `ZCOUNT`, `ZLEXCOUNT`, and the count half of any other range command.
469    pub fn zcount(&mut self, key: &[u8], q: &Query<'_>) -> Result<usize> {
470        Ok(self.zwindow(key, q)?.count)
471    }
472
473    /// `ZREMRANGEBYRANK`, `ZREMRANGEBYSCORE` and `ZREMRANGEBYLEX`, which differ
474    /// only in what the query was measured in.
475    ///
476    /// Answers how many went. The window is taken out from its high end down, so
477    /// that every rank still to be removed is the rank it was when the window
478    /// was worked out.
479    pub fn zremrange(&mut self, key: &[u8], q: &Query<'_>) -> Result<usize> {
480        let Some(at) = self.zset_slot(key)? else {
481            return Ok(0);
482        };
483        let w = window(self.zset_at(at), q);
484        let z = self
485            .zsets
486            .get_mut(at)
487            .expect("the record points at its body");
488        // A reverse window starts at its high end, so normalise to the low one
489        // and then count down from the top of it either way.
490        let low = if w.rev { w.from + 1 - w.count } else { w.from };
491        for i in (0..w.count).rev() {
492            z.remove_at(low + i);
493        }
494        if z.is_empty() {
495            self.drop_key(key);
496        }
497        Ok(w.count)
498    }
499
500    /// `ZPOPMIN key [count]` and `ZPOPMAX key [count]`.
501    ///
502    /// Nothing is collected. The member at the end is handed to `f`, which
503    /// writes it wherever it is going, and only then is it removed, which is why
504    /// this does not allocate where `SPOP` has to.
505    pub fn zpop<F>(&mut self, key: &[u8], end: From, count: usize, mut f: F) -> Result<usize>
506    where
507        F: FnMut(Member<'_>, f64),
508    {
509        let Some(at) = self.zset_slot(key)? else {
510            return Ok(0);
511        };
512        let z = self
513            .zsets
514            .get_mut(at)
515            .expect("the record points at its body");
516        let count = count.min(z.len());
517        for _ in 0..count {
518            // Always rank zero or the last rank, because taking one out moves
519            // everything above it down and the next one to go is at the same
520            // place again.
521            let rank = if end == From::Min { 0 } else { z.len() - 1 };
522            let Some((m, s)) = z.at(rank) else { break };
523            f(m, s);
524            z.remove_at(rank);
525        }
526        if z.is_empty() {
527            self.drop_key(key);
528        }
529        Ok(count)
530    }
531
532    /// `ZPOPMIN key` and `ZPOPMAX key`, as an owned member for a caller that has
533    /// nowhere to write it yet.
534    ///
535    /// This is what `BZPOPMIN` needs: a worker that has been parked has no reply
536    /// buffer open at the moment the member becomes available, so this one has
537    /// to allocate where [`Keyspace::zpop`] does not.
538    pub fn zpop_one(&mut self, key: &[u8], end: From) -> Result<Option<(Vec<u8>, f64)>> {
539        let mut got = None;
540        let mut name = [0u8; yo_common::num::DIGITS_MAX];
541        self.zpop(key, end, 1, |m, s| {
542            got = Some((member_bytes(m, &mut name).to_vec(), s));
543        })?;
544        Ok(got)
545    }
546
547    /// `ZRANDMEMBER key [count]`.
548    ///
549    /// A positive count draws without replacement and answers at most as many as
550    /// there are, and a negative one draws with replacement and answers exactly
551    /// as many as asked for, which is Redis's rule and is why the count is
552    /// signed here rather than paired with a flag.
553    ///
554    /// The draw without replacement is a partial shuffle of the row numbers and
555    /// not a retry loop, because a retry loop on a count near the size of the set
556    /// spends most of its time drawing members it already has.
557    pub fn zrandmember<F>(&mut self, key: &[u8], count: i64, mut f: F) -> Result<usize>
558    where
559        F: FnMut(Member<'_>, f64),
560    {
561        let Some(at) = self.zset_slot(key)? else {
562            return Ok(0);
563        };
564        let len = self.zset_at(at).len();
565        if len == 0 || count == 0 {
566            return Ok(0);
567        }
568        if count < 0 {
569            let want = count.unsigned_abs() as usize;
570            for _ in 0..want {
571                let pick = self.rng.below(len);
572                let Some((m, s)) = self.zset_at(at).pick(pick) else {
573                    break;
574                };
575                f(m, s);
576            }
577            return Ok(want);
578        }
579        let want = (count as usize).min(len);
580        // Whole set, in storage order, which is what Redis does for a count at
581        // or over the size and which saves shuffling in order to hand back
582        // everything anyway.
583        if want == len {
584            for i in 0..len {
585                let Some((m, s)) = self.zset_at(at).pick(i) else {
586                    break;
587                };
588                f(m, s);
589            }
590            return Ok(len);
591        }
592        // The database's index buffer rather than a fresh `Vec`, for the reason
593        // the byte scratch exists: a partial shuffle needs somewhere to hold the
594        // permutation while it draws from it, and building that somewhere out of
595        // the allocator on every `ZRANDMEMBER` is a malloc and a free on a
596        // command a sampler sends in a loop. Taken out and put back, so an early
597        // return leaves it as it was found.
598        // `yo_alloc::high_water` because this is the buffer reaching a size it
599        // has not been asked for before, which happens once per largest sorted
600        // set the database has been sampled from and never again.
601        let mut rows = std::mem::take(&mut self.rows);
602        rows.clear();
603        yo_alloc::high_water(|| rows.extend(0..len));
604        for i in 0..want {
605            let pick = i + self.rng.below(len - i);
606            rows.swap(i, pick);
607            let Some((m, s)) = self.zset_at(at).pick(rows[i]) else {
608                break;
609            };
610            f(m, s);
611        }
612        self.rows = rows;
613        Ok(want)
614    }
615
616    /// `ZSCAN key cursor [COUNT count]`.
617    ///
618    /// A small sorted set comes back whole with a cursor of [`Cursor::END`], the
619    /// same guarantee `SSCAN` and `HSCAN` give, because a listpack has no stable
620    /// position to resume from and 128 members is not worth a resume.
621    pub fn zscan<F>(&mut self, key: &[u8], cursor: Cursor, count: usize, f: F) -> Result<Cursor>
622    where
623        F: FnMut(Member<'_>, f64),
624    {
625        let Some(at) = self.zset_slot(key)? else {
626            return Ok(Cursor::END);
627        };
628        Ok(self.zset_at(at).scan(cursor, count, f))
629    }
630
631    /// `ZUNION`, `ZINTER` and `ZDIFF`, which differ only in `op`.
632    ///
633    /// The members come out in rank order, which means the result has to be put
634    /// in order before any of it can be handed over, and that is what the return
635    /// value's ordering costs. `ZINTERCARD` exists precisely because counting
636    /// does not need any of that, and it does not come through here.
637    pub fn zsetop<'k, F>(
638        &mut self,
639        op: Op,
640        keys: impl Iterator<Item = &'k [u8]>,
641        weights: &[f64],
642        agg: Aggregate,
643        f: F,
644    ) -> Result<usize>
645    where
646        F: FnMut(Member<'_>, f64),
647    {
648        let slots = self.operand_slots(keys)?;
649        let got = zsetops::gather(op, &self.operands_of(&slots), weights, agg);
650        let limits = self.zset_limits;
651        let Some(z) = Zset::from_elements(got, &limits) else {
652            return Ok(0);
653        };
654        let len = z.len();
655        z.walk(0, len, false, f);
656        Ok(len)
657    }
658
659    /// `ZUNIONSTORE`, `ZINTERSTORE` and `ZDIFFSTORE`.
660    ///
661    /// The destination is allowed to be one of the sources, which is safe for
662    /// the reason `SINTERSTORE` gives: the result is built whole before the
663    /// destination is touched, so nothing here writes over a body that is still
664    /// being read.
665    pub fn zsetop_store<'k>(
666        &mut self,
667        op: Op,
668        destination: &[u8],
669        keys: impl Iterator<Item = &'k [u8]>,
670        weights: &[f64],
671        agg: Aggregate,
672    ) -> Result<usize> {
673        let slots = self.operand_slots(keys)?;
674        let got = zsetops::gather(op, &self.operands_of(&slots), weights, agg);
675        let limits = self.zset_limits;
676        Ok(self.put_zset(destination, Zset::from_elements(got, &limits)))
677    }
678
679    /// `ZINTERCARD numkeys key [key ...] [LIMIT limit]`.
680    ///
681    /// Nothing is stored and no score is worked out, and a limit stops the walk
682    /// as soon as it is reached, which is the only reason this is not `ZINTER`
683    /// with the members thrown away.
684    pub fn zintercard<'k>(
685        &mut self,
686        keys: impl Iterator<Item = &'k [u8]>,
687        limit: usize,
688    ) -> Result<usize> {
689        let slots = self.operand_slots(keys)?;
690        Ok(zsetops::intercard(&self.operands_of(&slots), limit))
691    }
692
693    /// `ZRANGESTORE destination source <the arguments of ZRANGE>`.
694    ///
695    /// The window is copied rather than moved, because the destination may be
696    /// the source and because the source keeps its members either way. An empty
697    /// window deletes the destination, which is what `ZRANGESTORE d s 5 1` does
698    /// and is the same rule every store form follows.
699    pub fn zrangestore(
700        &mut self,
701        destination: &[u8],
702        source: &[u8],
703        q: &Query<'_>,
704    ) -> Result<usize> {
705        let built = match self.zset_slot(source)? {
706            None => None,
707            Some(at) => {
708                let z = self.zset_at(at);
709                let w = window(z, q);
710                let mut got = Elements::with_capacity(w.count.max(16));
711                let mut digits = [0u8; DIGITS_MAX];
712                // In rank order whichever way the query walked, because the
713                // destination is a sorted set and puts them in score order
714                // regardless. `REV` decides which members are taken, not what
715                // order they end up in.
716                let from = if w.rev { w.from + 1 - w.count } else { w.from };
717                z.walk(from, w.count, false, |m, s| {
718                    let _ = got.insert(member_bytes(m, &mut digits), s);
719                });
720                let limits = self.zset_limits;
721                Zset::from_elements(got, &limits)
722            }
723        };
724        Ok(self.put_zset(destination, built))
725    }
726
727    /// Where each key is and what type it holds, keeping the ones that are not
728    /// there in place.
729    ///
730    /// In place and not dropped, because `WEIGHTS` is positional: a missing
731    /// second key still has a second weight, and closing the gap would hand
732    /// every later input the wrong one.
733    fn operand_slots<'k>(
734        &mut self,
735        keys: impl Iterator<Item = &'k [u8]>,
736    ) -> Result<Vec<Option<(Kind, u32)>>> {
737        let mut out = Vec::with_capacity(keys.size_hint().0);
738        for key in keys {
739            out.push(self.live_slot_either(key, Kind::Zset, Kind::Set)?);
740        }
741        Ok(out)
742    }
743
744    /// The bodies those slots point at, as things the algebra can ask questions
745    /// of.
746    fn operands_of(&self, slots: &[Option<(Kind, u32)>]) -> Vec<Operand<'_>> {
747        slots
748            .iter()
749            .map(|got| match got {
750                Some((Kind::Zset, at)) => Operand::Zset(self.zset_at(*at)),
751                Some((Kind::Set, at)) => {
752                    Operand::Set(self.sets.get(*at).expect("the record points at its body"))
753                }
754                _ => Operand::Missing,
755            })
756            .collect()
757    }
758
759    /// Put a sorted set under `key`, replacing whatever was there.
760    ///
761    /// Nothing means delete the key, because an empty sorted set does not exist.
762    /// That is what makes `ZINTERSTORE d a b` with an empty intersection delete
763    /// `d` and answer zero rather than leave something `EXISTS` says one for.
764    ///
765    /// Whatever the key held is freed first, through the one funnel, and any
766    /// deadline it had goes with it, because the value under the key is not the
767    /// value the expiry was set on.
768    pub(crate) fn put_zset(&mut self, key: &[u8], z: Option<Zset>) -> usize {
769        let Some(z) = z else {
770            self.drop_key(key);
771            return 0;
772        };
773        self.free_body(key);
774        let len = z.len();
775        let at = self.zsets.insert(z);
776        let record = value::slot_record_len(false);
777        self.write_rec(key, record, |out| {
778            value::write_slot_record(out, Kind::Zset, at, None);
779        });
780        self.bodies += 1;
781        len
782    }
783
784    /// The slot `key`'s sorted set is in, or `None` if there is no such key.
785    #[inline]
786    pub(crate) fn zset_slot(&mut self, key: &[u8]) -> Result<Option<u32>> {
787        self.live_slot(key, Kind::Zset)
788    }
789
790    /// The body in a slot the record pointed at.
791    ///
792    /// Panicking here means a record outlived its body, which is the bug the
793    /// two invariants in [`crate::sets`] are there to make impossible.
794    #[inline]
795    pub(crate) fn zset_at(&self, at: u32) -> &Zset {
796        self.zsets.get(at).expect("the record points at its body")
797    }
798
799    /// Make an empty sorted set under `key` and answer which slot it went in.
800    ///
801    /// No hint, unlike a set. A sorted set starts packed whatever is going into
802    /// it, and the first `ZADD` that crosses either threshold promotes it, so
803    /// counting the pairs in advance would only move the same work earlier.
804    fn new_zset(&mut self, key: &[u8]) -> u32 {
805        // The body and, every so often, the slab that holds it. See
806        // `yo_alloc::first_touch` for why this is the one allocation a command
807        // is allowed to make.
808        let at = yo_alloc::first_touch(|| self.zsets.insert(Zset::new()));
809        let len = value::slot_record_len(false);
810        self.write_rec(key, len, |out| {
811            value::write_slot_record(out, Kind::Zset, at, None);
812        });
813        self.bodies += 1;
814        at
815    }
816}
817
818/// Where an input of the algebra lives on a striped database: which stripe it is
819/// on, what it is holding, and the slot the body is in.
820type Home = (usize, Kind, u32);
821
822impl Db {
823    /// `ZUNION`, `ZINTER` and `ZDIFF` over a database of any width.
824    ///
825    /// The keys are asked whether they share a stripe before anything else, and
826    /// when they do the whole command is handed to that stripe. A width one
827    /// database always takes that path and so does a hash tagged group on a wide
828    /// one, so only keys that are genuinely spread out pay for the two passes
829    /// below.
830    pub fn zsetop<'k, F>(
831        &self,
832        op: Op,
833        keys: impl Iterator<Item = &'k [u8]> + Clone,
834        weights: &[f64],
835        agg: Aggregate,
836        f: F,
837    ) -> Result<usize>
838    where
839        F: FnMut(Member<'_>, f64),
840    {
841        if let Some(home) = self.one_stripe(keys.clone()) {
842            return self.hold_stripe(home).zsetop(op, keys, weights, agg, f);
843        }
844        // Stripe zero comes along for the limits the result is built under,
845        // which belong to no key here. See [`Db::zset_limits`].
846        let mut held = self.hold_operands(keys.clone(), Some(0));
847        let slots = self.operand_slots(&mut held, keys)?;
848        let got = zsetops::gather(op, &operands_of(&held, &slots), weights, agg);
849        let limits = held.stripe(0).zset_limits;
850        let Some(z) = Zset::from_elements(got, &limits) else {
851            return Ok(0);
852        };
853        let len = z.len();
854        z.walk(0, len, false, f);
855        Ok(len)
856    }
857
858    /// `ZUNIONSTORE`, `ZINTERSTORE` and `ZDIFFSTORE`.
859    ///
860    /// The destination is allowed to be one of the sources here too, and for the
861    /// same reason: the whole result is built before the destination is touched,
862    /// so no body is written over while it is still being read, whichever stripe
863    /// it is on.
864    pub fn zsetop_store<'k>(
865        &self,
866        op: Op,
867        destination: &'k [u8],
868        keys: impl Iterator<Item = &'k [u8]> + Clone,
869        weights: &[f64],
870        agg: Aggregate,
871    ) -> Result<usize> {
872        if let Some(home) = self.one_stripe(std::iter::once(destination).chain(keys.clone())) {
873            return self
874                .hold_stripe(home)
875                .zsetop_store(op, destination, keys, weights, agg);
876        }
877        let onto = self.stripe_of(destination);
878        let mut held = self.hold_operands(keys.clone(), Some(onto));
879        let slots = self.operand_slots(&mut held, keys)?;
880        let got = zsetops::gather(op, &operands_of(&held, &slots), weights, agg);
881        // The destination's stripe's limits, since that is where the result is
882        // going to live.
883        let limits = held.stripe(onto).zset_limits;
884        let built = Zset::from_elements(got, &limits);
885        Ok(held.stripe_mut(onto).put_zset(destination, built))
886    }
887
888    /// `ZINTERCARD numkeys key [key ...] [LIMIT limit]`.
889    pub fn zintercard<'k>(
890        &self,
891        keys: impl Iterator<Item = &'k [u8]> + Clone,
892        limit: usize,
893    ) -> Result<usize> {
894        if let Some(home) = self.one_stripe(keys.clone()) {
895            return self.hold_stripe(home).zintercard(keys, limit);
896        }
897        let mut held = self.hold_operands(keys.clone(), None);
898        let slots = self.operand_slots(&mut held, keys)?;
899        Ok(zsetops::intercard(&operands_of(&held, &slots), limit))
900    }
901
902    /// `ZRANGESTORE destination source <the arguments of ZRANGE>`.
903    ///
904    /// Two keys and one window, so when they are on different stripes the window
905    /// is walked out of the source's stripe into a table of its own and the
906    /// sorted set that comes of it is put on the destination's. The source keeps
907    /// its members either way, which is what makes the copy the right shape even
908    /// when the two keys are the same key.
909    pub fn zrangestore(&self, destination: &[u8], source: &[u8], q: &Query<'_>) -> Result<usize> {
910        let (onto, home) = (self.stripe_of(destination), self.stripe_of(source));
911        if onto == home {
912            return self.hold_stripe(onto).zrangestore(destination, source, q);
913        }
914        // Both at once and in stripe order rather than one and then the other,
915        // because the source is read into the table, the destination's limits
916        // decide how the table is built, and the destination is written from
917        // it. Held for the whole of that, so the source cannot be added to
918        // after it was read.
919        let mut held = self.hold_many([home, onto].into_iter());
920        // Out of the stripe and then let go of, because a match keeps whatever
921        // it is looking at alive for the whole of itself and the arm below
922        // wants this same stripe again.
923        let slot = held.stripe_mut(home).zset_slot(source)?;
924        let built = match slot {
925            None => None,
926            Some(at) => {
927                let z = held.stripe(home).zset_at(at);
928                let w = window(z, q);
929                let mut got = Elements::with_capacity(w.count.max(16));
930                let mut digits = [0u8; DIGITS_MAX];
931                let from = if w.rev { w.from + 1 - w.count } else { w.from };
932                z.walk(from, w.count, false, |m, s| {
933                    let _ = got.insert(member_bytes(m, &mut digits), s);
934                });
935                let limits = held.stripe(onto).zset_limits;
936                Zset::from_elements(got, &limits)
937            }
938        };
939        Ok(held.stripe_mut(onto).put_zset(destination, built))
940    }
941
942    /// Every stripe the algebra names, held at once, in stripe order, with one
943    /// more for wherever the answer is going.
944    ///
945    /// Before anything is resolved rather than after, because a slot number is
946    /// only good while the stripe it came from is held: let go of it and the key
947    /// can be deleted and the slot handed to something else, and what was read
948    /// back would be a different body under the same number.
949    fn hold_operands<'k>(
950        &self,
951        keys: impl Iterator<Item = &'k [u8]>,
952        onto: Option<usize>,
953    ) -> Holds<'_> {
954        let named = keys.map(|key| self.stripe_of(key));
955        self.hold_many(named.chain(onto))
956    }
957
958    /// Reap and resolve every input key, in order, to the stripe and slot its
959    /// body is in.
960    ///
961    /// As [`Keyspace::operand_slots`], down to keeping a key that is not there in
962    /// place rather than dropping it, because `WEIGHTS` is positional. Each key
963    /// is resolved on the stripe it is on, out of the ones already held, which
964    /// is the only difference, and it is a pass of its own because reaping a key
965    /// wants its stripe mutably and reading a body wants it shared.
966    fn operand_slots<'k>(
967        &self,
968        held: &mut Holds<'_>,
969        keys: impl Iterator<Item = &'k [u8]>,
970    ) -> Result<Vec<Option<Home>>> {
971        let mut out = Vec::with_capacity(keys.size_hint().0);
972        for key in keys {
973            let stripe = self.stripe_of(key);
974            let got = held
975                .stripe_mut(stripe)
976                .live_slot_either(key, Kind::Zset, Kind::Set)?;
977            out.push(got.map(|(kind, at)| (stripe, kind, at)));
978        }
979        Ok(out)
980    }
981}
982
983/// The bodies those slots point at, as things the algebra can ask questions of.
984fn operands_of<'h>(held: &'h Holds<'_>, slots: &[Option<Home>]) -> Vec<Operand<'h>> {
985    slots
986        .iter()
987        .map(|got| match got {
988            Some((stripe, Kind::Zset, at)) => Operand::Zset(held.stripe(*stripe).zset_at(*at)),
989            Some((stripe, Kind::Set, at)) => Operand::Set(
990                held.stripe(*stripe)
991                    .sets
992                    .get(*at)
993                    .expect("the record points at its body"),
994            ),
995            _ => Operand::Missing,
996        })
997        .collect()
998}
999
1000/// What a `ZADD` of one pair would do, or nothing if a gate refuses it.
1001///
1002/// This is worked out before the add rather than from what the add answered,
1003/// because `GT` and `LT` have to see the old score to know whether the new one
1004/// is allowed at all, and by then the add has already stored it.
1005fn gated(z: &Zset, member: &[u8], score: f64, opts: ZAdd) -> Option<Added> {
1006    match z.score(member) {
1007        None => match opts.gate {
1008            Gate::IfPresent => None,
1009            // A new member is added whatever `GT` or `LT` say, because there is
1010            // no old score for them to be about.
1011            _ => Some(Added::New),
1012        },
1013        Some(was) => {
1014            if opts.gate == Gate::IfMissing {
1015                return None;
1016            }
1017            let ok = match opts.only {
1018                Move::Any => true,
1019                Move::Up => score > was,
1020                Move::Down => score < was,
1021            };
1022            if !ok {
1023                return None;
1024            }
1025            // An unchanged score is not a change, which is what stops `CH` from
1026            // counting a member that was written over with what it already had.
1027            Some(if score == was {
1028                Added::Same
1029            } else {
1030                Added::Changed
1031            })
1032        }
1033    }
1034}
1035
1036/// Turn a query into the run of ranks it covers.
1037fn window(z: &Zset, q: &Query<'_>) -> Window {
1038    let len = z.len();
1039    let range = match q.by {
1040        By::Rank { start, stop } => {
1041            // A rank range is already in the direction it will be walked, so a
1042            // reverse one counts its ends from the top rather than being found
1043            // and then flipped.
1044            let (from, count) = rank_span(start, stop, len);
1045            if q.rev {
1046                return apply_limit(
1047                    Window {
1048                        from: len - from - 1,
1049                        count,
1050                        rev: true,
1051                    },
1052                    q,
1053                    true,
1054                );
1055            }
1056            return apply_limit(
1057                Window {
1058                    from,
1059                    count,
1060                    rev: false,
1061                },
1062                q,
1063                true,
1064            );
1065        }
1066        By::Score { min, max } => z.window_by_score(min, max),
1067        By::Lex { min, max } => z.window_by_lex(min, max),
1068    };
1069    let count = range.end - range.start;
1070    let from = if q.rev {
1071        range.end.saturating_sub(1)
1072    } else {
1073        range.start
1074    };
1075    apply_limit(
1076        Window {
1077            from,
1078            count,
1079            rev: q.rev,
1080        },
1081        q,
1082        false,
1083    )
1084}
1085
1086/// Move a window along by `LIMIT`'s offset and cut it to its count.
1087///
1088/// `ZRANGE key 0 -1 REV LIMIT 1 2` is the second and third from the top, so the
1089/// offset walks in the direction of the walk and not up the ranks.
1090///
1091/// A rank range has already had its ends clamped, and Redis does not accept a
1092/// `LIMIT` on one anyway, so `skip` says to leave it alone rather than the
1093/// caller passing an offset of zero and a count of none and hoping.
1094fn apply_limit(w: Window, q: &Query<'_>, skip: bool) -> Window {
1095    if skip {
1096        return w;
1097    }
1098    let offset = q.offset.min(w.count);
1099    let count = q.count.unwrap_or(usize::MAX).min(w.count - offset);
1100    let from = if w.rev {
1101        w.from.saturating_sub(offset)
1102    } else {
1103        w.from + offset
1104    };
1105    Window {
1106        from,
1107        count,
1108        rev: w.rev,
1109    }
1110}
1111
1112/// Turn an inclusive `start` and `stop` into an offset from the front and a
1113/// count, clamping every out of range case rather than erroring.
1114///
1115/// The same rule `LRANGE` uses, and the same reason: a start before the front is
1116/// the front, a stop past the end is the end, and a start after the stop is
1117/// nothing at all.
1118fn rank_span(start: i64, stop: i64, len: usize) -> (usize, usize) {
1119    if len == 0 {
1120        return (0, 0);
1121    }
1122    let len = len as i64;
1123    let from = if start < 0 {
1124        (len + start).max(0)
1125    } else {
1126        start.min(len)
1127    };
1128    let to = if stop < 0 {
1129        len + stop
1130    } else {
1131        stop.min(len - 1)
1132    };
1133    if to < from {
1134        return (from as usize, 0);
1135    }
1136    (from as usize, (to - from + 1) as usize)
1137}
1138
1139/// The bytes of a member, which for one the listpack stored as an integer are
1140/// the caller's buffer.
1141pub(crate) fn member_bytes<'a>(
1142    m: Member<'a>,
1143    digits: &'a mut [u8; yo_common::num::DIGITS_MAX],
1144) -> &'a [u8] {
1145    match m {
1146        Member::Str(s) => s,
1147        Member::Int(n) => yo_common::num::i64_digits(digits, n),
1148    }
1149}
1150
1151#[cfg(test)]
1152mod tests {
1153    use super::*;
1154    use yo_common::num::DIGITS_MAX;
1155
1156    fn ks() -> Keyspace {
1157        Keyspace::new()
1158    }
1159
1160    /// Add a run of pairs the plain way.
1161    fn add(k: &mut Keyspace, key: &[u8], pairs: &[(f64, &[u8])]) -> usize {
1162        k.zadd(key, pairs.iter().copied(), ZAdd::default()).unwrap()
1163    }
1164
1165    /// Every member a query covers, as names.
1166    fn names(k: &mut Keyspace, key: &[u8], q: &Query<'_>) -> Vec<String> {
1167        let w = k.zwindow(key, q).unwrap();
1168        let mut out = Vec::new();
1169        let mut digits = [0u8; DIGITS_MAX];
1170        k.zwalk(key, w, |m, _| {
1171            out.push(String::from_utf8(member_bytes(m, &mut digits).to_vec()).unwrap());
1172        })
1173        .unwrap();
1174        assert_eq!(
1175            out.len(),
1176            w.count,
1177            "the window said {} and the walk gave {}",
1178            w.count,
1179            out.len()
1180        );
1181        out
1182    }
1183
1184    /// A positive count under the size of the set does a partial shuffle, and
1185    /// the permutation that needs used to be a fresh `Vec` every call. Sampling
1186    /// is a thing callers do in a loop, so the first call is allowed to grow the
1187    /// buffer and none after it may allocate at all.
1188    #[test]
1189    fn zrandmember_stops_allocating_once_its_buffer_is_grown() {
1190        let mut k = ks();
1191        add(
1192            &mut k,
1193            b"z",
1194            &[(1.0, b"a"), (2.0, b"b"), (3.0, b"c"), (4.0, b"d")],
1195        );
1196        k.zrandmember(b"z", 2, |_, _| {}).expect("a sorted set");
1197        let (_, allocs) = crate::tally::counted(|| {
1198            for _ in 0..100 {
1199                k.zrandmember(b"z", 2, |_, _| {}).expect("a sorted set");
1200            }
1201        });
1202        assert_eq!(
1203            allocs, 0,
1204            "zrandmember allocated {allocs} times in a hundred"
1205        );
1206    }
1207
1208    /// The buffer is put back on the way out, so the call after it sees a
1209    /// buffer rather than an empty one, and both answer with what they were
1210    /// asked for.
1211    #[test]
1212    fn zrandmember_hands_its_buffer_back() {
1213        let mut k = ks();
1214        add(&mut k, b"z", &[(1.0, b"a"), (2.0, b"b"), (3.0, b"c")]);
1215        for _ in 0..3 {
1216            let mut seen = Vec::new();
1217            let n = k
1218                .zrandmember(b"z", 2, |m, _| {
1219                    let mut digits = [0u8; DIGITS_MAX];
1220                    seen.push(String::from_utf8(member_bytes(m, &mut digits).to_vec()).unwrap());
1221                })
1222                .expect("a sorted set");
1223            assert_eq!(n, 2);
1224            assert_eq!(seen.len(), 2);
1225            assert_ne!(seen[0], seen[1], "a positive count draws without replacing");
1226        }
1227    }
1228
1229    #[test]
1230    fn a_missing_key_is_an_empty_sorted_set() {
1231        let mut k = ks();
1232        assert_eq!(k.zcard(b"nope").unwrap(), 0);
1233        assert_eq!(k.zscore(b"nope", b"m").unwrap(), None);
1234        assert_eq!(k.zrank(b"nope", b"m", false).unwrap(), None);
1235        assert_eq!(k.zrem(b"nope", [b"m".as_slice()].into_iter()).unwrap(), 0);
1236        assert_eq!(
1237            names(&mut k, b"nope", &Query::rank(0, -1)),
1238            Vec::<String>::new()
1239        );
1240        assert!(!k.exists(b"nope"));
1241    }
1242
1243    #[test]
1244    fn a_key_holding_something_else_is_a_wrongtype() {
1245        let mut k = ks();
1246        k.set(b"s", b"hello", strings::SetOptions::default())
1247            .unwrap();
1248        assert_eq!(k.zcard(b"s").unwrap_err().code(), Code::WrongType);
1249        assert_eq!(
1250            k.zadd(b"s", [(1.0, b"m".as_slice())].into_iter(), ZAdd::default())
1251                .unwrap_err()
1252                .code(),
1253            Code::WrongType
1254        );
1255        assert_eq!(k.zscore(b"s", b"m").unwrap_err().code(), Code::WrongType);
1256    }
1257
1258    #[test]
1259    fn adding_answers_how_many_were_new() {
1260        let mut k = ks();
1261        assert_eq!(add(&mut k, b"z", &[(1.0, b"a"), (2.0, b"b")]), 2);
1262        assert_eq!(add(&mut k, b"z", &[(1.0, b"a"), (3.0, b"c")]), 1);
1263        assert_eq!(k.zcard(b"z").unwrap(), 3);
1264        assert_eq!(k.zscore(b"z", b"c").unwrap(), Some(3.0));
1265        assert_eq!(k.kind_of(b"z").map(Kind::name), Some("zset"));
1266        assert_eq!(k.encoding_name(b"z"), Some("listpack"));
1267    }
1268
1269    #[test]
1270    fn the_ch_flag_counts_moved_scores_too() {
1271        let mut k = ks();
1272        add(&mut k, b"z", &[(1.0, b"a"), (2.0, b"b")]);
1273        let ch = ZAdd {
1274            changed: true,
1275            ..ZAdd::default()
1276        };
1277        // One score moved, one stayed, one member is new.
1278        let pairs = [
1279            (9.0, b"a".as_slice()),
1280            (2.0, b"b".as_slice()),
1281            (3.0, b"c".as_slice()),
1282        ];
1283        assert_eq!(k.zadd(b"z", pairs.into_iter(), ch).unwrap(), 2);
1284        assert_eq!(k.zadd(b"z", pairs.into_iter(), ZAdd::default()).unwrap(), 0);
1285    }
1286
1287    #[test]
1288    fn the_gates_let_the_right_members_through() {
1289        let mut k = ks();
1290        add(&mut k, b"z", &[(1.0, b"a")]);
1291        let nx = ZAdd {
1292            gate: Gate::IfMissing,
1293            ..ZAdd::default()
1294        };
1295        let xx = ZAdd {
1296            gate: Gate::IfPresent,
1297            ..ZAdd::default()
1298        };
1299        assert_eq!(
1300            k.zadd(b"z", [(5.0, b"a".as_slice())].into_iter(), nx)
1301                .unwrap(),
1302            0
1303        );
1304        assert_eq!(k.zscore(b"z", b"a").unwrap(), Some(1.0));
1305        assert_eq!(
1306            k.zadd(b"z", [(5.0, b"b".as_slice())].into_iter(), nx)
1307                .unwrap(),
1308            1
1309        );
1310        assert_eq!(
1311            k.zadd(b"z", [(7.0, b"c".as_slice())].into_iter(), xx)
1312                .unwrap(),
1313            0
1314        );
1315        assert_eq!(k.zscore(b"z", b"c").unwrap(), None);
1316        assert_eq!(
1317            k.zadd(b"z", [(7.0, b"a".as_slice())].into_iter(), xx)
1318                .unwrap(),
1319            0
1320        );
1321        assert_eq!(k.zscore(b"z", b"a").unwrap(), Some(7.0));
1322        // XX on a key that does not exist does not create it.
1323        assert_eq!(
1324            k.zadd(b"gone", [(1.0, b"a".as_slice())].into_iter(), xx)
1325                .unwrap(),
1326            0
1327        );
1328        assert!(!k.exists(b"gone"));
1329    }
1330
1331    #[test]
1332    fn gt_and_lt_only_move_a_score_one_way() {
1333        let mut k = ks();
1334        add(&mut k, b"z", &[(5.0, b"a")]);
1335        let gt = ZAdd {
1336            only: Move::Up,
1337            changed: true,
1338            ..ZAdd::default()
1339        };
1340        let lt = ZAdd {
1341            only: Move::Down,
1342            changed: true,
1343            ..ZAdd::default()
1344        };
1345        assert_eq!(
1346            k.zadd(b"z", [(3.0, b"a".as_slice())].into_iter(), gt)
1347                .unwrap(),
1348            0
1349        );
1350        assert_eq!(k.zscore(b"z", b"a").unwrap(), Some(5.0));
1351        assert_eq!(
1352            k.zadd(b"z", [(9.0, b"a".as_slice())].into_iter(), gt)
1353                .unwrap(),
1354            1
1355        );
1356        assert_eq!(k.zscore(b"z", b"a").unwrap(), Some(9.0));
1357        assert_eq!(
1358            k.zadd(b"z", [(9.0, b"a".as_slice())].into_iter(), lt)
1359                .unwrap(),
1360            0
1361        );
1362        assert_eq!(
1363            k.zadd(b"z", [(2.0, b"a".as_slice())].into_iter(), lt)
1364                .unwrap(),
1365            1
1366        );
1367        assert_eq!(k.zscore(b"z", b"a").unwrap(), Some(2.0));
1368        // A member that is not there is added whatever GT says.
1369        assert_eq!(
1370            k.zadd(b"z", [(1.0, b"new".as_slice())].into_iter(), gt)
1371                .unwrap(),
1372            1
1373        );
1374    }
1375
1376    #[test]
1377    fn incrementing_adds_to_what_is_there_or_to_nothing() {
1378        let mut k = ks();
1379        let plain = ZAdd::default();
1380        assert_eq!(k.zincrby(b"z", b"a", 5.0, plain).unwrap(), Some(5.0));
1381        assert_eq!(k.zincrby(b"z", b"a", -2.5, plain).unwrap(), Some(2.5));
1382        assert_eq!(k.zscore(b"z", b"a").unwrap(), Some(2.5));
1383        let nx = ZAdd {
1384            gate: Gate::IfMissing,
1385            ..ZAdd::default()
1386        };
1387        assert_eq!(k.zincrby(b"z", b"a", 1.0, nx).unwrap(), None);
1388        assert_eq!(k.zscore(b"z", b"a").unwrap(), Some(2.5));
1389        let xx = ZAdd {
1390            gate: Gate::IfPresent,
1391            ..ZAdd::default()
1392        };
1393        assert_eq!(k.zincrby(b"z", b"never", 1.0, xx).unwrap(), None);
1394        assert!(k.zscore(b"z", b"never").unwrap().is_none());
1395        let gt = ZAdd {
1396            only: Move::Up,
1397            ..ZAdd::default()
1398        };
1399        assert_eq!(k.zincrby(b"z", b"a", -1.0, gt).unwrap(), None);
1400        assert_eq!(k.zincrby(b"z", b"a", 1.0, gt).unwrap(), Some(3.5));
1401    }
1402
1403    #[test]
1404    fn a_score_that_is_not_a_number_is_refused() {
1405        let mut k = ks();
1406        let plain = ZAdd::default();
1407        assert_eq!(
1408            k.zadd(b"z", [(f64::NAN, b"a".as_slice())].into_iter(), plain)
1409                .unwrap_err()
1410                .code(),
1411            Code::Invalid
1412        );
1413        assert!(!k.exists(b"z"));
1414        k.zincrby(b"z", b"a", f64::INFINITY, plain).unwrap();
1415        let err = k.zincrby(b"z", b"a", f64::NEG_INFINITY, plain).unwrap_err();
1416        assert_eq!(err.code(), Code::Invalid);
1417        // The score is left exactly as it was rather than stored as a NaN.
1418        assert_eq!(k.zscore(b"z", b"a").unwrap(), Some(f64::INFINITY));
1419    }
1420
1421    #[test]
1422    fn a_sorted_set_that_loses_its_last_member_loses_its_key() {
1423        let mut k = ks();
1424        add(&mut k, b"z", &[(1.0, b"a"), (2.0, b"b")]);
1425        assert_eq!(
1426            k.zrem(b"z", [b"a".as_slice(), b"b".as_slice()].into_iter())
1427                .unwrap(),
1428            2
1429        );
1430        assert!(!k.exists(b"z"));
1431        assert_eq!(k.zcard(b"z").unwrap(), 0);
1432    }
1433
1434    #[test]
1435    fn ranks_count_from_both_ends() {
1436        let mut k = ks();
1437        add(&mut k, b"z", &[(1.0, b"a"), (2.0, b"b"), (3.0, b"c")]);
1438        assert_eq!(k.zrank(b"z", b"a", false).unwrap(), Some((0, 1.0)));
1439        assert_eq!(k.zrank(b"z", b"c", false).unwrap(), Some((2, 3.0)));
1440        assert_eq!(k.zrank(b"z", b"a", true).unwrap(), Some((2, 1.0)));
1441        assert_eq!(k.zrank(b"z", b"c", true).unwrap(), Some((0, 3.0)));
1442        assert_eq!(k.zrank(b"z", b"nope", false).unwrap(), None);
1443    }
1444
1445    #[test]
1446    fn a_rank_range_clamps_every_way_it_can_be_wrong() {
1447        let mut k = ks();
1448        add(&mut k, b"z", &[(1.0, b"a"), (2.0, b"b"), (3.0, b"c")]);
1449        assert_eq!(names(&mut k, b"z", &Query::rank(0, -1)), ["a", "b", "c"]);
1450        assert_eq!(names(&mut k, b"z", &Query::rank(1, 1)), ["b"]);
1451        assert_eq!(names(&mut k, b"z", &Query::rank(-2, -1)), ["b", "c"]);
1452        assert_eq!(names(&mut k, b"z", &Query::rank(-99, 99)), ["a", "b", "c"]);
1453        assert_eq!(
1454            names(&mut k, b"z", &Query::rank(2, 1)),
1455            Vec::<String>::new()
1456        );
1457        assert_eq!(
1458            names(&mut k, b"z", &Query::rank(5, 9)),
1459            Vec::<String>::new()
1460        );
1461        assert_eq!(
1462            names(&mut k, b"z", &Query::rank(0, -1).rev(true)),
1463            ["c", "b", "a"]
1464        );
1465        assert_eq!(
1466            names(&mut k, b"z", &Query::rank(0, 1).rev(true)),
1467            ["c", "b"]
1468        );
1469    }
1470
1471    #[test]
1472    fn a_score_range_walks_either_way_and_takes_a_limit() {
1473        let mut k = ks();
1474        add(
1475            &mut k,
1476            b"z",
1477            &[(1.0, b"a"), (2.0, b"b"), (3.0, b"c"), (4.0, b"d")],
1478        );
1479        let all = Query::score(
1480            Bound::closed(f64::NEG_INFINITY),
1481            Bound::closed(f64::INFINITY),
1482        );
1483        assert_eq!(names(&mut k, b"z", &all), ["a", "b", "c", "d"]);
1484        assert_eq!(names(&mut k, b"z", &all.rev(true)), ["d", "c", "b", "a"]);
1485        assert_eq!(
1486            names(
1487                &mut k,
1488                b"z",
1489                &Query::score(Bound::closed(2.0), Bound::closed(3.0))
1490            ),
1491            ["b", "c"]
1492        );
1493        assert_eq!(
1494            names(
1495                &mut k,
1496                b"z",
1497                &Query::score(Bound::open(2.0), Bound::open(4.0))
1498            ),
1499            ["c"]
1500        );
1501        // LIMIT walks in the direction of the walk, so the reverse one skips
1502        // from the top.
1503        assert_eq!(names(&mut k, b"z", &all.limit(1, Some(2))), ["b", "c"]);
1504        assert_eq!(
1505            names(&mut k, b"z", &all.rev(true).limit(1, Some(2))),
1506            ["c", "b"]
1507        );
1508        assert_eq!(
1509            names(&mut k, b"z", &all.limit(9, Some(2))),
1510            Vec::<String>::new()
1511        );
1512        assert_eq!(names(&mut k, b"z", &all.limit(2, None)), ["c", "d"]);
1513        assert_eq!(
1514            k.zcount(b"z", &Query::score(Bound::closed(2.0), Bound::closed(3.0)))
1515                .unwrap(),
1516            2
1517        );
1518    }
1519
1520    #[test]
1521    fn a_member_range_orders_by_member_when_every_score_is_the_same() {
1522        let mut k = ks();
1523        add(
1524            &mut k,
1525            b"z",
1526            &[(0.0, b"a"), (0.0, b"b"), (0.0, b"c"), (0.0, b"d")],
1527        );
1528        assert_eq!(
1529            names(&mut k, b"z", &Query::lex(Lex::Min, Lex::Max)),
1530            ["a", "b", "c", "d"]
1531        );
1532        assert_eq!(
1533            names(&mut k, b"z", &Query::lex(Lex::Incl(b"b"), Lex::Incl(b"c"))),
1534            ["b", "c"]
1535        );
1536        assert_eq!(
1537            names(&mut k, b"z", &Query::lex(Lex::Excl(b"a"), Lex::Excl(b"d"))),
1538            ["b", "c"]
1539        );
1540        assert_eq!(
1541            names(&mut k, b"z", &Query::lex(Lex::Min, Lex::Max).rev(true)),
1542            ["d", "c", "b", "a"]
1543        );
1544        assert_eq!(
1545            k.zcount(b"z", &Query::lex(Lex::Incl(b"b"), Lex::Max))
1546                .unwrap(),
1547            3
1548        );
1549    }
1550
1551    #[test]
1552    fn removing_a_range_takes_out_exactly_what_the_walk_would_have_given() {
1553        let mut k = ks();
1554        for (q, left) in [
1555            (Query::rank(0, 1), vec!["c", "d", "e"]),
1556            (Query::rank(-2, -1), vec!["a", "b", "c"]),
1557            (
1558                Query::score(Bound::closed(2.0), Bound::closed(4.0)),
1559                vec!["a", "e"],
1560            ),
1561            (
1562                Query::lex(Lex::Incl(b"b"), Lex::Excl(b"d")),
1563                vec!["a", "d", "e"],
1564            ),
1565            (Query::rank(0, -1).rev(true), vec![]),
1566        ] {
1567            k.del(b"z");
1568            add(
1569                &mut k,
1570                b"z",
1571                &[
1572                    (1.0, b"a"),
1573                    (2.0, b"b"),
1574                    (3.0, b"c"),
1575                    (4.0, b"d"),
1576                    (5.0, b"e"),
1577                ],
1578            );
1579            let want = names(&mut k, b"z", &q).len();
1580            assert_eq!(k.zremrange(b"z", &q).unwrap(), want, "{q:?}");
1581            assert_eq!(names(&mut k, b"z", &Query::rank(0, -1)), left, "{q:?}");
1582        }
1583        // The key went with the last member.
1584        assert!(!k.exists(b"z"));
1585    }
1586
1587    #[test]
1588    fn popping_takes_from_the_end_it_was_told_to() {
1589        let mut k = ks();
1590        add(&mut k, b"z", &[(1.0, b"a"), (2.0, b"b"), (3.0, b"c")]);
1591        let mut got = Vec::new();
1592        let mut digits = [0u8; DIGITS_MAX];
1593        k.zpop(b"z", From::Min, 2, |m, s| {
1594            got.push((
1595                String::from_utf8(member_bytes(m, &mut digits).to_vec()).unwrap(),
1596                s,
1597            ));
1598        })
1599        .unwrap();
1600        assert_eq!(got, [("a".to_string(), 1.0), ("b".to_string(), 2.0)]);
1601        assert_eq!(
1602            k.zpop_one(b"z", From::Max).unwrap(),
1603            Some((b"c".to_vec(), 3.0))
1604        );
1605        assert!(!k.exists(b"z"));
1606        // A count past the end takes what there is and no more.
1607        add(&mut k, b"z", &[(1.0, b"a")]);
1608        assert_eq!(k.zpop(b"z", From::Max, 99, |_, _| {}).unwrap(), 1);
1609        assert!(!k.exists(b"z"));
1610        assert_eq!(k.zpop_one(b"z", From::Min).unwrap(), None);
1611    }
1612
1613    #[test]
1614    fn a_random_draw_is_with_or_without_replacement_by_the_sign_of_the_count() {
1615        let mut k = ks();
1616        add(&mut k, b"z", &[(1.0, b"a"), (2.0, b"b"), (3.0, b"c")]);
1617        let mut digits = [0u8; DIGITS_MAX];
1618        let mut seen = Vec::new();
1619        k.zrandmember(b"z", 2, |m, _| {
1620            seen.push(String::from_utf8(member_bytes(m, &mut digits).to_vec()).unwrap());
1621        })
1622        .unwrap();
1623        assert_eq!(seen.len(), 2);
1624        seen.sort();
1625        seen.dedup();
1626        assert_eq!(seen.len(), 2, "a positive count does not repeat a member");
1627        // A count past the size is the whole set and not a repeat of it.
1628        let mut all = Vec::new();
1629        k.zrandmember(b"z", 99, |m, _| {
1630            all.push(String::from_utf8(member_bytes(m, &mut digits).to_vec()).unwrap());
1631        })
1632        .unwrap();
1633        all.sort();
1634        assert_eq!(all, ["a", "b", "c"]);
1635        // A negative count answers exactly as many as asked for and may repeat.
1636        let mut many = 0;
1637        k.zrandmember(b"z", -10, |_, _| many += 1).unwrap();
1638        assert_eq!(many, 10);
1639        assert_eq!(k.zrandmember(b"nope", 3, |_, _| {}).unwrap(), 0);
1640    }
1641
1642    #[test]
1643    fn a_scan_of_either_band_walks_every_member_once() {
1644        let mut k = ks();
1645        for entries in [8usize, 4096] {
1646            k.del(b"z");
1647            let pairs: Vec<(f64, Vec<u8>)> = (0..entries)
1648                .map(|i| (i as f64, format!("m{i:05}").into_bytes()))
1649                .collect();
1650            k.zadd(
1651                b"z",
1652                pairs.iter().map(|(s, m)| (*s, m.as_slice())),
1653                ZAdd::default(),
1654            )
1655            .unwrap();
1656            let mut seen = Vec::new();
1657            let mut digits = [0u8; DIGITS_MAX];
1658            let mut cursor = Cursor::START;
1659            loop {
1660                cursor = k
1661                    .zscan(b"z", cursor, 16, |m, _| {
1662                        seen.push(
1663                            String::from_utf8(member_bytes(m, &mut digits).to_vec()).unwrap(),
1664                        );
1665                    })
1666                    .unwrap();
1667                if cursor.is_end() {
1668                    break;
1669                }
1670            }
1671            seen.sort();
1672            seen.dedup();
1673            assert_eq!(seen.len(), entries, "{entries} members");
1674        }
1675    }
1676
1677    #[test]
1678    fn a_big_sorted_set_promotes_and_still_answers_every_rank() {
1679        let mut k = ks();
1680        let pairs: Vec<(f64, Vec<u8>)> = (0..5_000)
1681            .map(|i| (f64::from(i), format!("m{i:05}").into_bytes()))
1682            .collect();
1683        assert_eq!(
1684            k.zadd(
1685                b"z",
1686                pairs.iter().map(|(s, m)| (*s, m.as_slice())),
1687                ZAdd::default()
1688            )
1689            .unwrap(),
1690            5_000
1691        );
1692        assert_eq!(k.encoding_name(b"z"), Some("skiplist"));
1693        assert_eq!(k.zcard(b"z").unwrap(), 5_000);
1694        assert_eq!(
1695            k.zrank(b"z", b"m02500", false).unwrap(),
1696            Some((2_500, 2500.0))
1697        );
1698        let q = Query::score(Bound::closed(1000.0), Bound::open(1010.0));
1699        assert_eq!(k.zcount(b"z", &q).unwrap(), 10);
1700        assert_eq!(
1701            names(&mut k, b"z", &q).first().map(String::as_str),
1702            Some("m01000")
1703        );
1704        // Everything still lines up after a few thousand removals.
1705        assert_eq!(k.zremrange(b"z", &Query::rank(0, 2_499)).unwrap(), 2_500);
1706        assert_eq!(k.zcard(b"z").unwrap(), 2_500);
1707        assert_eq!(k.zrank(b"z", b"m02500", false).unwrap(), Some((0, 2500.0)));
1708        assert_eq!(k.zrank(b"z", b"m00000", false).unwrap(), None);
1709    }
1710
1711    #[test]
1712    fn a_deadline_on_a_sorted_set_leaves_its_members_alone() {
1713        let mut k = ks();
1714        add(&mut k, b"z", &[(1.0, b"a"), (2.0, b"b")]);
1715        assert!(k.set_expiry(b"z", Some(u64::MAX / 2)));
1716        assert_eq!(k.zcard(b"z").unwrap(), 2);
1717        assert_eq!(k.zscore(b"z", b"b").unwrap(), Some(2.0));
1718        assert!(k.set_expiry(b"z", None));
1719        assert_eq!(k.zcard(b"z").unwrap(), 2);
1720    }
1721
1722    #[test]
1723    fn writing_a_string_over_a_sorted_set_gives_its_body_back() {
1724        let mut k = ks();
1725        add(&mut k, b"z", &[(1.0, b"a")]);
1726        let held = k.memory_bytes();
1727        k.set(b"z", b"now a string", strings::SetOptions::default())
1728            .unwrap();
1729        assert_eq!(k.kind_of(b"z").map(Kind::name), Some("string"));
1730        assert!(
1731            k.memory_bytes() < held,
1732            "the sorted set's body was not freed"
1733        );
1734    }
1735
1736    /// Everything in a key, in rank order, as pairs.
1737    fn all(k: &mut Keyspace, key: &[u8]) -> Vec<(String, f64)> {
1738        let q = Query::rank(0, -1);
1739        let w = k.zwindow(key, &q).unwrap();
1740        let mut out = Vec::new();
1741        let mut digits = [0u8; DIGITS_MAX];
1742        k.zwalk(key, w, |m, s| {
1743            out.push((
1744                String::from_utf8(member_bytes(m, &mut digits).to_vec()).unwrap(),
1745                s,
1746            ));
1747        })
1748        .unwrap();
1749        out
1750    }
1751
1752    /// What one of the non storing forms answers, in the order it answered.
1753    fn got(
1754        k: &mut Keyspace,
1755        op: Op,
1756        keys: &[&[u8]],
1757        weights: &[f64],
1758        agg: Aggregate,
1759    ) -> Vec<(String, f64)> {
1760        let mut out = Vec::new();
1761        let mut digits = [0u8; DIGITS_MAX];
1762        let n = k
1763            .zsetop(op, keys.iter().copied(), weights, agg, |m, s| {
1764                out.push((
1765                    String::from_utf8(member_bytes(m, &mut digits).to_vec()).unwrap(),
1766                    s,
1767                ));
1768            })
1769            .unwrap();
1770        assert_eq!(out.len(), n, "the count and the walk disagree");
1771        out
1772    }
1773
1774    #[test]
1775    fn a_union_adds_the_scores_of_a_member_that_is_in_both() {
1776        let mut k = ks();
1777        add(&mut k, b"a", &[(1.0, b"x"), (2.0, b"y")]);
1778        add(&mut k, b"b", &[(10.0, b"y"), (3.0, b"z")]);
1779        assert_eq!(
1780            got(&mut k, Op::Union, &[b"a", b"b"], &[], Aggregate::Sum),
1781            [("x".into(), 1.0), ("z".into(), 3.0), ("y".into(), 12.0)]
1782        );
1783    }
1784
1785    #[test]
1786    fn an_intersection_keeps_only_what_every_input_has() {
1787        let mut k = ks();
1788        add(&mut k, b"a", &[(1.0, b"x"), (2.0, b"y"), (3.0, b"z")]);
1789        add(&mut k, b"b", &[(5.0, b"y"), (5.0, b"z")]);
1790        add(&mut k, b"c", &[(7.0, b"z")]);
1791        assert_eq!(
1792            got(&mut k, Op::Inter, &[b"a", b"b", b"c"], &[], Aggregate::Sum),
1793            [("z".into(), 15.0)]
1794        );
1795        assert_eq!(
1796            got(&mut k, Op::Inter, &[b"a", b"b", b"c"], &[], Aggregate::Min),
1797            [("z".into(), 3.0)]
1798        );
1799        assert_eq!(
1800            got(&mut k, Op::Inter, &[b"a", b"b", b"c"], &[], Aggregate::Max),
1801            [("z".into(), 7.0)]
1802        );
1803    }
1804
1805    #[test]
1806    fn a_difference_takes_the_first_input_and_removes_the_rest() {
1807        let mut k = ks();
1808        add(&mut k, b"a", &[(1.0, b"x"), (2.0, b"y"), (3.0, b"z")]);
1809        add(&mut k, b"b", &[(99.0, b"y")]);
1810        assert_eq!(
1811            got(&mut k, Op::Diff, &[b"a", b"b"], &[], Aggregate::Sum),
1812            [("x".into(), 1.0), ("z".into(), 3.0)]
1813        );
1814    }
1815
1816    #[test]
1817    fn a_missing_key_keeps_its_place_so_the_weights_stay_lined_up() {
1818        let mut k = ks();
1819        add(&mut k, b"a", &[(1.0, b"x")]);
1820        add(&mut k, b"c", &[(1.0, b"x")]);
1821        // The second weight belongs to the key that is not there, and the third
1822        // to `c`. Dropping the gap would give `c` the 10 and answer 11.
1823        let out = got(
1824            &mut k,
1825            Op::Union,
1826            &[b"a", b"gone", b"c"],
1827            &[2.0, 10.0, 3.0],
1828            Aggregate::Sum,
1829        );
1830        assert_eq!(out, [("x".into(), 5.0)]);
1831    }
1832
1833    #[test]
1834    fn a_plain_set_counts_as_every_score_being_one() {
1835        let mut k = ks();
1836        add(&mut k, b"z", &[(5.0, b"x")]);
1837        k.sadd(b"s", [b"x".as_slice(), b"y".as_slice()].into_iter())
1838            .unwrap();
1839        assert_eq!(
1840            got(&mut k, Op::Union, &[b"z", b"s"], &[], Aggregate::Sum),
1841            [("y".into(), 1.0), ("x".into(), 6.0)]
1842        );
1843        assert_eq!(
1844            got(&mut k, Op::Inter, &[b"z", b"s"], &[], Aggregate::Min),
1845            [("x".into(), 1.0)]
1846        );
1847    }
1848
1849    #[test]
1850    fn a_store_writes_the_result_and_answers_its_size() {
1851        let mut k = ks();
1852        add(&mut k, b"a", &[(1.0, b"x"), (2.0, b"y")]);
1853        add(&mut k, b"b", &[(10.0, b"y")]);
1854        assert_eq!(
1855            k.zsetop_store(
1856                Op::Union,
1857                b"d",
1858                [b"a".as_slice(), b"b".as_slice()].into_iter(),
1859                &[],
1860                Aggregate::Sum
1861            )
1862            .unwrap(),
1863            2
1864        );
1865        assert_eq!(all(&mut k, b"d"), [("x".into(), 1.0), ("y".into(), 12.0)]);
1866        assert_eq!(k.zscore(b"d", b"y").unwrap(), Some(12.0));
1867        assert_eq!(k.zrank(b"d", b"y", false).unwrap(), Some((1, 12.0)));
1868    }
1869
1870    #[test]
1871    fn a_store_onto_one_of_its_own_sources_still_reads_the_old_body() {
1872        let mut k = ks();
1873        add(&mut k, b"a", &[(1.0, b"x"), (2.0, b"y")]);
1874        add(&mut k, b"b", &[(10.0, b"y")]);
1875        assert_eq!(
1876            k.zsetop_store(
1877                Op::Union,
1878                b"a",
1879                [b"a".as_slice(), b"b".as_slice()].into_iter(),
1880                &[],
1881                Aggregate::Sum
1882            )
1883            .unwrap(),
1884            2
1885        );
1886        assert_eq!(all(&mut k, b"a"), [("x".into(), 1.0), ("y".into(), 12.0)]);
1887    }
1888
1889    #[test]
1890    fn a_store_with_nothing_in_it_deletes_the_destination() {
1891        let mut k = ks();
1892        add(&mut k, b"a", &[(1.0, b"x")]);
1893        add(&mut k, b"b", &[(1.0, b"y")]);
1894        add(&mut k, b"d", &[(1.0, b"old")]);
1895        assert_eq!(
1896            k.zsetop_store(
1897                Op::Inter,
1898                b"d",
1899                [b"a".as_slice(), b"b".as_slice()].into_iter(),
1900                &[],
1901                Aggregate::Sum
1902            )
1903            .unwrap(),
1904            0
1905        );
1906        assert!(!k.exists(b"d"));
1907    }
1908
1909    #[test]
1910    fn a_store_clears_the_deadline_the_destination_was_carrying() {
1911        let mut k = ks();
1912        add(&mut k, b"a", &[(1.0, b"x")]);
1913        add(&mut k, b"d", &[(1.0, b"old")]);
1914        assert!(k.set_expiry(b"d", Some(u64::MAX / 2)));
1915        k.zsetop_store(
1916            Op::Union,
1917            b"d",
1918            [b"a".as_slice()].into_iter(),
1919            &[],
1920            Aggregate::Sum,
1921        )
1922        .unwrap();
1923        assert_eq!(all(&mut k, b"d"), [("x".into(), 1.0)]);
1924        assert_eq!(k.expire_at(b"d"), None);
1925    }
1926
1927    #[test]
1928    fn the_algebra_refuses_a_key_holding_something_that_is_neither() {
1929        let mut k = ks();
1930        add(&mut k, b"a", &[(1.0, b"x")]);
1931        k.set(b"s", b"hello", strings::SetOptions::default())
1932            .unwrap();
1933        assert_eq!(
1934            k.zsetop(
1935                Op::Union,
1936                [b"a".as_slice(), b"s".as_slice()].into_iter(),
1937                &[],
1938                Aggregate::Sum,
1939                |_, _| {}
1940            )
1941            .unwrap_err()
1942            .code(),
1943            Code::WrongType
1944        );
1945        // And the destination is untouched, because the keys are all resolved
1946        // before anything is built.
1947        assert!(!k.exists(b"d"));
1948        assert_eq!(
1949            k.zsetop_store(
1950                Op::Union,
1951                b"d",
1952                [b"a".as_slice(), b"s".as_slice()].into_iter(),
1953                &[],
1954                Aggregate::Sum
1955            )
1956            .unwrap_err()
1957            .code(),
1958            Code::WrongType
1959        );
1960        assert!(!k.exists(b"d"));
1961    }
1962
1963    #[test]
1964    fn intercard_counts_without_building_anything_and_stops_at_the_limit() {
1965        let mut k = ks();
1966        add(
1967            &mut k,
1968            b"a",
1969            &[(1.0, b"w"), (2.0, b"x"), (3.0, b"y"), (4.0, b"z")],
1970        );
1971        add(&mut k, b"b", &[(1.0, b"x"), (1.0, b"y"), (1.0, b"z")]);
1972        assert_eq!(
1973            k.zintercard([b"a".as_slice(), b"b".as_slice()].into_iter(), 0)
1974                .unwrap(),
1975            3
1976        );
1977        assert_eq!(
1978            k.zintercard([b"a".as_slice(), b"b".as_slice()].into_iter(), 2)
1979                .unwrap(),
1980            2
1981        );
1982        assert_eq!(
1983            k.zintercard([b"a".as_slice(), b"gone".as_slice()].into_iter(), 0)
1984                .unwrap(),
1985            0
1986        );
1987    }
1988
1989    #[test]
1990    fn a_range_store_copies_a_window_and_leaves_the_source_alone() {
1991        let mut k = ks();
1992        add(
1993            &mut k,
1994            b"z",
1995            &[(1.0, b"a"), (2.0, b"b"), (3.0, b"c"), (4.0, b"d")],
1996        );
1997        assert_eq!(k.zrangestore(b"d", b"z", &Query::rank(1, 2)).unwrap(), 2);
1998        assert_eq!(all(&mut k, b"d"), [("b".into(), 2.0), ("c".into(), 3.0)]);
1999        assert_eq!(k.zcard(b"z").unwrap(), 4);
2000    }
2001
2002    #[test]
2003    fn a_reverse_range_store_picks_the_same_members_and_orders_them_the_same() {
2004        let mut k = ks();
2005        add(
2006            &mut k,
2007            b"z",
2008            &[(1.0, b"a"), (2.0, b"b"), (3.0, b"c"), (4.0, b"d")],
2009        );
2010        // REV picks which members, not what order they end up in, because the
2011        // destination is a sorted set and a sorted set has one order.
2012        assert_eq!(
2013            k.zrangestore(b"d", b"z", &Query::rank(0, 1).rev(true))
2014                .unwrap(),
2015            2
2016        );
2017        assert_eq!(all(&mut k, b"d"), [("c".into(), 3.0), ("d".into(), 4.0)]);
2018    }
2019
2020    #[test]
2021    fn a_range_store_by_score_takes_the_bounds_the_range_would_have() {
2022        let mut k = ks();
2023        add(
2024            &mut k,
2025            b"z",
2026            &[(1.0, b"a"), (2.0, b"b"), (3.0, b"c"), (4.0, b"d")],
2027        );
2028        let q = Query::score(Bound::closed(2.0), Bound::open(4.0));
2029        assert_eq!(k.zrangestore(b"d", b"z", &q).unwrap(), 2);
2030        assert_eq!(all(&mut k, b"d"), [("b".into(), 2.0), ("c".into(), 3.0)]);
2031    }
2032
2033    #[test]
2034    fn a_range_store_of_an_empty_window_deletes_the_destination() {
2035        let mut k = ks();
2036        add(&mut k, b"z", &[(1.0, b"a")]);
2037        add(&mut k, b"d", &[(1.0, b"old")]);
2038        assert_eq!(k.zrangestore(b"d", b"z", &Query::rank(5, 9)).unwrap(), 0);
2039        assert!(!k.exists(b"d"));
2040        assert_eq!(
2041            k.zrangestore(b"d", b"gone", &Query::rank(0, -1)).unwrap(),
2042            0
2043        );
2044        assert!(!k.exists(b"d"));
2045    }
2046
2047    #[test]
2048    fn a_range_store_onto_its_own_source_keeps_the_window() {
2049        let mut k = ks();
2050        add(
2051            &mut k,
2052            b"z",
2053            &[(1.0, b"a"), (2.0, b"b"), (3.0, b"c"), (4.0, b"d")],
2054        );
2055        assert_eq!(k.zrangestore(b"z", b"z", &Query::rank(1, 2)).unwrap(), 2);
2056        assert_eq!(all(&mut k, b"z"), [("b".into(), 2.0), ("c".into(), 3.0)]);
2057    }
2058
2059    #[test]
2060    fn a_big_result_comes_out_on_the_table_band_in_the_right_order() {
2061        let mut k = ks();
2062        let names: Vec<String> = (0..3000).map(|i| format!("member-{i:05}")).collect();
2063        for (i, name) in names.iter().enumerate() {
2064            add(&mut k, b"a", &[((i % 17) as f64, name.as_bytes())]);
2065        }
2066        for name in names.iter().step_by(3) {
2067            add(&mut k, b"b", &[(100.0, name.as_bytes())]);
2068        }
2069        let n = k
2070            .zsetop_store(
2071                Op::Union,
2072                b"d",
2073                [b"a".as_slice(), b"b".as_slice()].into_iter(),
2074                &[],
2075                Aggregate::Sum,
2076            )
2077            .unwrap();
2078        assert_eq!(n, 3000);
2079        assert_eq!(
2080            k.zset_encoding(b"d").map(crate::zset::Encoding::name),
2081            Some("skiplist")
2082        );
2083        let out = all(&mut k, b"d");
2084        assert_eq!(out.len(), 3000);
2085        let mut want = out.clone();
2086        want.sort_by(|x, y| x.1.partial_cmp(&y.1).unwrap().then_with(|| x.0.cmp(&y.0)));
2087        assert_eq!(out, want, "the result came out in the wrong order");
2088        for (i, name) in names.iter().enumerate() {
2089            let base = (i % 17) as f64;
2090            let want = if i % 3 == 0 { base + 100.0 } else { base };
2091            assert_eq!(k.zscore(b"d", name.as_bytes()).unwrap(), Some(want));
2092        }
2093    }
2094}