Skip to main content

yo_kv/
zset.rs

1//! A sorted set: an element table for the score of a member, and a counted tree
2//! for the rank of a score.
3//!
4//! Every sorted set command is one of two questions and this is why there are
5//! two structures. `ZSCORE` and `ZADD` ask what a member's score is, which is a
6//! probe by name and wants a hash table. `ZRANK`, `ZRANGE` and everything with a
7//! range in it ask where a score sits among the others, which is order
8//! statistics and wants a tree. Redis answers both out of a skiplist plus a
9//! dictionary; the dictionary is the same answer as ours and the skiplist is
10//! where `08` section 5's memory went.
11//!
12//! ```text
13//!   small                       everything else
14//! +------------------+   +-------------------------------------+
15//! | listpack         |-->| element table   +   counted tree     |
16//! | member, score,   |   | member -> score     rank -> row      |
17//! | member, score... |   |                                      |
18//! +------------------+   +-------------------------------------+
19//!    to 128 members         one probe                one descent
20//! ```
21//!
22//! # What the tree holds, and what it does not
23//!
24//! The tree holds row numbers, three bytes each, and nothing else. It does not
25//! hold the score. A search asks it for a position and it asks the caller, on
26//! each comparison, where the thing being looked for sits against the element in
27//! a given row, and the caller answers by reading that row's score out of the
28//! element table.
29//!
30//! That is the trade `Y14` asks for. A tree with the score beside the row would
31//! finish a search without leaving the node it is in, and it would cost eleven
32//! bytes an element rather than three, which is a fail however fast it is. What
33//! it costs instead is that a descent touches a handful of rows that are not
34//! next to each other, and the rows it touches on a zipfian draw, which is the
35//! gate cell aki lost, are the hot ones that are in cache anyway.
36//!
37//! # Ties
38//!
39//! Order is by score and then by member, comparing the member bytes, which is
40//! Redis's rule and the reason `ZRANGEBYLEX` works at all: give every member the
41//! same score and the set is ordered by member alone. Two members never tie
42//! completely, so every element has exactly one position and a search for one
43//! lands on it rather than on a run.
44//!
45//! # Renumbering
46//!
47//! The element table is dense, so taking a row out of it moves the last row into
48//! the hole, and one element that nobody asked about gets a new row number on
49//! every removal. The tree is told through [`Rank::set_at`], which needs the
50//! moved element's position, so a `ZREM` is two descents rather than one. The
51//! alternative is holes in the element table, and then the uniform draw
52//! `ZRANDMEMBER` wants has to retry until it lands on a live row, which is fine
53//! at nine tenths full and unbounded on a set that has been drained.
54
55use core::cmp::Ordering;
56use core::ops::Range;
57
58use yo_common::num::{DIGITS_MAX, DOUBLE_MAX, i64_digits, parse_f64, write_double};
59
60use crate::elem::Elements;
61use crate::frozen::{self, Broken};
62use crate::listpack::{self, Listpack};
63use crate::rank::Rank;
64use crate::scan::Cursor;
65
66/// A member: bytes as they lie, or an integer not yet formatted.
67pub type Member<'a> = listpack::Entry<'a>;
68
69/// The packed band, which is Redis's `ZSET_LISTPACK`.
70const FORM_PACKED: u8 = 1;
71/// The table and its tree, written out as members in rank order.
72const FORM_MEMBERS: u8 = 2;
73
74/// Where a sorted set stops being one packed blob.
75#[derive(Debug, Clone, Copy)]
76pub struct Limits {
77    /// At this many members a sorted set stops being a listpack.
78    pub max_listpack_entries: usize,
79    /// A member longer than this cannot go in a listpack.
80    pub max_listpack_value: usize,
81}
82
83impl Limits {
84    /// Redis's defaults: 128 and 64.
85    pub const DEFAULT: Limits = Limits {
86        max_listpack_entries: 128,
87        max_listpack_value: 64,
88    };
89}
90
91impl Default for Limits {
92    fn default() -> Limits {
93        Limits::DEFAULT
94    }
95}
96
97/// Which of the two a sorted set is in, which is what `OBJECT ENCODING` reports.
98#[derive(Debug, Clone, Copy, PartialEq, Eq)]
99pub enum Encoding {
100    /// One packed blob, walked linearly.
101    Listpack,
102    /// The element table and the tree.
103    ///
104    /// The name is Redis's and the structure is not. `12` section 2 records the
105    /// divergence: `OBJECT ENCODING` returns a name clients test for, not a
106    /// claim about what is underneath, and answering with a word no client has
107    /// heard of breaks tools for nothing.
108    Skiplist,
109}
110
111impl Encoding {
112    /// The word `OBJECT ENCODING` replies with.
113    #[must_use]
114    pub const fn name(self) -> &'static str {
115        match self {
116            Encoding::Listpack => "listpack",
117            Encoding::Skiplist => "skiplist",
118        }
119    }
120}
121
122/// What an add did.
123#[derive(Debug, Clone, Copy, PartialEq, Eq)]
124pub enum Added {
125    /// The member was not there and now is. This is what `ZADD` counts.
126    New,
127    /// The member was there with a different score.
128    Changed,
129    /// The member was there with this score already.
130    Same,
131    /// The element table is full, so nothing happened.
132    Full,
133}
134
135/// One end of a score range.
136#[derive(Debug, Clone, Copy)]
137pub struct Bound {
138    /// The score itself, which may be infinite.
139    pub at: f64,
140    /// Whether the score itself is outside the range, which is `ZRANGEBYSCORE`'s
141    /// parenthesis.
142    pub open: bool,
143}
144
145impl Bound {
146    /// A bound that includes its score.
147    #[must_use]
148    pub const fn closed(at: f64) -> Bound {
149        Bound { at, open: false }
150    }
151
152    /// A bound that excludes its score.
153    #[must_use]
154    pub const fn open(at: f64) -> Bound {
155        Bound { at, open: true }
156    }
157}
158
159/// One end of a member range, for the commands that order by member alone.
160#[derive(Debug, Clone, Copy)]
161pub enum Lex<'a> {
162    /// Before every member, which is `-`.
163    Min,
164    /// After every member, which is `+`.
165    Max,
166    /// This member and everything after it, which is `[`.
167    Incl(&'a [u8]),
168    /// Everything after this member, which is `(`.
169    Excl(&'a [u8]),
170}
171
172/// The element table and the tree over it.
173#[derive(Debug, Clone)]
174struct Table {
175    members: Elements<f64>,
176    order: Rank,
177}
178
179/// The two representations.
180#[derive(Debug, Clone)]
181enum Body {
182    /// Member, score, member, score, in order.
183    Packed(Listpack),
184    Table(Table),
185}
186
187/// A set of members, each with a score, ordered by score and then by member.
188#[derive(Debug, Clone)]
189pub struct Zset {
190    body: Body,
191}
192
193impl Default for Zset {
194    fn default() -> Self {
195        Self::new()
196    }
197}
198
199/// Scores compare the way Redis compares them, which is not the way `f64`
200/// compares by default.
201///
202/// `total_cmp` puts negative zero below positive zero and Redis does not, and a
203/// `NaN` never gets here because every command that takes a score refuses one
204/// before this is reached.
205#[inline]
206fn cmp_score(a: f64, b: f64) -> Ordering {
207    a.partial_cmp(&b).unwrap_or(Ordering::Equal)
208}
209
210/// The order the whole structure is in: score first, then the member bytes.
211#[inline]
212fn cmp_key(score: f64, member: &[u8], other_score: f64, other_member: &[u8]) -> Ordering {
213    match cmp_score(score, other_score) {
214        Ordering::Equal => member.cmp(other_member),
215        other => other,
216    }
217}
218
219/// The score an entry holds, whichever way the listpack stored it.
220fn score_of(entry: Member<'_>) -> f64 {
221    match entry {
222        Member::Int(n) => n as f64,
223        // A score that went in as text came from `push_double`, so it parses.
224        // A listpack that did not come from here is checked when it is loaded.
225        Member::Str(s) => parse_f64(s).unwrap_or(0.0),
226    }
227}
228
229/// The bytes of an entry, which for an integer member are the caller's buffer.
230fn bytes_of<'a>(entry: Member<'a>, digits: &'a mut [u8; DIGITS_MAX]) -> &'a [u8] {
231    match entry {
232        Member::Str(s) => s,
233        Member::Int(n) => i64_digits(digits, n),
234    }
235}
236
237impl Zset {
238    /// An empty sorted set, in the packed band.
239    #[must_use]
240    pub fn new() -> Zset {
241        Zset {
242            body: Body::Packed(Listpack::new()),
243        }
244    }
245
246    /// An empty sorted set with room for `hint` members.
247    ///
248    /// A caller that knows the count up front, which is every `RESTORE`, should
249    /// not fill the packed band to its limit and then promote the lot into a
250    /// table. That pays a scan for the member and a scan for the position on
251    /// every one of the first hundred and twenty eight, and then throws the
252    /// listpack away.
253    ///
254    /// The hint is only a hint. Being wrong about it costs a table with more
255    /// room than it needed rather than anything incorrect, and a hint under the
256    /// band limit still gets the band, because a sorted set that turns out to be
257    /// small is the common one and it should stay packed.
258    #[must_use]
259    pub fn with_hint(hint: usize, limits: &Limits) -> Zset {
260        if hint <= limits.max_listpack_entries {
261            return Zset::new();
262        }
263        Zset {
264            body: Body::Table(Table {
265                members: Elements::with_capacity(hint),
266                order: Rank::new(),
267            }),
268        }
269    }
270
271    /// Take a listpack that is already in this band's layout, if it really is.
272    ///
273    /// The payload a `RESTORE` hands over for a small sorted set is Redis's own
274    /// `ZSET_LISTPACK`, which is byte for byte the layout this band uses, so the
275    /// whole load can be the blob moving in rather than a member at a time. That
276    /// is the same argument [`Zset::packed_bytes`] makes on the way out, run
277    /// backwards.
278    ///
279    /// It is worth more coming in than going out. Adding a member at a time
280    /// costs a scan to see whether the member is already there and a second scan
281    /// to find where it belongs, both of them over everything added so far, so a
282    /// hundred member sorted set took nine times as long to restore as a hundred
283    /// field hash did.
284    ///
285    /// The blob comes back on refusal, so a caller that has to walk it after all
286    /// does not have to parse it twice. Refusal covers a count or a member past
287    /// the limits, and it covers a blob that is not in order. This band answers
288    /// a rank query by position and nothing else, so a payload that says it is a
289    /// sorted set while not being sorted has to be rebuilt rather than trusted.
290    /// Checking that is one pass and it is the same pass that rules out
291    /// duplicates, since strictly increasing means no two members compare equal
292    /// on the score and then equal on the bytes.
293    pub(crate) fn from_packed(lp: Listpack, limits: &Limits) -> Result<Zset, Listpack> {
294        let n = lp.len();
295        if n == 0 || !n.is_multiple_of(2) || n / 2 > limits.max_listpack_entries {
296            return Err(lp);
297        }
298        let ok = {
299            let mut walk = lp.iter();
300            let mut prev: Option<(f64, Member<'_>)> = None;
301            let mut before_buf = [0u8; DIGITS_MAX];
302            let mut member_buf = [0u8; DIGITS_MAX];
303            loop {
304                let Some(member) = walk.next() else {
305                    break true;
306                };
307                // The count is even, checked above, so there is always a score
308                // behind a member.
309                let Some(entry) = walk.next() else {
310                    break false;
311                };
312                let score = match entry {
313                    Member::Int(v) => v as f64,
314                    // `score_of` reads a text score with `unwrap_or(0.0)`, which
315                    // is right for a blob this band wrote and wrong for one that
316                    // arrived over the wire, so it is checked here once rather
317                    // than guessed at on every read after.
318                    Member::Str(s) => match parse_f64(s) {
319                        Some(v) => v,
320                        None => break false,
321                    },
322                };
323                let bytes = bytes_of(member, &mut member_buf);
324                if bytes.len() > limits.max_listpack_value {
325                    break false;
326                }
327                if let Some((before, was)) = prev {
328                    let was = bytes_of(was, &mut before_buf);
329                    if cmp_key(before, was, score, bytes) != Ordering::Less {
330                        break false;
331                    }
332                }
333                prev = Some((score, member));
334            }
335        };
336        if ok {
337            Ok(Zset {
338                body: Body::Packed(lp),
339            })
340        } else {
341            Err(lp)
342        }
343    }
344
345    /// Write this sorted set out in a form a device can hold.
346    ///
347    /// The packed band goes out as its own bytes, and the table goes out as its
348    /// members in rank order, each one followed by its score. Rank order is the
349    /// point of writing it that way: the order is the expensive half of a sorted
350    /// set, and a body written in order comes back without a single comparison.
351    ///
352    /// A score is eight raw bytes rather than the text a listpack holds, because
353    /// the table has the double already and formatting it here only to parse it
354    /// on the way back would cost two conversions for no saving.
355    pub fn freeze(&self, out: &mut Vec<u8>) {
356        match &self.body {
357            Body::Packed(lp) => {
358                out.push(FORM_PACKED);
359                out.extend_from_slice(lp.as_bytes());
360            }
361            Body::Table(t) => {
362                out.push(FORM_MEMBERS);
363                frozen::put_uint(out, t.members.len() as u64);
364                // Through the tree's own walk and not through `Zset::at`, which
365                // would descend from the root for every rank and turn a pass
366                // into a count times a depth.
367                for row in t.order.iter_from(0) {
368                    let Some((name, score)) = t.members.at(row as usize) else {
369                        continue;
370                    };
371                    frozen::put_bytes(out, name);
372                    frozen::put_f64(out, *score);
373                }
374            }
375        }
376    }
377
378    /// Read back what [`Zset::freeze`] wrote.
379    ///
380    /// The band a sorted set left in is the band it comes back in, so a value
381    /// that was quiet long enough to be moved out answers `OBJECT ENCODING` with
382    /// the same word it answered before.
383    pub fn thaw(bytes: &[u8]) -> Result<Zset, Broken> {
384        let mut cut = frozen::Cut::new(bytes);
385        match cut.byte()? {
386            FORM_PACKED => Ok(Zset {
387                body: Body::Packed(Listpack::from_bytes(cut.rest()).map_err(|_| Broken::Body)?),
388            }),
389            FORM_MEMBERS => {
390                let n = usize::try_from(cut.uint()?).map_err(|_| Broken::Short)?;
391                // A member costs a length byte and a score costs eight, so a
392                // count larger than what is left cannot be honest and is not
393                // worth an allocation.
394                if n > cut.rest().len() {
395                    return Err(Broken::Body);
396                }
397                let mut table = Table {
398                    members: Elements::with_capacity(n),
399                    order: Rank::new(),
400                };
401                for _ in 0..n {
402                    let name = cut.bytes()?;
403                    let score = cut.f64()?;
404                    let row = table.members.len() as u32;
405                    // A member twice over would leave the table one row short of
406                    // the tree and every rank after it wrong, so a body that
407                    // repeats one is refused rather than half built.
408                    if !matches!(table.members.insert(name, score), Ok(None)) {
409                        return Err(Broken::Body);
410                    }
411                    // Written in rank order, so every member goes on the end and
412                    // the tree compares nothing.
413                    table.order.insert_at(row as usize, row);
414                }
415                Ok(Zset {
416                    body: Body::Table(table),
417                })
418            }
419            _ => Err(Broken::Form),
420        }
421    }
422
423    /// How many members are in here.
424    #[must_use]
425    pub fn len(&self) -> usize {
426        match &self.body {
427            Body::Packed(lp) => lp.len() / 2,
428            Body::Table(t) => t.members.len(),
429        }
430    }
431
432    /// Whether there are no members in here.
433    #[must_use]
434    pub fn is_empty(&self) -> bool {
435        self.len() == 0
436    }
437
438    /// Which representation this is on.
439    #[must_use]
440    pub const fn encoding(&self) -> Encoding {
441        match &self.body {
442            Body::Packed(_) => Encoding::Listpack,
443            Body::Table(_) => Encoding::Skiplist,
444        }
445    }
446
447    /// The bytes behind a sorted set on the packed band, for `DUMP` to copy.
448    ///
449    /// Member and score alternate in here exactly as `ZSET_LISTPACK` wants them,
450    /// because the band was built to Redis's layout, so the payload is these
451    /// bytes with a length in front. `None` on the tree, where there is no blob
452    /// and the members have to be walked.
453    #[inline]
454    pub(crate) fn packed_bytes(&self) -> Option<&[u8]> {
455        match &self.body {
456            Body::Packed(lp) => Some(lp.as_bytes()),
457            Body::Table(_) => None,
458        }
459    }
460
461    /// What this is holding on to, in bytes.
462    #[must_use]
463    pub fn memory_bytes(&self) -> usize {
464        match &self.body {
465            Body::Packed(lp) => lp.byte_len(),
466            Body::Table(t) => t.members.memory_bytes() + t.order.bytes(),
467        }
468    }
469
470    /// The score of a member, or `None` if it is not in here.
471    ///
472    /// `ZSCORE`, and the first half of every `ZADD`.
473    #[must_use]
474    pub fn score(&self, member: &[u8]) -> Option<f64> {
475        match &self.body {
476            Body::Packed(lp) => {
477                let at = lp.find(member, 2)?;
478                lp.get(at + 1).map(score_of)
479            }
480            Body::Table(t) => t.members.get(member).copied(),
481        }
482    }
483
484    /// Put a member in, or move one that is already there.
485    pub fn add(&mut self, member: &[u8], score: f64, limits: &Limits) -> Added {
486        if let Body::Packed(lp) = &mut self.body {
487            if let Some(at) = lp.find(member, 2) {
488                let old = lp.get(at + 1).map_or(0.0, score_of);
489                if cmp_score(old, score) == Ordering::Equal {
490                    return Added::Same;
491                }
492                lp.delete(at, 2);
493                packed_insert(lp, member, score);
494                return Added::Changed;
495            }
496            if lp.len() / 2 < limits.max_listpack_entries
497                && member.len() <= limits.max_listpack_value
498            {
499                packed_insert(lp, member, score);
500                return Added::New;
501            }
502            self.promote();
503        }
504        let Body::Table(t) = &mut self.body else {
505            unreachable!("promoted above")
506        };
507        t.add(member, score)
508    }
509
510    /// Take a member out.
511    ///
512    /// `ZREM`, and the way `ZADD GT` gets rid of nothing at all.
513    pub fn remove(&mut self, member: &[u8]) -> bool {
514        match &mut self.body {
515            Body::Packed(lp) => match lp.find(member, 2) {
516                Some(at) => lp.delete(at, 2),
517                None => false,
518            },
519            Body::Table(t) => {
520                let Some(row) = t.members.index_of(member) else {
521                    return false;
522                };
523                let score = *t.members.get(member).expect("just found");
524                let rank = t.rank_of(row as u32, score, member);
525                t.take(rank, row);
526                true
527            }
528        }
529    }
530
531    /// Where a member sits, counting from the lowest score.
532    ///
533    /// `ZRANK`, and `ZREVRANK` by taking it from the length.
534    #[must_use]
535    pub fn rank(&self, member: &[u8]) -> Option<usize> {
536        match &self.body {
537            Body::Packed(lp) => lp.find(member, 2).map(|at| at / 2),
538            Body::Table(t) => {
539                let row = t.members.index_of(member)?;
540                let score = *t.members.get(member)?;
541                Some(t.rank_of(row as u32, score, member))
542            }
543        }
544    }
545
546    /// The member and score at a rank.
547    ///
548    /// The member is not copied. `ZPOPMIN` writes it into the reply and then
549    /// calls [`Zset::remove_at`] with the same rank, which is why these are two
550    /// methods and not one that hands back an owned name.
551    #[must_use]
552    pub fn at(&self, rank: usize) -> Option<(Member<'_>, f64)> {
553        match &self.body {
554            Body::Packed(lp) => {
555                let member = lp.get(rank * 2)?;
556                let score = lp.get(rank * 2 + 1).map(score_of)?;
557                Some((member, score))
558            }
559            Body::Table(t) => {
560                let row = t.order.row_at(rank)?;
561                let (name, score) = t.members.at(row as usize)?;
562                Some((Member::Str(name), *score))
563            }
564        }
565    }
566
567    /// Take out whatever is at a rank.
568    pub fn remove_at(&mut self, rank: usize) -> bool {
569        match &mut self.body {
570            Body::Packed(lp) => lp.delete(rank * 2, 2),
571            Body::Table(t) => {
572                let Some(row) = t.order.row_at(rank) else {
573                    return false;
574                };
575                t.take(rank, row as usize);
576                true
577            }
578        }
579    }
580
581    /// A member by position in no particular order, for a uniform draw.
582    ///
583    /// `ZRANDMEMBER` wants any member with equal probability and does not care
584    /// which, so on the table this reads a row straight out of the dense array
585    /// rather than descending the tree for a rank nobody asked for.
586    #[must_use]
587    pub fn pick(&self, at: usize) -> Option<(Member<'_>, f64)> {
588        match &self.body {
589            Body::Packed(_) => self.at(at),
590            Body::Table(t) => {
591                let (name, score) = t.members.at(at)?;
592                Some((Member::Str(name), *score))
593            }
594        }
595    }
596
597    /// Walk members in rank order, from a rank, for a count.
598    ///
599    /// Every range command comes through here after working out which ranks it
600    /// wants, because a range by score and a range by member are the same walk
601    /// once the two ends have been found.
602    pub fn walk<F: FnMut(Member<'_>, f64)>(&self, from: usize, count: usize, rev: bool, mut f: F) {
603        let len = self.len();
604        if from >= len || count == 0 {
605            return;
606        }
607        let count = count.min(if rev { from + 1 } else { len - from });
608        match &self.body {
609            Body::Packed(lp) => {
610                for i in 0..count {
611                    let at = if rev { from - i } else { from + i };
612                    let (Some(m), Some(s)) = (lp.get(at * 2), lp.get(at * 2 + 1)) else {
613                        return;
614                    };
615                    f(m, score_of(s));
616                }
617            }
618            Body::Table(t) => {
619                if rev {
620                    for row in t.order.iter_back_from(from).take(count) {
621                        let Some((name, score)) = t.members.at(row as usize) else {
622                            return;
623                        };
624                        f(Member::Str(name), *score);
625                    }
626                } else {
627                    for row in t.order.iter_from(from).take(count) {
628                        let Some((name, score)) = t.members.at(row as usize) else {
629                            return;
630                        };
631                        f(Member::Str(name), *score);
632                    }
633                }
634            }
635        }
636    }
637
638    /// The ranks whose scores fall inside a range.
639    ///
640    /// `ZRANGEBYSCORE`, `ZCOUNT` and `ZREMRANGEBYSCORE` are all this plus a walk
641    /// or a count of what it returns. An empty range comes back as an empty one
642    /// rather than as a pair that has to be checked by the caller.
643    #[must_use]
644    pub fn window_by_score(&self, min: Bound, max: Bound) -> Range<usize> {
645        let start = self.seek(|score, _| {
646            // Everything below the bottom of the range is behind us.
647            let before = match cmp_score(score, min.at) {
648                Ordering::Less => true,
649                Ordering::Equal => min.open,
650                Ordering::Greater => false,
651            };
652            if before {
653                Ordering::Greater
654            } else {
655                Ordering::Less
656            }
657        });
658        let end = self.seek(|score, _| {
659            let inside = match cmp_score(score, max.at) {
660                Ordering::Less => true,
661                Ordering::Equal => !max.open,
662                Ordering::Greater => false,
663            };
664            if inside {
665                Ordering::Greater
666            } else {
667                Ordering::Less
668            }
669        });
670        start..end.max(start)
671    }
672
673    /// The ranks whose members fall inside a range, ignoring scores.
674    ///
675    /// `ZRANGEBYLEX`, which is only meaningful when every member has the same
676    /// score and is nonsense otherwise, exactly as it is in Redis.
677    #[must_use]
678    pub fn window_by_lex(&self, min: Lex<'_>, max: Lex<'_>) -> Range<usize> {
679        let start = self.seek(|_, member| match min {
680            Lex::Min => Ordering::Less,
681            Lex::Max => Ordering::Greater,
682            Lex::Incl(at) => {
683                if member < at {
684                    Ordering::Greater
685                } else {
686                    Ordering::Less
687                }
688            }
689            Lex::Excl(at) => {
690                if member <= at {
691                    Ordering::Greater
692                } else {
693                    Ordering::Less
694                }
695            }
696        });
697        let end = self.seek(|_, member| match max {
698            Lex::Min => Ordering::Less,
699            Lex::Max => Ordering::Greater,
700            Lex::Incl(at) => {
701                if member <= at {
702                    Ordering::Greater
703                } else {
704                    Ordering::Less
705                }
706            }
707            Lex::Excl(at) => {
708                if member < at {
709                    Ordering::Greater
710                } else {
711                    Ordering::Less
712                }
713            }
714        });
715        start..end.max(start)
716    }
717
718    /// How many elements a probe leaves behind it.
719    ///
720    /// The probe is given a score and a member and says `Greater` while the
721    /// thing being looked for is still ahead. On the table this is one descent
722    /// and on a listpack it is a walk, which is the same thing at 128 members.
723    fn seek<F: FnMut(f64, &[u8]) -> Ordering>(&self, mut probe: F) -> usize {
724        match &self.body {
725            Body::Packed(lp) => {
726                let mut digits = [0u8; DIGITS_MAX];
727                let mut at = 0;
728                while let (Some(m), Some(s)) = (lp.get(at * 2), lp.get(at * 2 + 1)) {
729                    let bytes = bytes_of(m, &mut digits);
730                    if probe(score_of(s), bytes) != Ordering::Greater {
731                        break;
732                    }
733                    at += 1;
734                }
735                at
736            }
737            Body::Table(t) => t.order.seek(|row| {
738                let (name, score) = t.members.at(row as usize).expect("a row the tree holds");
739                probe(*score, name)
740            }),
741        }
742    }
743
744    /// Walk members for `ZSCAN`, in whatever order the storage has them in.
745    pub fn scan<F: FnMut(Member<'_>, f64)>(
746        &self,
747        cursor: Cursor,
748        count: usize,
749        mut f: F,
750    ) -> Cursor {
751        match &self.body {
752            Body::Packed(lp) => {
753                for at in 0..lp.len() / 2 {
754                    let (Some(m), Some(s)) = (lp.get(at * 2), lp.get(at * 2 + 1)) else {
755                        break;
756                    };
757                    f(m, score_of(s));
758                }
759                Cursor::END
760            }
761            Body::Table(t) => t.members.scan(cursor, count, |name, score| {
762                f(Member::Str(name), *score);
763            }),
764        }
765    }
766
767    /// Build a sorted set out of a member to score table that is in no order.
768    ///
769    /// This is what the algebra next door hands back. `ZUNIONSTORE` works out
770    /// every member's final score in a table that knows nothing about order,
771    /// because a member appearing in a fourth input should be a hash probe and
772    /// not a pair of tree descents, and then this puts the whole thing in order
773    /// once at the end.
774    ///
775    /// The table is not read and copied, it is moved in and becomes the sorted
776    /// set. Every member's bytes were written when the first input holding that
777    /// member was walked and they are never touched again, which is the thing
778    /// that makes a union of four large sets one pass over each of them and one
779    /// sort, rather than a pass and a rebuild.
780    ///
781    /// Answers nothing for an empty table, because an empty sorted set does not
782    /// exist and the caller's key should be deleted rather than made.
783    #[must_use]
784    pub fn from_elements(members: Elements<f64>, limits: &Limits) -> Option<Zset> {
785        let n = members.len();
786        if n == 0 {
787            return None;
788        }
789        // Row numbers put in order, which is a sort of four byte integers with
790        // an indirection in the comparison rather than a sort of the members
791        // themselves. Nothing is moved but the numbers.
792        let mut rows: Vec<u32> = (0..n as u32).collect();
793        rows.sort_unstable_by(|&a, &b| {
794            let (a_name, a_score) = members.at(a as usize).expect("in range");
795            let (b_name, b_score) = members.at(b as usize).expect("in range");
796            cmp_key(*a_score, a_name, *b_score, b_name)
797        });
798        // Small enough to stay packed, which is worth checking because a
799        // `ZINTERSTORE` of two large sets very often produces a small one.
800        let packable = n <= limits.max_listpack_entries
801            && rows.iter().all(|&r| {
802                members.at(r as usize).expect("in range").0.len() <= limits.max_listpack_value
803            });
804        if packable {
805            let mut lp = Listpack::new();
806            let mut score_buf = [0u8; DOUBLE_MAX];
807            for &row in &rows {
808                let (name, score) = members.at(row as usize).expect("in range");
809                lp.push(name);
810                lp.push(write_double(&mut score_buf, *score));
811            }
812            return Some(Zset {
813                body: Body::Packed(lp),
814            });
815        }
816        let mut order = Rank::new();
817        // Already in order, so every row goes on the end and the tree never
818        // compares anything.
819        for (at, &row) in rows.iter().enumerate() {
820            order.insert_at(at, row);
821        }
822        Some(Zset {
823            body: Body::Table(Table { members, order }),
824        })
825    }
826
827    /// Move to the table, which is one way and does not come back.
828    fn promote(&mut self) {
829        let Body::Packed(lp) = &self.body else { return };
830        let n = lp.len() / 2;
831        let mut table = Table {
832            members: Elements::with_capacity(n.next_power_of_two().max(16)),
833            order: Rank::new(),
834        };
835        // The listpack is already in order, so every member goes on the end of
836        // the tree and no comparison is needed. That is what makes a promotion a
837        // walk rather than 128 descents.
838        //
839        // Walked and not indexed. There is no offset table in a listpack, so
840        // asking it for element `i` costs a walk from the front and asking it
841        // for every element in turn costs the square of the count.
842        let mut digits = [0u8; DIGITS_MAX];
843        let mut steps = lp.iter();
844        while let (Some(m), Some(s)) = (steps.next(), steps.next()) {
845            let bytes = bytes_of(m, &mut digits);
846            let row = table.members.len() as u32;
847            if table.members.insert(bytes, score_of(s)).is_err() {
848                break;
849            }
850            table.order.insert_at(row as usize, row);
851        }
852        self.body = Body::Table(table);
853    }
854}
855
856/// Put a member and its score into a listpack at the position it belongs.
857fn packed_insert(lp: &mut Listpack, member: &[u8], score: f64) {
858    // On the stack, not in a `Vec`. A double is at most `DOUBLE_MAX` bytes and
859    // this runs on every packed `ZADD` and `ZINCRBY`, so a fresh allocation here
860    // is a malloc and a free per member added to a small sorted set.
861    let mut score_buf = [0u8; DOUBLE_MAX];
862    let text = write_double(&mut score_buf, score);
863    let mut digits = [0u8; DIGITS_MAX];
864    let mut at = 0;
865    while let (Some(m), Some(s)) = (lp.get(at * 2), lp.get(at * 2 + 1)) {
866        let bytes = bytes_of(m, &mut digits);
867        if cmp_key(score, member, score_of(s), bytes) == Ordering::Less {
868            break;
869        }
870        at += 1;
871    }
872    if at * 2 == lp.len() {
873        lp.push(member);
874        lp.push(text);
875    } else {
876        lp.insert(at * 2, member);
877        lp.insert(at * 2 + 1, text);
878    }
879}
880
881impl Table {
882    /// Where an element sits, given everything already known about it.
883    ///
884    /// The key is unique, so the lower bound is the element itself and there is
885    /// no run to walk past.
886    fn rank_of(&self, row: u32, score: f64, member: &[u8]) -> usize {
887        let members = &self.members;
888        self.order.seek(|other| {
889            if other == row {
890                return Ordering::Equal;
891            }
892            let (name, at) = members.at(other as usize).expect("a row the tree holds");
893            cmp_key(score, member, *at, name)
894        })
895    }
896
897    fn add(&mut self, member: &[u8], score: f64) -> Added {
898        if let Some(row) = self.members.index_of(member) {
899            let old = *self.members.at(row).map_or(&0.0, |(_, s)| s);
900            if cmp_score(old, score) == Ordering::Equal {
901                return Added::Same;
902            }
903            // The member has not changed and the score has, so this is a move
904            // rather than an add: out of the tree at the old rank, back in at
905            // the new one, and the element table's row number is untouched.
906            let was = self.rank_of(row as u32, old, member);
907            self.order.remove_at(was);
908            if let Some(at) = self.members.at_mut(row) {
909                *at = score;
910            }
911            let now = self.rank_of(row as u32, score, member);
912            self.order.insert_at(now, row as u32);
913            return Added::Changed;
914        }
915        let row = self.members.len() as u32;
916        if self.members.insert(member, score).is_err() {
917            return Added::Full;
918        }
919        let at = self.rank_of(row, score, member);
920        self.order.insert_at(at, row);
921        Added::New
922    }
923
924    /// Take out the element at a rank, which is known to be in a row.
925    fn take(&mut self, rank: usize, row: usize) {
926        let last = self.members.len() - 1;
927        // Where the row that is about to be renumbered sits, found before
928        // anything moves, because afterwards its score is in a different row and
929        // the tree still holds the old number.
930        let moved = if last == row {
931            None
932        } else {
933            let (name, score) = self.members.at(last).expect("the last row");
934            // The name is borrowed from the element table and the search reads
935            // the same table, so both borrows are shared and neither outlives
936            // the call.
937            let at = {
938                let members = &self.members;
939                let score = *score;
940                self.order.seek(|other| {
941                    if other as usize == last {
942                        return Ordering::Equal;
943                    }
944                    let (other_name, other_score) =
945                        members.at(other as usize).expect("a row the tree holds");
946                    cmp_key(score, name, *other_score, other_name)
947                })
948            };
949            Some(at)
950        };
951        self.order.remove_at(rank);
952        self.members.remove_at(row);
953        if let Some(at) = moved {
954            // Everything above the hole shifted down by one when the row came
955            // out of the tree.
956            let at = if at > rank { at - 1 } else { at };
957            self.order.set_at(at, row as u32);
958        }
959    }
960}
961
962#[cfg(test)]
963mod tests {
964    use super::*;
965
966    /// A listpack holding the member and score pairs given, in the order given.
967    fn packed(pairs: &[(&[u8], &str)]) -> Listpack {
968        let mut lp = Listpack::new();
969        for (member, score) in pairs {
970            lp.push(member);
971            lp.push(score.as_bytes());
972        }
973        lp
974    }
975
976    /// A payload already in this band's layout is taken whole.
977    ///
978    /// This is the whole point of [`Zset::from_packed`]: a hundred member sorted
979    /// set restored in 534 us by adding a member at a time and restores in
980    /// 4.6 us by moving the blob in, because adding costs a scan for the member
981    /// and a scan for the position on every one of them.
982    #[test]
983    fn a_payload_in_this_layout_is_taken_whole() {
984        let lp = packed(&[(b"a", "1"), (b"b", "2"), (b"c", "2.5")]);
985        let z = Zset::from_packed(lp, &Limits::DEFAULT).expect("in order and inside the limits");
986        assert_eq!(z.encoding(), Encoding::Listpack);
987        assert_eq!(z.len(), 3);
988        assert_eq!(z.score(b"a"), Some(1.0));
989        assert_eq!(z.score(b"c"), Some(2.5));
990        assert_eq!(z.rank(b"b"), Some(1));
991        assert_eq!(z.score(b"missing"), None);
992
993        // A member that looks like a number goes into a listpack as a number and
994        // not as its digits, so it is worth pinning that one is still found. A
995        // blob from another server is full of these and the failure would be a
996        // member that is in the set and cannot be looked up.
997        let lp = packed(&[(b"10", "1"), (b"9", "2")]);
998        let z = Zset::from_packed(lp, &Limits::DEFAULT).expect("sorted by score");
999        assert_eq!(z.score(b"10"), Some(1.0));
1000        assert_eq!(z.score(b"9"), Some(2.0));
1001        assert_eq!(z.rank(b"9"), Some(1));
1002    }
1003
1004    /// A blob this band cannot hold comes back rather than being taken on trust.
1005    ///
1006    /// The caller walks it after that, so none of these is an error, and getting
1007    /// any of them wrong would be: the band answers a rank query by position and
1008    /// nothing else, so a payload claiming to be sorted while not being sorted
1009    /// would answer `ZRANGE` with the wrong members and never say why.
1010    #[test]
1011    fn a_blob_this_band_cannot_hold_is_handed_back() {
1012        let small = Limits {
1013            max_listpack_entries: 4,
1014            max_listpack_value: 8,
1015        };
1016        for (why, lp) in [
1017            ("out of order by score", packed(&[(b"a", "2"), (b"b", "1")])),
1018            (
1019                "out of order by member",
1020                packed(&[(b"b", "1"), (b"a", "1")]),
1021            ),
1022            (
1023                "the same member twice",
1024                packed(&[(b"a", "1"), (b"a", "1"), (b"b", "2")]),
1025            ),
1026            ("a score that is not a number", packed(&[(b"a", "no")])),
1027            (
1028                "a member past the value limit",
1029                packed(&[(b"aaaaaaaaaa", "1")]),
1030            ),
1031            (
1032                "more members than the band takes",
1033                packed(&[
1034                    (b"a", "1"),
1035                    (b"b", "2"),
1036                    (b"c", "3"),
1037                    (b"d", "4"),
1038                    (b"e", "5"),
1039                ]),
1040            ),
1041            ("nothing in it at all", packed(&[])),
1042        ] {
1043            assert!(
1044                Zset::from_packed(lp, &small).is_err(),
1045                "{why} should be handed back"
1046            );
1047        }
1048
1049        // An odd count has no last score, which is malformed rather than merely
1050        // outside the band.
1051        let mut odd = Listpack::new();
1052        odd.push(b"a");
1053        assert!(Zset::from_packed(odd, &Limits::DEFAULT).is_err());
1054    }
1055
1056    /// A hint past the band starts on the table instead of filling the band up.
1057    #[test]
1058    fn a_hint_past_the_band_starts_on_the_table() {
1059        let big = Zset::with_hint(Limits::DEFAULT.max_listpack_entries + 1, &Limits::DEFAULT);
1060        assert_eq!(big.encoding(), Encoding::Skiplist);
1061        assert!(big.is_empty());
1062
1063        // At the limit it is still packed, matching what the add path does.
1064        let small = Zset::with_hint(Limits::DEFAULT.max_listpack_entries, &Limits::DEFAULT);
1065        assert_eq!(small.encoding(), Encoding::Listpack);
1066
1067        // And a hint is only a hint, so the table takes members like anything
1068        // else and reports them in order.
1069        let mut z = Zset::with_hint(1_000_000, &Limits::DEFAULT);
1070        z.add(b"b", 2.0, &Limits::DEFAULT);
1071        z.add(b"a", 1.0, &Limits::DEFAULT);
1072        assert_eq!(z.len(), 2);
1073        assert_eq!(z.rank(b"a"), Some(0));
1074        assert_eq!(z.rank(b"b"), Some(1));
1075    }
1076
1077    /// What a sorted set actually costs per member, which is M4's exit gate and
1078    /// was an argument rather than a number until this was written.
1079    ///
1080    /// Run it with `cargo test -p yo-kv --release measure_bytes_per_entry --
1081    /// --ignored --nocapture`. It is ignored because a million members is not
1082    /// something every `cargo test` should pay for, and it prints rather than
1083    /// asserts because the number it prints is the thing being reported. The
1084    /// bound that guards against a regression is
1085    /// [`a_large_sorted_set_does_not_hold_much_more_than_it_stores`], which is
1086    /// small enough to run every time.
1087    #[test]
1088    #[ignore = "a measurement, run it by name"]
1089    fn measure_bytes_per_entry() {
1090        // The packed band first, at the largest size it is allowed to reach.
1091        let mut lp = Zset::new();
1092        let mut lp_payload = 0usize;
1093        for i in 0..128 {
1094            let m = format!("member:{i:09}");
1095            lp_payload += m.len();
1096            lp.add(m.as_bytes(), i as f64, &Limits::DEFAULT);
1097        }
1098        println!(
1099            "packed n=128 total={} payload={lp_payload} overhead_per_entry={:.2}",
1100            lp.memory_bytes(),
1101            (lp.memory_bytes() as f64 - lp_payload as f64) / 128.0
1102        );
1103        // And the table band, at four sizes, because the answer used to depend
1104        // on how close the count was to a power of two and that is exactly the
1105        // thing being fixed.
1106        for n in [10_000usize, 100_000, 600_000, 1_000_000] {
1107            let (z, payload) = filled(n);
1108            let total = z.memory_bytes();
1109            let scores = n * 8;
1110            let (slots, rows, names, tree) = match &z.body {
1111                Body::Table(t) => (
1112                    t.members.slot_bytes(),
1113                    t.members.row_bytes(),
1114                    t.members.name_bytes(),
1115                    t.order.bytes(),
1116                ),
1117                Body::Packed(_) => (0, 0, 0, 0),
1118            };
1119            let per = |b: usize| b as f64 / n as f64;
1120            println!(
1121                "table n={n} total={total} slots={:.2}/e rows={:.2}/e names={:.2}/e tree={:.2}/e overhead_per_entry={:.2}",
1122                per(slots),
1123                per(rows),
1124                per(names),
1125                per(tree),
1126                (total as f64 - payload as f64 - scores as f64) / n as f64
1127            );
1128        }
1129    }
1130
1131    /// `n` members named `member:` and nine digits, so sixteen bytes each, and
1132    /// what those names weigh.
1133    fn filled(n: usize) -> (Zset, usize) {
1134        let mut z = Zset::new();
1135        let mut payload = 0usize;
1136        for i in 0..n {
1137            let m = format!("member:{i:09}");
1138            payload += m.len();
1139            z.add(m.as_bytes(), i as f64, &Limits::DEFAULT);
1140        }
1141        (z, payload)
1142    }
1143
1144    /// The guard on the measurement above.
1145    ///
1146    /// Forty thousand members is a count nowhere near a power of two, which is
1147    /// the case that used to be worst: a row array and a name blob that had both
1148    /// just doubled held nearly twice what they were storing, and the slack was
1149    /// more than everything else in the structure put together. Both grow by a
1150    /// quarter now, so the bound below is one a doubling array cannot meet and
1151    /// this test fails if either of them goes back to `Vec`'s policy.
1152    #[test]
1153    fn a_large_sorted_set_does_not_hold_much_more_than_it_stores() {
1154        let n = 40_000usize;
1155        let (z, payload) = filled(n);
1156        let Body::Table(t) = &z.body else {
1157            panic!("forty thousand members is not a listpack");
1158        };
1159        // Twelve bytes of row and eight of score, so twenty five leaves room for
1160        // the growth policy's quarter and none for a payload that has crept back
1161        // inside the row and brought four bytes of padding with it.
1162        assert!(
1163            t.members.row_bytes() < n * 25,
1164            "the row and score arrays hold {} for {n} members",
1165            t.members.row_bytes()
1166        );
1167        assert!(
1168            t.members.name_bytes() < payload + payload / 4,
1169            "the name blob holds {} for {payload} bytes of names",
1170            t.members.name_bytes()
1171        );
1172        // The tree is the part that already meets the gate and the part most
1173        // likely to be quietly broken by a change to the element table, so it
1174        // is worth pinning: three and a bit bytes a member, which is the row
1175        // number plus its share of a branch node at a fanout of a hundred and
1176        // twenty eight.
1177        assert!(
1178            t.order.bytes() < n * 4,
1179            "the tree holds {} for {n} members",
1180            t.order.bytes()
1181        );
1182    }
1183
1184    /// A set built by adding in whatever order, checked against a model.
1185    fn built(pairs: &[(&str, f64)], limits: &Limits) -> Zset {
1186        let mut z = Zset::new();
1187        for (m, s) in pairs {
1188            z.add(m.as_bytes(), *s, limits);
1189        }
1190        z
1191    }
1192
1193    /// Every member and score in rank order.
1194    fn listed(z: &Zset) -> Vec<(String, f64)> {
1195        let mut out = Vec::new();
1196        let mut digits = [0u8; DIGITS_MAX];
1197        z.walk(0, z.len(), false, |m, s| {
1198            let bytes = bytes_of(m, &mut digits).to_vec();
1199            out.push((String::from_utf8(bytes).unwrap(), s));
1200        });
1201        out
1202    }
1203
1204    /// What the model says the order is.
1205    fn model(pairs: &[(&str, f64)]) -> Vec<(String, f64)> {
1206        let mut last: Vec<(String, f64)> = Vec::new();
1207        for (m, s) in pairs {
1208            match last.iter_mut().find(|(name, _)| name == m) {
1209                Some(row) => row.1 = *s,
1210                None => last.push(((*m).to_string(), *s)),
1211            }
1212        }
1213        last.sort_by(|a, b| cmp_key(a.1, a.0.as_bytes(), b.1, b.0.as_bytes()));
1214        last
1215    }
1216
1217    const PACKED: Limits = Limits::DEFAULT;
1218    const TABLE: Limits = Limits {
1219        max_listpack_entries: 0,
1220        max_listpack_value: 64,
1221    };
1222
1223    #[test]
1224    fn an_empty_set_answers_nothing() {
1225        let z = Zset::new();
1226        assert_eq!(z.len(), 0);
1227        assert!(z.is_empty());
1228        assert_eq!(z.encoding(), Encoding::Listpack);
1229        assert_eq!(z.score(b"nobody"), None);
1230        assert_eq!(z.rank(b"nobody"), None);
1231        assert!(z.at(0).is_none());
1232        assert_eq!(
1233            z.window_by_score(Bound::closed(f64::NEG_INFINITY), Bound::closed(0.0)),
1234            0..0
1235        );
1236    }
1237
1238    #[test]
1239    fn both_bands_put_members_in_the_same_order() {
1240        let pairs = [("c", 3.0), ("a", 1.0), ("b", 2.0), ("d", 2.0), ("e", -1.5)];
1241        for limits in [&PACKED, &TABLE] {
1242            let z = built(&pairs, limits);
1243            assert_eq!(listed(&z), model(&pairs), "{:?}", z.encoding());
1244            assert_eq!(z.len(), 5);
1245        }
1246    }
1247
1248    #[test]
1249    fn a_tie_on_score_is_broken_by_the_member() {
1250        for limits in [&PACKED, &TABLE] {
1251            let pairs = [
1252                ("beta", 1.0),
1253                ("alpha", 1.0),
1254                ("gamma", 1.0),
1255                ("Alpha", 1.0),
1256            ];
1257            let z = built(&pairs, limits);
1258            let names: Vec<String> = listed(&z).into_iter().map(|(m, _)| m).collect();
1259            assert_eq!(names, ["Alpha", "alpha", "beta", "gamma"]);
1260        }
1261    }
1262
1263    #[test]
1264    fn adding_a_member_again_moves_it_rather_than_adding_it() {
1265        for limits in [&PACKED, &TABLE] {
1266            let mut z = built(&[("a", 1.0), ("b", 2.0), ("c", 3.0)], limits);
1267            assert_eq!(z.add(b"a", 1.0, limits), Added::Same);
1268            assert_eq!(z.add(b"a", 9.0, limits), Added::Changed);
1269            assert_eq!(z.len(), 3);
1270            assert_eq!(z.score(b"a"), Some(9.0));
1271            assert_eq!(z.rank(b"a"), Some(2));
1272            assert_eq!(listed(&z).last().unwrap().0, "a");
1273        }
1274    }
1275
1276    #[test]
1277    fn a_removal_leaves_every_other_rank_right() {
1278        for limits in [&PACKED, &TABLE] {
1279            let pairs: Vec<(&str, f64)> =
1280                vec![("a", 1.0), ("b", 2.0), ("c", 3.0), ("d", 4.0), ("e", 5.0)];
1281            let mut z = built(&pairs, limits);
1282            assert!(z.remove(b"c"));
1283            assert!(!z.remove(b"c"));
1284            assert_eq!(z.len(), 4);
1285            assert_eq!(z.rank(b"a"), Some(0));
1286            assert_eq!(z.rank(b"d"), Some(2));
1287            assert_eq!(z.rank(b"e"), Some(3));
1288            assert_eq!(z.score(b"c"), None);
1289        }
1290    }
1291
1292    /// The dense element table moves the last row into the hole, so a removal
1293    /// renumbers an element nobody asked about. This is the test that fails if
1294    /// the tree is not told.
1295    #[test]
1296    fn removing_from_the_middle_renumbers_the_last_element() {
1297        let limits = &TABLE;
1298        let mut z = Zset::new();
1299        for i in 0..64u32 {
1300            z.add(format!("m{i:03}").as_bytes(), f64::from(i), limits);
1301        }
1302        // Take them out from the front, which moves the last row into row zero
1303        // every single time.
1304        for i in 0..63u32 {
1305            assert!(z.remove(format!("m{i:03}").as_bytes()));
1306            assert_eq!(z.len() as u32, 63 - i);
1307            // Every member still in here has to be findable, in the right place,
1308            // with the right score.
1309            for j in i + 1..64 {
1310                let name = format!("m{j:03}");
1311                assert_eq!(
1312                    z.score(name.as_bytes()),
1313                    Some(f64::from(j)),
1314                    "score of {name}"
1315                );
1316                assert_eq!(
1317                    z.rank(name.as_bytes()),
1318                    Some((j - i - 1) as usize),
1319                    "rank of {name} after {i}"
1320                );
1321            }
1322        }
1323    }
1324
1325    #[test]
1326    fn a_set_promotes_when_it_outgrows_the_packed_band() {
1327        let limits = Limits {
1328            max_listpack_entries: 4,
1329            max_listpack_value: 64,
1330        };
1331        let mut z = Zset::new();
1332        for i in 0..4u32 {
1333            z.add(format!("m{i}").as_bytes(), f64::from(i), &limits);
1334        }
1335        assert_eq!(z.encoding(), Encoding::Listpack);
1336        z.add(b"m4", 4.0, &limits);
1337        assert_eq!(z.encoding(), Encoding::Skiplist);
1338        assert_eq!(z.len(), 5);
1339        let names: Vec<String> = listed(&z).into_iter().map(|(m, _)| m).collect();
1340        assert_eq!(names, ["m0", "m1", "m2", "m3", "m4"]);
1341        // A promotion is one way, and a set that shrinks stays where it is.
1342        z.remove(b"m4");
1343        z.remove(b"m3");
1344        assert_eq!(z.encoding(), Encoding::Skiplist);
1345    }
1346
1347    #[test]
1348    fn a_member_too_long_for_the_packed_band_promotes_on_its_own() {
1349        let limits = Limits {
1350            max_listpack_entries: 128,
1351            max_listpack_value: 8,
1352        };
1353        let mut z = Zset::new();
1354        z.add(b"short", 1.0, &limits);
1355        assert_eq!(z.encoding(), Encoding::Listpack);
1356        z.add(b"a member well past eight bytes", 2.0, &limits);
1357        assert_eq!(z.encoding(), Encoding::Skiplist);
1358        assert_eq!(z.len(), 2);
1359        assert_eq!(z.score(b"a member well past eight bytes"), Some(2.0));
1360    }
1361
1362    #[test]
1363    fn a_score_range_finds_both_of_its_ends() {
1364        for limits in [&PACKED, &TABLE] {
1365            let pairs = [("a", 1.0), ("b", 2.0), ("c", 2.0), ("d", 3.0), ("e", 4.0)];
1366            let z = built(&pairs, limits);
1367            assert_eq!(
1368                z.window_by_score(Bound::closed(2.0), Bound::closed(3.0)),
1369                1..4
1370            );
1371            assert_eq!(
1372                z.window_by_score(Bound::open(2.0), Bound::closed(3.0)),
1373                3..4
1374            );
1375            assert_eq!(
1376                z.window_by_score(Bound::closed(2.0), Bound::open(3.0)),
1377                1..3
1378            );
1379            assert_eq!(z.window_by_score(Bound::open(1.0), Bound::open(4.0)), 1..4);
1380            assert_eq!(
1381                z.window_by_score(
1382                    Bound::closed(f64::NEG_INFINITY),
1383                    Bound::closed(f64::INFINITY)
1384                ),
1385                0..5
1386            );
1387            // A range with nothing in it is empty and not backwards.
1388            assert_eq!(
1389                z.window_by_score(Bound::closed(3.0), Bound::closed(2.0)),
1390                3..3
1391            );
1392            assert_eq!(
1393                z.window_by_score(Bound::closed(9.0), Bound::closed(10.0)),
1394                5..5
1395            );
1396        }
1397    }
1398
1399    #[test]
1400    fn a_member_range_orders_by_member_when_the_scores_are_equal() {
1401        for limits in [&PACKED, &TABLE] {
1402            let pairs = [("a", 0.0), ("b", 0.0), ("c", 0.0), ("d", 0.0), ("e", 0.0)];
1403            let z = built(&pairs, limits);
1404            assert_eq!(z.window_by_lex(Lex::Min, Lex::Max), 0..5);
1405            assert_eq!(z.window_by_lex(Lex::Incl(b"b"), Lex::Incl(b"d")), 1..4);
1406            assert_eq!(z.window_by_lex(Lex::Excl(b"b"), Lex::Excl(b"d")), 2..3);
1407            assert_eq!(z.window_by_lex(Lex::Incl(b"b"), Lex::Excl(b"c")), 1..2);
1408            assert_eq!(z.window_by_lex(Lex::Excl(b"e"), Lex::Max), 5..5);
1409            assert_eq!(z.window_by_lex(Lex::Min, Lex::Excl(b"a")), 0..0);
1410        }
1411    }
1412
1413    #[test]
1414    fn a_walk_can_go_backwards_and_stops_where_it_is_told() {
1415        for limits in [&PACKED, &TABLE] {
1416            let pairs = [("a", 1.0), ("b", 2.0), ("c", 3.0), ("d", 4.0)];
1417            let z = built(&pairs, limits);
1418            let mut seen = Vec::new();
1419            let mut digits = [0u8; DIGITS_MAX];
1420            z.walk(3, 2, true, |m, _| {
1421                seen.push(String::from_utf8(bytes_of(m, &mut digits).to_vec()).unwrap());
1422            });
1423            assert_eq!(seen, ["d", "c"]);
1424            let mut seen = Vec::new();
1425            z.walk(1, 99, false, |m, _| {
1426                seen.push(String::from_utf8(bytes_of(m, &mut digits).to_vec()).unwrap());
1427            });
1428            assert_eq!(seen, ["b", "c", "d"]);
1429            // A rank past the end is nothing rather than a panic.
1430            let mut count = 0;
1431            z.walk(4, 1, false, |_, _| count += 1);
1432            assert_eq!(count, 0);
1433        }
1434    }
1435
1436    #[test]
1437    fn taking_from_a_rank_takes_the_right_one() {
1438        for limits in [&PACKED, &TABLE] {
1439            let pairs = [("a", 1.0), ("b", 2.0), ("c", 3.0)];
1440            let mut z = built(&pairs, limits);
1441            assert!(z.remove_at(0));
1442            assert_eq!(z.len(), 2);
1443            assert_eq!(z.rank(b"b"), Some(0));
1444            assert!(z.remove_at(1));
1445            assert_eq!(z.score(b"c"), None);
1446            assert!(!z.remove_at(5));
1447        }
1448    }
1449
1450    /// The one that actually catches things: a few thousand adds, moves and
1451    /// removals against a `Vec` that says what the answer is.
1452    #[test]
1453    fn a_run_of_everything_agrees_with_a_model() {
1454        for limits in [
1455            &PACKED,
1456            &Limits {
1457                max_listpack_entries: 8,
1458                max_listpack_value: 64,
1459            },
1460        ] {
1461            let mut z = Zset::new();
1462            let mut model: Vec<(String, f64)> = Vec::new();
1463            let mut seed = 0x8765_4321_9ABC_DEF0u64;
1464            let mut roll = || {
1465                seed ^= seed << 13;
1466                seed ^= seed >> 7;
1467                seed ^= seed << 17;
1468                seed
1469            };
1470            for round in 0..3_000 {
1471                let name = format!("m{:02}", roll() % 40);
1472                // Scores land on a handful of values so that ties are the rule
1473                // and not an accident.
1474                let score = (roll() % 7) as f64 - 3.0;
1475                if round % 5 == 4 {
1476                    let hit = z.remove(name.as_bytes());
1477                    let was = model.iter().position(|(m, _)| *m == name);
1478                    assert_eq!(hit, was.is_some());
1479                    if let Some(at) = was {
1480                        model.remove(at);
1481                    }
1482                } else {
1483                    z.add(name.as_bytes(), score, limits);
1484                    match model.iter_mut().find(|(m, _)| *m == name) {
1485                        Some(row) => row.1 = score,
1486                        None => model.push((name, score)),
1487                    }
1488                }
1489                if round % 97 == 0 {
1490                    let mut want = model.clone();
1491                    want.sort_by(|a, b| cmp_key(a.1, a.0.as_bytes(), b.1, b.0.as_bytes()));
1492                    assert_eq!(listed(&z), want, "round {round}");
1493                    for (at, (m, s)) in want.iter().enumerate() {
1494                        assert_eq!(z.rank(m.as_bytes()), Some(at), "rank of {m}");
1495                        assert_eq!(z.score(m.as_bytes()), Some(*s), "score of {m}");
1496                    }
1497                }
1498            }
1499        }
1500    }
1501
1502    #[test]
1503    fn a_big_set_costs_what_the_tree_said_it_would() {
1504        let limits = &TABLE;
1505        let mut z = Zset::new();
1506        let n = 100_000u32;
1507        for i in 0..n {
1508            z.add(format!("member:{i:08}").as_bytes(), f64::from(i), limits);
1509        }
1510        assert_eq!(z.len(), n as usize);
1511        assert_eq!(z.rank(b"member:00050000"), Some(50_000));
1512        assert_eq!(z.at(0).map(|(_, s)| s), Some(0.0));
1513        assert_eq!(z.at(n as usize - 1).map(|(_, s)| s), Some(f64::from(n - 1)));
1514        // What the gate is about is the ordering structure, which is the tree,
1515        // and it holds three bytes an element and change. The rest is the
1516        // element table, which a hash and a set pay the same way: twenty four
1517        // bytes a row because a score is eight and forces the row to align to
1518        // eight, four more for the open addressed slot, and the member bytes
1519        // themselves. The row array and the slot array are both powers of two,
1520        // so at a hundred thousand members a third of both is slack.
1521        let Body::Table(t) = &z.body else {
1522            panic!("a hundred thousand members is not a listpack")
1523        };
1524        let per_order = t.order.bytes() as f64 / f64::from(n);
1525        assert!(per_order < 3.4, "{per_order} bytes an element in the tree");
1526        let per = z.memory_bytes() as f64 / f64::from(n);
1527        assert!(per < 70.0, "{per} bytes a member all in");
1528    }
1529
1530    /// Freeze a sorted set, read it back, and check that nothing about it moved.
1531    fn round_trip(z: &Zset) -> Zset {
1532        let mut buf = Vec::new();
1533        z.freeze(&mut buf);
1534        let back = Zset::thaw(&buf).expect("what freeze wrote");
1535        assert_eq!(back.len(), z.len(), "the member count");
1536        assert_eq!(back.encoding(), z.encoding(), "the band");
1537        assert_eq!(listed(&back), listed(z), "the members in rank order");
1538        back
1539    }
1540
1541    #[test]
1542    fn a_frozen_sorted_set_comes_back_in_the_band_it_left() {
1543        let pairs: Vec<(String, f64)> = (0..40)
1544            .map(|i| (format!("member:{i:04}"), f64::from(i) * 1.5 - 12.0))
1545            .collect();
1546        let refs: Vec<(&str, f64)> = pairs.iter().map(|(m, s)| (m.as_str(), *s)).collect();
1547
1548        let packed = round_trip(&built(&refs, &PACKED));
1549        assert_eq!(packed.encoding(), Encoding::Listpack);
1550        let table = round_trip(&built(&refs, &TABLE));
1551        assert_eq!(table.encoding(), Encoding::Skiplist);
1552
1553        // The two bands agree with each other and with the model, so a value
1554        // that was frozen on one and read on the other would have been caught.
1555        assert_eq!(listed(&packed), model(&refs));
1556        assert_eq!(listed(&table), model(&refs));
1557
1558        round_trip(&Zset::new());
1559    }
1560
1561    #[test]
1562    fn a_score_survives_the_trip_exactly() {
1563        // Scores that a text round trip through a listpack would be at risk of
1564        // rounding, plus the two that a zigzagged integer encoding would get
1565        // wrong. Every one of them has to come back bit for bit, because a
1566        // score is what the whole order is built on.
1567        let scores = [
1568            0.0,
1569            -0.0,
1570            1.0 / 3.0,
1571            -1.0 / 3.0,
1572            f64::MIN_POSITIVE,
1573            f64::MAX,
1574            f64::MIN,
1575            9_007_199_254_740_993.0,
1576            f64::INFINITY,
1577            f64::NEG_INFINITY,
1578        ];
1579        let mut z = Zset::new();
1580        for (i, s) in scores.iter().enumerate() {
1581            z.add(format!("member:{i:04}").as_bytes(), *s, &TABLE);
1582        }
1583        let back = round_trip(&z);
1584        for (i, s) in scores.iter().enumerate() {
1585            let member = format!("member:{i:04}");
1586            assert_eq!(
1587                back.score(member.as_bytes()).map(f64::to_bits),
1588                Some(s.to_bits()),
1589                "the score of {member}"
1590            );
1591        }
1592    }
1593
1594    #[test]
1595    fn a_sorted_set_that_came_back_still_takes_members_and_ranks_them() {
1596        let pairs: Vec<(String, f64)> = (0..200)
1597            .map(|i| (format!("member:{i:04}"), f64::from(i)))
1598            .collect();
1599        let refs: Vec<(&str, f64)> = pairs.iter().map(|(m, s)| (m.as_str(), *s)).collect();
1600        let mut back = round_trip(&built(&refs, &TABLE));
1601
1602        // One in the middle, one at each end, and one that is already in here
1603        // and only moves.
1604        assert_eq!(back.add(b"middle", 99.5, &TABLE), Added::New);
1605        assert_eq!(back.add(b"first", -1.0, &TABLE), Added::New);
1606        assert_eq!(back.add(b"last", 1000.0, &TABLE), Added::New);
1607        assert_eq!(back.add(b"member:0000", 500.0, &TABLE), Added::Changed);
1608
1609        assert_eq!(back.len(), 203);
1610        assert_eq!(back.rank(b"first"), Some(0));
1611        assert_eq!(back.rank(b"last"), Some(202));
1612        // `first`, then member:0001 through member:0099, because member:0000
1613        // moved up to five hundred and left the front.
1614        assert_eq!(back.rank(b"middle"), Some(100));
1615        assert_eq!(back.score(b"member:0000"), Some(500.0));
1616        assert!(back.remove(b"middle"));
1617        assert_eq!(back.rank(b"last"), Some(201));
1618    }
1619
1620    #[test]
1621    fn a_frozen_sorted_set_that_arrives_damaged_is_an_error_and_not_a_panic() {
1622        let pairs: Vec<(String, f64)> = (0..200)
1623            .map(|i| (format!("member:{i:04}"), f64::from(i)))
1624            .collect();
1625        let refs: Vec<(&str, f64)> = pairs.iter().map(|(m, s)| (m.as_str(), *s)).collect();
1626        let z = built(&refs, &TABLE);
1627        let mut buf = Vec::new();
1628        z.freeze(&mut buf);
1629
1630        assert!(Zset::thaw(&[]).is_err(), "nothing at all");
1631        assert!(Zset::thaw(&[99]).is_err(), "a form nobody wrote");
1632        for cut in 1..buf.len().min(64) {
1633            assert!(Zset::thaw(&buf[..cut]).is_err(), "cut at {cut}");
1634        }
1635        // A count that claims far more members than there are bytes behind it.
1636        let mut lying = vec![FORM_MEMBERS];
1637        frozen::put_uint(&mut lying, u64::MAX);
1638        assert!(Zset::thaw(&lying).is_err(), "a count nobody could hold");
1639        // The same member twice, which would leave the tree longer than the
1640        // table and every rank after it wrong.
1641        let mut twice = vec![FORM_MEMBERS];
1642        frozen::put_uint(&mut twice, 2);
1643        for _ in 0..2 {
1644            frozen::put_bytes(&mut twice, b"member");
1645            frozen::put_f64(&mut twice, 1.0);
1646        }
1647        assert!(Zset::thaw(&twice).is_err(), "a member written twice");
1648    }
1649}