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