Skip to main content

yo_kv/
set.rs

1//! A set, in whichever representation currently fits it.
2//!
3//! A set is one of an [`Intset`], a [`Listpack`], an [`Elements`] table or a
4//! [`Parts`] band, and which one is not a choice this file gets to make freely.
5//! `OBJECT ENCODING` has to answer `intset`, `listpack` or `hashtable` at exactly
6//! the sizes a real server answers them, because clients and test suites read it
7//! (`08` ยง1), so the promotion rules here are Redis's rules and they were read
8//! off `t_set.c` in the 8.10.1 tarball rather than reasoned out from what each
9//! structure is good at.
10//!
11//! ```text
12//!  all integers        small, any bytes      everything else       large
13//! +-------------+   +------------------+   +---------------+   +-----------+
14//! | intset      |-->| listpack         |-->| element table |-->| partition |
15//! | 2 B member  |   | ~2 B + payload   |   | one probe     |   | band      |
16//! +-------------+   +------------------+   +---------------+   +-----------+
17//!  to 512 members     to 128 members                            past 262144
18//!                                          <------- both say `hashtable` ---->
19//! ```
20//!
21//! The fourth step is the one Redis does not have, and it is invisible on
22//! purpose. A set past 262,144 members becomes several element tables rather than
23//! one, which is what makes the merges and the growth pauses bounded, and it
24//! still answers `hashtable` because a client that gets a fourth word back from
25//! `OBJECT ENCODING` is a client whose assertions break. What partitioning
26//! changes is how a large set is stored, not what it is. See [`crate::parts`].
27//!
28//! Promotion is one-way and upward, which is Y4. A set that has been a hash
29//! table does not go back to an intset when it shrinks, and neither does Redis's:
30//! a set that demoted on the way down would rewrite itself on every second
31//! operation for a workload that adds and removes across a threshold. The band
32//! follows the same rule for the same reason, so a set hovering at the partition
33//! threshold rehashes once rather than on every other `SREM`.
34//!
35//! # The rules, and the two that are not obvious
36//!
37//! Adding an integer to an intset keeps it an intset until it holds more than
38//! `set-max-intset-entries`, and then it becomes a **hash table** and not a
39//! listpack, because the intset ceiling is 512 and the listpack ceiling is 128
40//! and something over the first is well over the second.
41//!
42//! Adding a non integer to an intset is the asymmetric one. It becomes a
43//! listpack only if the intset is currently under the *listpack* ceiling of 128,
44//! so an intset of 200 integers that receives one string goes straight to a hash
45//! table and is never a listpack at all. Reading that off the source was worth
46//! more than reasoning about it, because the natural implementation converts to
47//! a listpack whenever the members would fit and gets a different encoding name
48//! than a real server for a shape a test suite actually builds.
49//!
50//! # Members
51//!
52//! A member comes back as a [`Member`], which is [`listpack::Entry`] under
53//! another name: either the bytes as they lie or an integer that has not been
54//! formatted yet. All three representations can produce one without copying, and
55//! the formatting happens once, into the reply buffer, at the moment the reply is
56//! built. That is Y18, and it is the same reason [`crate::value::Str`] has two
57//! arms.
58
59use crate::elem::Elements;
60use crate::frozen::{self, Broken};
61use crate::intset::Intset;
62use crate::listpack::{self, Listpack};
63use crate::parts::{PARTITION_AT, Parts, parts_for};
64use crate::scan::Cursor;
65use yo_common::num::{DIGITS_MAX, i64_digits, i64_len, parse_i64};
66
67/// A set member: bytes as they lie, or an integer not yet formatted.
68pub type Member<'a> = listpack::Entry<'a>;
69
70/// A single run intset, written as the Redis layout it already is.
71const FORM_INTSET: u8 = 1;
72/// A listpack, written as the layout it already is.
73const FORM_PACKED: u8 = 2;
74/// An intset that has split into runs, written as ascending gaps.
75const FORM_RUNS: u8 = 3;
76/// A table or a band, written as its members.
77const FORM_MEMBERS: u8 = 4;
78/// The top bit of the form byte, which carries [`Set::ints_past_limit`].
79///
80/// In the form byte rather than a byte of its own because it only ever means
81/// anything for the two integer forms, and a set that has to be told what to call
82/// itself has already been told what it is.
83const PAST_LIMIT: u8 = 0x80;
84
85/// A member on its way to being asked about, in every form the three
86/// representations want it in.
87///
88/// Set algebra walks one set and asks every other set the same question about
89/// each member, and the three representations do not want the question in the
90/// same shape. An intset wants a number, a listpack wants bytes and a number
91/// because it holds both kinds, and an element table wants bytes and their
92/// hash. Asking through [`Set::contains`] would redo all of that per question:
93/// a parse for the intset, another parse inside the listpack, and a hash per
94/// table. This does each once per member and then asks `k - 1` times.
95///
96/// The hash is computed whether or not any operand is a table, which is waste
97/// when none is. It is waste worth taking, because the sets where it is wasted
98/// are an intset or a listpack, which are capped at a few hundred members, and
99/// an operation over sets that small is finished before the saving could have
100/// been measured. The sets where the hash pays are the large ones, and those
101/// are tables by definition.
102#[derive(Debug, Clone, Copy)]
103pub struct Needle<'a> {
104    /// The member as bytes, which for an integer member is a caller's buffer.
105    bytes: &'a [u8],
106    /// The number it is, if it is one, under the same rule that decides whether
107    /// a set stores it as one.
108    int: Option<i64>,
109    /// What an element table would key it under.
110    hash: u64,
111}
112
113impl<'a> Needle<'a> {
114    /// A needle from bytes, which is what a command line argument is.
115    #[must_use]
116    pub fn new(bytes: &'a [u8]) -> Needle<'a> {
117        Needle {
118            bytes,
119            int: parse_i64(bytes),
120            hash: Elements::<()>::hash_of(bytes),
121        }
122    }
123
124    /// A needle from a member walked out of a set.
125    ///
126    /// `digits` is where an integer member's text goes, because an intset holds
127    /// the number and the digits do not exist anywhere until somebody writes
128    /// them. It is the caller's buffer rather than a field so that the needle
129    /// stays a borrow and the buffer is written once per member rather than
130    /// allocated once per member.
131    ///
132    /// A member that came out as bytes is still parsed, because the set being
133    /// asked may be an intset and `SINTER ints strings` has to find the members
134    /// they share. A member that came out as a number is not, which is the
135    /// whole saving.
136    #[must_use]
137    pub fn of(member: Member<'a>, digits: &'a mut [u8; DIGITS_MAX]) -> Needle<'a> {
138        match member {
139            Member::Str(s) => Needle::new(s),
140            Member::Int(n) => {
141                let bytes = i64_digits(digits, n);
142                Needle {
143                    bytes,
144                    int: Some(n),
145                    hash: Elements::<()>::hash_of(bytes),
146                }
147            }
148        }
149    }
150
151    /// The member as bytes, which is what a caller collecting an answer wants.
152    #[must_use]
153    pub const fn bytes(&self) -> &'a [u8] {
154        self.bytes
155    }
156}
157
158/// Where the encodings change over.
159///
160/// These are `set-max-intset-entries`, `set-max-listpack-entries` and
161/// `set-max-listpack-value`, and they are runtime configuration in Redis, so
162/// they are a value passed in here rather than three constants. The defaults are
163/// Redis's defaults and a client that never touches `CONFIG SET` sees exactly
164/// the encodings a real server would give it.
165#[derive(Debug, Clone, Copy, PartialEq, Eq)]
166pub struct Limits {
167    /// Past this many members an all integer set stops being an intset.
168    pub max_intset_entries: usize,
169    /// At this many members a set stops being a listpack.
170    pub max_listpack_entries: usize,
171    /// A member longer than this cannot go in a listpack.
172    pub max_listpack_value: usize,
173}
174
175impl Limits {
176    /// Redis's defaults: 512, 128 and 64.
177    pub const DEFAULT: Limits = Limits {
178        max_intset_entries: 512,
179        max_listpack_entries: 128,
180        max_listpack_value: 64,
181    };
182}
183
184impl Default for Limits {
185    fn default() -> Limits {
186        Limits::DEFAULT
187    }
188}
189
190/// Which of the three a set is in, which is what `OBJECT ENCODING` reports.
191#[derive(Debug, Clone, Copy, PartialEq, Eq)]
192pub enum Encoding {
193    /// All members are integers and there are few enough of them.
194    Intset,
195    /// One packed blob, walked linearly.
196    Listpack,
197    /// The element table.
198    Hashtable,
199}
200
201impl Encoding {
202    /// The word `OBJECT ENCODING` replies with.
203    #[inline]
204    pub const fn name(self) -> &'static str {
205        match self {
206            Encoding::Intset => "intset",
207            Encoding::Listpack => "listpack",
208            Encoding::Hashtable => "hashtable",
209        }
210    }
211}
212
213/// The four representations, of which `OBJECT ENCODING` can see three.
214///
215/// [`Body::Split`] is the partitioned band and it is deliberately invisible from
216/// outside. Redis has three set encodings and a client that gets a fourth word
217/// back from `OBJECT ENCODING` is a client whose assertions break, so a split set
218/// answers `hashtable` like the table it was. What partitioning changes is how a
219/// large set is stored and merged, not what it is.
220#[derive(Debug, Clone)]
221enum Body {
222    Ints(Intset),
223    Packed(Listpack),
224    Table(Elements<()>),
225    Split(Parts<()>),
226}
227
228/// A set of members.
229#[derive(Debug, Clone)]
230pub struct Set {
231    body: Body,
232    /// Whether an all integer set has passed `set-max-intset-entries`.
233    ///
234    /// This is where Redis rehashes the members into a dictionary and starts
235    /// answering `hashtable` to `OBJECT ENCODING`. Nothing is rehashed here,
236    /// because [`Intset`] holds a large set in runs and stays at two to eight
237    /// bytes a member where a table would cost thirty, so the only thing the
238    /// ceiling still decides is the word.
239    ///
240    /// It is a flag and not a comparison against the length because the ceiling
241    /// is configurable and [`Set::encoding`] is not handed the configuration. It
242    /// is also one way, like every promotion here: a set that has been called a
243    /// hashtable once does not go back to being called an intset when members
244    /// are removed, which is Y4 and is what Redis does.
245    ///
246    /// Only ever true while [`Body::Ints`] is the body. Every other body already
247    /// knows what to answer.
248    ints_past_limit: bool,
249}
250
251impl Set {
252    /// An empty set, which starts as an intset.
253    ///
254    /// This is what `SADD` on a missing key creates when it has no size hint,
255    /// and the first member decides nothing: an intset that receives a string
256    /// converts on the spot, and it costs a conversion of nothing.
257    #[must_use]
258    pub fn new() -> Set {
259        Set {
260            body: Body::Ints(Intset::new()),
261            ints_past_limit: false,
262        }
263    }
264
265    /// An empty set sized for what is about to go in it.
266    ///
267    /// Redis's `setTypeCreate`, which picks the representation from the first
268    /// member and the count the caller expects, so that `SADD k a b c ...` with
269    /// a thousand arguments builds a table once rather than converting twice on
270    /// the way there. `hint` is only a hint and being wrong about it costs a
271    /// conversion and no correctness.
272    ///
273    /// An integer first member sends it to the intset even when the count is
274    /// past `set-max-intset-entries`, where Redis would go straight to a
275    /// dictionary. A large set of integers is the case the runs exist for, and
276    /// building a table and never leaving it would give up the whole saving on
277    /// the one call that said in advance it was going to matter. The listpack
278    /// band in between is still honoured, because a set that small has nothing
279    /// to save and a server configured that way expects a listpack.
280    #[must_use]
281    pub fn with_hint(first: &[u8], hint: usize, limits: &Limits) -> Set {
282        let ints = parse_i64(first).is_some()
283            && (hint <= limits.max_intset_entries || hint > limits.max_listpack_entries);
284        if ints {
285            Set {
286                body: Body::Ints(Intset::with_capacity(hint)),
287                ints_past_limit: hint > limits.max_intset_entries,
288            }
289        } else if hint <= limits.max_listpack_entries {
290            Set {
291                body: Body::Packed(Listpack::new()),
292                ints_past_limit: false,
293            }
294        } else if hint > PARTITION_AT {
295            // A caller that says up front it is about to load a million members
296            // should not build one table, fill it past the threshold and then
297            // rehash the lot into partitions. The hint is only a hint, and being
298            // wrong about it here costs a set with more partitions than it needs
299            // rather than anything incorrect.
300            Set {
301                body: Body::Split(Parts::with_parts(parts_for(hint))),
302                ints_past_limit: false,
303            }
304        } else {
305            Set {
306                body: Body::Table(Elements::with_capacity(hint)),
307                ints_past_limit: false,
308            }
309        }
310    }
311
312    /// Which representation this is in.
313    ///
314    /// Four bodies and three words, and the intset accounts for two of the
315    /// missing ones. The partitioned body answers `hashtable` because it is one,
316    /// and an intset past `set-max-intset-entries` answers `hashtable` because
317    /// that is what a real server would have turned into by then, even though
318    /// nothing here was rehashed.
319    #[inline]
320    #[must_use]
321    pub const fn encoding(&self) -> Encoding {
322        match self.body {
323            Body::Ints(_) if self.ints_past_limit => Encoding::Hashtable,
324            Body::Ints(_) => Encoding::Intset,
325            Body::Packed(_) => Encoding::Listpack,
326            Body::Table(_) | Body::Split(_) => Encoding::Hashtable,
327        }
328    }
329
330    /// The bytes behind a set that is stored the way Redis stores it.
331    ///
332    /// `DUMP` writes these straight out instead of walking the members, so this
333    /// exists for [`crate::rdb`] and for nothing else. The word this answers
334    /// against is [`Set::encoding`] and not the body, so a set that calls itself
335    /// a hashtable is walked even on the rare occasion its members are still in
336    /// one intset run. Keeping those two in step is what makes the rule sayable:
337    /// whatever `OBJECT ENCODING` says, that is the type byte the payload gets.
338    ///
339    /// `None` for a set with no such shape, which is the table, the partitioned
340    /// body and an intset that has split into runs.
341    #[inline]
342    pub(crate) fn packed_bytes(&self) -> Option<&[u8]> {
343        match &self.body {
344            Body::Ints(s) if !self.ints_past_limit => s.as_bytes(),
345            Body::Packed(lp) => Some(lp.as_bytes()),
346            _ => None,
347        }
348    }
349
350    /// Write the set out as the bytes it becomes when it leaves memory.
351    ///
352    /// One form byte and then whatever that form needs. The two representations
353    /// that already have a flat layout, the single run intset and the listpack,
354    /// are written as themselves and cost one byte, because those bytes are
355    /// exactly what has to come back. The other two are written as their members.
356    ///
357    /// The form carries enough to land back in the same representation, which
358    /// [`crate::rdb`] deliberately does not: `ints_past_limit` rides in the top
359    /// bit of the form byte, and a set of integers that has split into runs is
360    /// its own form rather than a list of members, so that a million integers do
361    /// not come back as a hash table at thirty bytes a member. A value that
362    /// changed encoding because it was quiet long enough to be demoted would be a
363    /// value whose `OBJECT ENCODING` depends on memory pressure.
364    ///
365    /// See [`Set::thaw`], which is the other half and reads exactly this.
366    pub fn freeze(&self, out: &mut Vec<u8>) {
367        let past = if self.ints_past_limit { PAST_LIMIT } else { 0 };
368        match &self.body {
369            // `as_bytes` and not `packed_bytes`, because the ceiling only decides
370            // what the encoding is called and the flag above carries that. A set
371            // over the ceiling still has the Redis layout while it is one run,
372            // and writing the layout it has beats writing its members.
373            Body::Ints(s) => match s.as_bytes() {
374                Some(bytes) => {
375                    out.push(FORM_INTSET | past);
376                    out.extend_from_slice(bytes);
377                }
378                // Split into runs, so there is no one blob. Ascending, so the
379                // gaps are what gets written: a dense set of integers is a byte
380                // a member here where its members as digits would be several.
381                None => {
382                    out.push(FORM_RUNS | past);
383                    frozen::put_uint(out, s.len() as u64);
384                    let mut prev = 0u64;
385                    for i in 0..s.len() {
386                        let v = s.at(i) as u64;
387                        frozen::put_uint(out, v.wrapping_sub(prev));
388                        prev = v;
389                    }
390                }
391            },
392            Body::Packed(lp) => {
393                out.push(FORM_PACKED);
394                out.extend_from_slice(lp.as_bytes());
395            }
396            Body::Table(_) | Body::Split(_) => {
397                out.push(FORM_MEMBERS);
398                frozen::put_uint(out, self.len() as u64);
399                let mut digits = [0u8; DIGITS_MAX];
400                for m in self.iter() {
401                    match m {
402                        Member::Str(bytes) => frozen::put_bytes(out, bytes),
403                        // Neither of these two bodies stores an integer as one,
404                        // so this arm is for the type and not for a case that
405                        // arrives.
406                        Member::Int(n) => frozen::put_bytes(out, i64_digits(&mut digits, n)),
407                    }
408                }
409            }
410        }
411    }
412
413    /// Read back a set written by [`Set::freeze`].
414    ///
415    /// The count in the member form is a capacity and not a promise. It picks the
416    /// body and sizes it, and then the members decide the length, so a count that
417    /// disagrees with what follows costs a rehash and never a wrong set.
418    pub fn thaw(bytes: &[u8]) -> Result<Set, Broken> {
419        let mut cut = frozen::Cut::new(bytes);
420        let tag = cut.byte()?;
421        let ints_past_limit = tag & PAST_LIMIT != 0;
422        match tag & !PAST_LIMIT {
423            FORM_INTSET => Ok(Set {
424                body: Body::Ints(Intset::from_bytes(cut.rest()).map_err(|_| Broken::Body)?),
425                ints_past_limit,
426            }),
427            FORM_PACKED => Ok(Set {
428                body: Body::Packed(Listpack::from_bytes(cut.rest()).map_err(|_| Broken::Body)?),
429                ints_past_limit: false,
430            }),
431            FORM_RUNS => {
432                let n = usize::try_from(cut.uint()?).map_err(|_| Broken::Short)?;
433                let mut s = Intset::with_capacity(n);
434                let mut prev = 0u64;
435                for _ in 0..n {
436                    prev = prev.wrapping_add(cut.uint()?);
437                    s.add(prev as i64);
438                }
439                Ok(Set {
440                    body: Body::Ints(s),
441                    ints_past_limit,
442                })
443            }
444            FORM_MEMBERS => {
445                let n = usize::try_from(cut.uint()?).map_err(|_| Broken::Short)?;
446                // A member is at least one byte, so a count larger than what is
447                // left cannot be honest and is not worth an allocation.
448                if n > cut.rest().len() {
449                    return Err(Broken::Body);
450                }
451                // The same threshold the live set splits at, so a band that was
452                // demoted comes back a band. One that has shrunk under the
453                // threshold since comes back a table, which is the one place a
454                // round trip changes the body, and both of them answer
455                // `hashtable` so nothing outside can tell.
456                let mut set = Set {
457                    body: if n > PARTITION_AT {
458                        Body::Split(Parts::with_parts(parts_for(n)))
459                    } else {
460                        Body::Table(Elements::with_capacity(n))
461                    },
462                    ints_past_limit: false,
463                };
464                for _ in 0..n {
465                    let member = cut.bytes()?;
466                    match &mut set.body {
467                        Body::Table(t) => t.insert(member, ()).map_err(|_| Broken::Body)?,
468                        Body::Split(p) => p.insert(member, ()).map_err(|_| Broken::Body)?,
469                        _ => unreachable!("the body was just built as one of those two"),
470                    };
471                }
472                Ok(set)
473            }
474            _ => Err(Broken::Form),
475        }
476    }
477
478    /// How many members. This is `SCARD`.
479    #[inline]
480    pub fn len(&self) -> usize {
481        match &self.body {
482            Body::Ints(s) => s.len(),
483            Body::Packed(lp) => lp.len(),
484            Body::Table(t) => t.len(),
485            Body::Split(p) => p.len(),
486        }
487    }
488
489    /// Whether there are none.
490    ///
491    /// An empty set does not exist in Redis, so the caller deletes the key when
492    /// this turns true rather than storing an empty one.
493    #[inline]
494    pub fn is_empty(&self) -> bool {
495        self.len() == 0
496    }
497
498    /// Whether `member` is in the set. This is `SISMEMBER`.
499    #[must_use]
500    pub fn contains(&self, member: &[u8]) -> bool {
501        match &self.body {
502            // A member that is not an integer cannot be in a set of integers,
503            // and answering that costs a parse rather than a search.
504            Body::Ints(s) => parse_i64(member).is_some_and(|v| s.contains(v)),
505            Body::Packed(lp) => lp.find(member, 1).is_some(),
506            Body::Table(t) => t.contains(member),
507            Body::Split(p) => p.contains(member),
508        }
509    }
510
511    /// The same question asked with the work already done. See [`Needle`].
512    ///
513    /// This is what set algebra probes with. Every arm is the arm
514    /// [`Set::contains`] would have taken, with the parse and the hash lifted
515    /// out of it, so the two cannot disagree about what a member is.
516    #[must_use]
517    #[inline]
518    pub fn has(&self, needle: &Needle<'_>) -> bool {
519        match &self.body {
520            Body::Ints(s) => needle.int.is_some_and(|v| s.contains(v)),
521            Body::Packed(lp) => lp.find_parsed(needle.bytes, needle.int, 1).is_some(),
522            Body::Table(t) => t.contains_hashed(needle.hash, needle.bytes),
523            Body::Split(p) => p.contains_hashed(needle.hash, needle.bytes),
524        }
525    }
526
527    /// The member at `index`, in whatever order the representation holds them.
528    ///
529    /// Ascending for an intset, insertion order for the other two. Redis makes
530    /// no promise about set order and neither does this, but a uniform draw
531    /// needs positions and this is what gives it them (K9).
532    #[must_use]
533    pub fn at(&self, index: usize) -> Option<Member<'_>> {
534        match &self.body {
535            Body::Ints(s) => s.get(index).map(Member::Int),
536            Body::Packed(lp) => lp.get(index),
537            Body::Table(t) => t.at(index).map(|(name, _)| Member::Str(name)),
538            Body::Split(p) => p.at(index).map(|(name, _)| Member::Str(name)),
539        }
540    }
541
542    /// Every member.
543    pub fn iter(&self) -> impl Iterator<Item = Member<'_>> {
544        (0..self.len()).map(|i| self.at(i).expect("index is under the length"))
545    }
546
547    /// The members as a sorted array of integers, when that is what this is.
548    ///
549    /// The one place the representation is not an implementation detail, and it
550    /// is here for [`crate::setops`]: two sorted arrays intersect by stepping
551    /// through both of them with no hash anywhere, which is a different order of
552    /// cost from asking a table a question per member. That was worth nothing
553    /// while an all integer set turned into a table at five hundred and twelve
554    /// members, and it is worth a great deal now that it does not.
555    #[inline]
556    #[must_use]
557    pub const fn ints(&self) -> Option<&Intset> {
558        match &self.body {
559            Body::Ints(s) => Some(s),
560            _ => None,
561        }
562    }
563
564    /// Walk part of the set and say where to resume. This is `SSCAN`.
565    ///
566    /// Only the table and the partitioned band walk in windows. An intset or a
567    /// listpack hands back
568    /// every member in one call and a cursor of [`Cursor::END`], ignoring the
569    /// cursor it was given, which is what Redis does for the same two encodings
570    /// and for the same reason: a hundred and twenty eight members is smaller
571    /// than the reply header arithmetic to split them up, and a set that small
572    /// cannot block the loop long enough for the split to be worth anything.
573    ///
574    /// Ignoring the cursor is safe rather than merely convenient, because
575    /// promotion is one way. A set that gave a client a listpack cursor is not
576    /// going to be a listpack again, so the only way to arrive at those two arms
577    /// with a cursor from somewhere else is for the key to have been deleted and
578    /// remade underneath the scan, and returning everything to that client
579    /// returns a member twice at worst, which the guarantee allows.
580    ///
581    /// A table cursor arriving at the band is the one crossing that does happen,
582    /// because a set can split part way through a client's scan. That is handled
583    /// rather than ignored: a table cursor names one partition, and
584    /// [`Cursor::rebase`] reads the widening and restarts the walk at the top of
585    /// the new layout, so the client sees some members a second time and misses
586    /// none. Repeats are what the `SCAN` guarantee gives up in exchange for
587    /// surviving a resize, and a set only splits once.
588    pub fn scan<F>(&self, cursor: Cursor, count: usize, mut f: F) -> Cursor
589    where
590        F: FnMut(Member<'_>),
591    {
592        match &self.body {
593            Body::Table(t) => t.scan(cursor, count, |name, ()| f(Member::Str(name))),
594            Body::Split(p) => p.scan(cursor, count, |name, ()| f(Member::Str(name))),
595            // An intset past the ceiling is the one thing outside the table
596            // band that a single reply cannot hold. Redis has a dictionary by
597            // this point and walks it in windows, and a set of a million
598            // integers answering `SSCAN` with a million members in one go would
599            // be a several megabyte reply and a loop iteration nobody could
600            // measure. So it walks in windows too, and by index, which the runs
601            // answer in a walk down their tree rather than a walk along them.
602            //
603            // Downward, which is the direction the element table walks and for
604            // the same reason. Positions in a sorted array shift when a member
605            // below them goes, so an upward walk would miss a member for every
606            // one removed behind it, and `SSCAN` followed by `SREM` on what it
607            // found is the commonest thing anyone does with this command.
608            // Walking down means those removals are all above the cursor, where
609            // they cost nothing.
610            Body::Ints(s) if self.ints_past_limit && !s.is_empty() => {
611                let top = s.len() - 1;
612                let mut at = match cursor.rebase(1).idx() {
613                    Some(idx) => (idx as usize).min(top),
614                    None => top,
615                };
616                for _ in 0..count.max(1) {
617                    f(Member::Int(s.at(at)));
618                    if at == 0 {
619                        return Cursor::END;
620                    }
621                    at -= 1;
622                }
623                Cursor::at(1, 0, at as u64)
624            }
625            _ => {
626                for m in self.iter() {
627                    f(m);
628                }
629                Cursor::END
630            }
631        }
632    }
633
634    /// Bytes held by whichever representation this is.
635    #[must_use]
636    pub fn memory_bytes(&self) -> usize {
637        match &self.body {
638            Body::Ints(s) => s.memory_bytes(),
639            Body::Packed(lp) => lp.byte_len(),
640            Body::Table(t) => t.memory_bytes(),
641            Body::Split(p) => p.memory_bytes(),
642        }
643    }
644
645    /// What the slot array costs on its own, or nothing if there is not one.
646    ///
647    /// The three of these split [`Set::memory_bytes`] into the arrays it is made
648    /// of, which is what turns an argument about where the memory went into a
649    /// number. An intset and a listpack are one allocation with no index over
650    /// it, so they answer all of it under names and nothing under the other two.
651    #[must_use]
652    pub fn slot_bytes(&self) -> usize {
653        match &self.body {
654            Body::Ints(_) | Body::Packed(_) => 0,
655            Body::Table(t) => t.slot_bytes(),
656            Body::Split(p) => p.slot_bytes(),
657        }
658    }
659
660    /// What the row array costs on its own, capacity and not length.
661    #[must_use]
662    pub fn row_bytes(&self) -> usize {
663        match &self.body {
664            Body::Ints(_) | Body::Packed(_) => 0,
665            Body::Table(t) => t.row_bytes(),
666            Body::Split(p) => p.row_bytes(),
667        }
668    }
669
670    /// What the member bytes cost, live ones and dead ones together.
671    #[must_use]
672    pub fn name_bytes(&self) -> usize {
673        match &self.body {
674            Body::Ints(s) => s.memory_bytes(),
675            Body::Packed(lp) => lp.byte_len(),
676            Body::Table(t) => t.name_bytes(),
677            Body::Split(p) => p.name_bytes(),
678        }
679    }
680
681    /// Add `member`, promoting if it no longer fits. Answers whether it was new.
682    ///
683    /// This is `setTypeAdd`, arm for arm.
684    pub fn add(&mut self, member: &[u8], limits: &Limits) -> bool {
685        match &mut self.body {
686            Body::Table(t) => {
687                let new = t.insert(member, ()).is_ok_and(|old| old.is_none());
688                // Checked after the insert rather than before, so the set that
689                // splits is the one that has actually outgrown a table and not
690                // the one that is about to.
691                if t.len() > PARTITION_AT {
692                    self.become_split();
693                }
694                return new;
695            }
696            Body::Split(p) => {
697                let new = p.insert(member, ()).is_ok_and(|old| old.is_none());
698                // Asked rather than decided, because growing is a rehash of the
699                // whole set and the band leaves the timing to whoever knows
700                // whether this is one write or the middle of a bulk load.
701                if let Some(want) = p.wants_parts() {
702                    p.grow_to(want);
703                }
704                return new;
705            }
706            Body::Packed(lp) => {
707                if lp.find(member, 1).is_some() {
708                    return false;
709                }
710                if lp.len() < limits.max_listpack_entries
711                    && member.len() <= limits.max_listpack_value
712                {
713                    lp.push(member);
714                    return true;
715                }
716                // Too many members, or one too long. It falls out to a table.
717            }
718            Body::Ints(s) => {
719                if let Some(v) = parse_i64(member) {
720                    if !s.add(v) {
721                        return false;
722                    }
723                    // Strictly greater, so the 512th member is still an intset
724                    // and the 513th is what a real server would call a
725                    // hashtable. Nothing is rewritten, only the word changes.
726                    //
727                    // Unless the ceilings have been configured the wrong way
728                    // round, where a set past the intset ceiling is still under
729                    // the listpack one and a real server puts it in a listpack.
730                    // That set is a handful of members and there is no memory
731                    // argument for keeping it here, so it goes where it would
732                    // have gone.
733                    if s.len() > limits.max_intset_entries {
734                        if self.ints_fit_a_listpack_alone(limits) {
735                            self.become_listpack();
736                        } else {
737                            self.ints_past_limit = true;
738                        }
739                    }
740                    return true;
741                }
742                // Not an integer, so it is certainly not in a set of integers
743                // already. If the set is still small enough it becomes a
744                // listpack, and otherwise it falls out to a table.
745                if self.ints_fit_a_listpack(member, limits) {
746                    self.become_listpack();
747                    self.push_new(member);
748                    return true;
749                }
750            }
751        }
752        self.become_table(1);
753        self.push_new(member);
754        true
755    }
756
757    /// Put in a member already known to be new and known to fit where it is.
758    ///
759    /// Only ever called on the far side of a promotion, where both of those are
760    /// facts the promotion established and not things worth establishing twice.
761    fn push_new(&mut self, member: &[u8]) {
762        match &mut self.body {
763            Body::Packed(lp) => lp.push(member),
764            Body::Table(t) => {
765                t.insert(member, ())
766                    .expect("the table was sized for this one");
767            }
768            Body::Split(p) => {
769                p.insert(member, ())
770                    .expect("the band was sized for this one");
771            }
772            Body::Ints(_) => unreachable!("no promotion ever lands on an intset"),
773        }
774    }
775
776    /// Remove `member`. Answers whether it was there.
777    ///
778    /// Never demotes, which is Y4's one-way rule and Redis's behaviour.
779    pub fn remove(&mut self, member: &[u8]) -> bool {
780        match &mut self.body {
781            Body::Ints(s) => parse_i64(member).is_some_and(|v| s.remove(v)),
782            Body::Packed(lp) => match lp.find(member, 1) {
783                Some(at) => lp.delete(at, 1),
784                None => false,
785            },
786            Body::Table(t) => t.remove(member).is_some(),
787            Body::Split(p) => p.remove(member).is_some(),
788        }
789    }
790
791    /// Take out the member at `index` and hand it back.
792    ///
793    /// This is what `SPOP` runs on top of. The table moves its last row into the
794    /// hole rather than shifting, so the position of every other member is
795    /// stable except for one; the other two shift. Neither is a promise a caller
796    /// can lean on, and `SPOP` does not need one because it draws again from the
797    /// new length each time.
798    pub fn remove_at(&mut self, index: usize) -> Option<Vec<u8>> {
799        match &mut self.body {
800            Body::Ints(s) => {
801                let v = s.get(index)?;
802                s.remove(v);
803                let mut out = Vec::with_capacity(i64_len(v));
804                Member::Int(v).write_to(&mut out);
805                Some(out)
806            }
807            Body::Packed(lp) => {
808                let out = lp.get(index)?.to_vec();
809                lp.delete(index, 1);
810                Some(out)
811            }
812            Body::Table(t) => t.take_at(index).map(|(name, ())| name),
813            Body::Split(p) => p.take_at(index).map(|(name, ())| name),
814        }
815    }
816
817    /// Take out the member at `index` without building it into a `Vec` first.
818    ///
819    /// The same removal as [`Set::remove_at`] for a caller that has already read
820    /// the member and does not need it handed back. That caller is `SPOP` on the
821    /// wire, which reads with [`Set::at`], writes the bytes straight into the
822    /// reply buffer, and only then calls this. It is an allocation a member
823    /// saved on the one command in the set whose whole cost is the allocating.
824    ///
825    /// [`Set::remove_at`] stays for the embedded API, where the caller wants the
826    /// bytes and has nowhere to put them.
827    pub fn drop_at(&mut self, index: usize) -> bool {
828        match &mut self.body {
829            Body::Ints(s) => match s.get(index) {
830                Some(v) => {
831                    s.remove(v);
832                    true
833                }
834                None => false,
835            },
836            Body::Packed(lp) => {
837                if index >= lp.len() {
838                    return false;
839                }
840                lp.delete(index, 1);
841                true
842            }
843            Body::Table(t) => t.remove_at(index).is_some(),
844            Body::Split(p) => p.remove_at(index).is_some(),
845        }
846    }
847
848    /// Whether an intset plus one non integer member would still be a listpack.
849    ///
850    /// Three tests, and the first is the asymmetric one: the count is compared
851    /// against the *listpack* ceiling and not the intset one, so an intset with
852    /// two hundred members is already too big to become a listpack even though
853    /// it is a perfectly legal intset. The other two are the new member's length
854    /// and the longest existing member's length once it is written as digits,
855    /// which only bites when `set-max-listpack-value` has been turned down,
856    /// because no integer is more than twenty characters.
857    fn ints_fit_a_listpack(&self, member: &[u8], limits: &Limits) -> bool {
858        let Body::Ints(s) = &self.body else {
859            return false;
860        };
861        s.len() < limits.max_listpack_entries
862            && member.len() <= limits.max_listpack_value
863            && self.ints_are_short_enough(limits)
864    }
865
866    /// Whether this intset on its own would fit a listpack.
867    ///
868    /// The same question with no new member in it, which is what an intset that
869    /// has just passed `set-max-intset-entries` asks. It only ever answers yes
870    /// when the two ceilings have been configured the wrong way round, because
871    /// 512 is not under 128, and a server run that way expects a listpack there.
872    fn ints_fit_a_listpack_alone(&self, limits: &Limits) -> bool {
873        let Body::Ints(s) = &self.body else {
874            return false;
875        };
876        s.len() <= limits.max_listpack_entries && self.ints_are_short_enough(limits)
877    }
878
879    /// Whether every member, written as digits, is under the listpack ceiling.
880    fn ints_are_short_enough(&self, limits: &Limits) -> bool {
881        let Body::Ints(s) = &self.body else {
882            return false;
883        };
884        // The two ends bound the digits of everything between them, so there is
885        // nothing to walk.
886        let widest = s
887            .min()
888            .map(i64_len)
889            .unwrap_or(0)
890            .max(s.max().map(i64_len).unwrap_or(0));
891        widest <= limits.max_listpack_value
892    }
893
894    /// Rewrite as a listpack, which only an intset ever does.
895    fn become_listpack(&mut self) {
896        let Body::Ints(s) = &self.body else {
897            return;
898        };
899        let mut lp = Listpack::new();
900        let mut buf = Vec::with_capacity(20);
901        for v in s.iter() {
902            buf.clear();
903            Member::Int(v).write_to(&mut buf);
904            lp.push(&buf);
905        }
906        self.body = Body::Packed(lp);
907    }
908
909    /// Rewrite as an element table, with room for `extra` more members.
910    fn become_table(&mut self, extra: usize) {
911        let mut t = Elements::with_capacity(self.len() + extra);
912        let mut buf = Vec::with_capacity(20);
913        for m in self.iter() {
914            match m {
915                Member::Str(b) => {
916                    t.insert(b, ()).expect("room, and every member was unique");
917                }
918                Member::Int(v) => {
919                    buf.clear();
920                    Member::Int(v).write_to(&mut buf);
921                    t.insert(&buf, ())
922                        .expect("room, and every member was unique");
923                }
924            }
925        }
926        self.body = Body::Table(t);
927    }
928
929    /// Spread an element table over partitions.
930    ///
931    /// One way, like every other promotion here. A set that drops back under the
932    /// threshold keeps its partitions, which is Y4's rule and Redis's behaviour
933    /// for the encodings it does expose: the cost of a representation is paid
934    /// when it is entered, and paying it again on the way back out turns one
935    /// `SREM` at the boundary into a rehash of the whole set.
936    fn become_split(&mut self) {
937        if let Body::Table(t) = &self.body {
938            let p = Parts::from_table(t, parts_for(t.len()));
939            self.body = Body::Split(p);
940        }
941    }
942}
943
944impl Default for Set {
945    fn default() -> Set {
946        Set::new()
947    }
948}
949
950#[cfg(test)]
951mod tests {
952    use super::*;
953
954    fn of(members: &[&str]) -> Set {
955        let mut s = Set::new();
956        for m in members {
957            assert!(s.add(m.as_bytes(), &Limits::DEFAULT), "{m} was new");
958        }
959        s
960    }
961
962    /// Freeze a set, thaw it, and check that what came back is the same set in
963    /// the same representation.
964    ///
965    /// Members are compared as bytes and as a sorted list, because two of the
966    /// four bodies hold them in insertion order and a round trip through the
967    /// member form does not promise to preserve it. What it does promise is the
968    /// membership and the word `OBJECT ENCODING` answers, and those are what this
969    /// checks.
970    fn round_trip(set: &Set) -> Set {
971        let mut out = Vec::new();
972        set.freeze(&mut out);
973        let back = Set::thaw(&out).expect("what freeze wrote, thaw reads");
974        assert_eq!(back.len(), set.len(), "member count");
975        assert_eq!(back.encoding(), set.encoding(), "encoding");
976        let mut was: Vec<Vec<u8>> = set.iter().map(bytes_of).collect();
977        let mut now: Vec<Vec<u8>> = back.iter().map(bytes_of).collect();
978        was.sort_unstable();
979        now.sort_unstable();
980        assert_eq!(was, now, "members");
981        back
982    }
983
984    fn bytes_of(m: Member<'_>) -> Vec<u8> {
985        match m {
986            Member::Str(b) => b.to_vec(),
987            Member::Int(n) => n.to_string().into_bytes(),
988        }
989    }
990
991    #[test]
992    fn a_frozen_set_comes_back_in_the_body_it_left() {
993        // An intset, which is one run and goes out as the Redis layout.
994        let mut ints = Set::new();
995        for i in 0..100i64 {
996            ints.add(i.to_string().as_bytes(), &Limits::DEFAULT);
997        }
998        assert_eq!(ints.encoding(), Encoding::Intset);
999        round_trip(&ints);
1000
1001        // A listpack, same story with the other flat layout.
1002        let packed = of(&["one", "two", "three"]);
1003        assert_eq!(packed.encoding(), Encoding::Listpack);
1004        round_trip(&packed);
1005
1006        // A table, which goes out as its members.
1007        let mut table = Set::new();
1008        for i in 0..500 {
1009            table.add(format!("member:{i}").as_bytes(), &Limits::DEFAULT);
1010        }
1011        assert_eq!(table.encoding(), Encoding::Hashtable);
1012        round_trip(&table);
1013
1014        // An empty set, which the keyspace never stores but which is one byte
1015        // and should not need a special case anywhere.
1016        round_trip(&Set::new());
1017    }
1018
1019    #[test]
1020    fn an_intset_past_its_ceiling_still_calls_itself_a_hashtable_after_a_round_trip() {
1021        let mut s = Set::new();
1022        for i in 0..2_000i64 {
1023            s.add(i.to_string().as_bytes(), &Limits::DEFAULT);
1024        }
1025        // Past `set-max-intset-entries`, so it answers `hashtable` while its
1026        // body is still an intset. That flag is the whole reason the form byte
1027        // has a spare bit, and losing it would change what a client sees.
1028        assert_eq!(s.encoding(), Encoding::Hashtable);
1029        let back = round_trip(&s);
1030        assert!(back.ints().is_some(), "still an intset underneath");
1031    }
1032
1033    #[test]
1034    fn a_large_set_of_integers_does_not_come_back_as_a_table() {
1035        // Past the point where the runs split, which is the case the gap form
1036        // exists for. Through a table it would be thirty bytes a member.
1037        let mut s = Set::new();
1038        for i in 0..200_000i64 {
1039            s.add((i * 3).to_string().as_bytes(), &Limits::DEFAULT);
1040        }
1041        assert!(s.ints().is_some_and(|i| i.as_bytes().is_none()), "split");
1042        let mut out = Vec::new();
1043        s.freeze(&mut out);
1044        let back = round_trip(&s);
1045        assert!(back.ints().is_some(), "came back as integers");
1046        // Gaps of three, so two bytes a member is the ceiling here and the
1047        // digits would have been six. The check is against the loose bound
1048        // because the point is the order of magnitude and not the byte.
1049        assert!(out.len() < s.len() * 3, "{} bytes", out.len());
1050    }
1051
1052    #[test]
1053    fn a_partitioned_set_comes_back_partitioned() {
1054        let mut s = Set::new();
1055        for i in 0..=PARTITION_AT {
1056            s.add(format!("member:{i}").as_bytes(), &Limits::DEFAULT);
1057        }
1058        assert!(matches!(s.body, Body::Split(_)), "split");
1059        let back = round_trip(&s);
1060        assert!(matches!(back.body, Body::Split(_)), "still split");
1061    }
1062
1063    #[test]
1064    fn a_frozen_set_that_arrives_damaged_is_an_error_and_not_a_panic() {
1065        let s = of(&["one", "two", "three"]);
1066        let mut out = Vec::new();
1067        s.freeze(&mut out);
1068        for cut in 1..out.len() {
1069            // Every prefix, because a short read off the store can end anywhere
1070            // and none of them may panic.
1071            let _ = Set::thaw(&out[..cut]);
1072        }
1073        assert_eq!(Set::thaw(&[]).err(), Some(Broken::Short));
1074        assert_eq!(Set::thaw(&[9]).err(), Some(Broken::Form));
1075        assert_eq!(Set::thaw(&[FORM_INTSET, 1, 2]).err(), Some(Broken::Body));
1076    }
1077
1078    /// What a set actually costs per member, which is half of M3's memory gate
1079    /// row and was an argument rather than a number until this was written.
1080    ///
1081    /// Run it with `cargo test -p yo-kv --release measure_bytes_per_member --
1082    /// --ignored --nocapture`. Ignored because a million members is not
1083    /// something every `cargo test` should pay for, and it prints rather than
1084    /// asserts because the number it prints is the thing being reported.
1085    ///
1086    /// Three shapes, because the gate names one of them and the other two are
1087    /// what most sets actually hold. Integers first, at every band an all
1088    /// integer set passes through, then the same counts of integers scattered
1089    /// over a wide range, then strings, which never see the intset at all.
1090    ///
1091    /// The scattered row is there because the dense one is the most favourable
1092    /// input this structure will ever be given and a gate read off it would be
1093    /// a gate read off the easy case. `0..n` is sorted, contiguous and arrives
1094    /// in order, so every run fills to `RUN_MAX` and never splits. Real integer
1095    /// sets are ids with holes in them, arriving in whatever order the writer
1096    /// had, and the number that matters is the one they produce.
1097    #[test]
1098    #[ignore = "a measurement, run it by name"]
1099    fn measure_bytes_per_member() {
1100        let limits = Limits::DEFAULT;
1101        for n in [512usize, 1_000, 100_000, 1_000_000] {
1102            let mut s = Set::new();
1103            for i in 0..n {
1104                s.add(i.to_string().as_bytes(), &limits);
1105            }
1106            println!(
1107                "int   n={n:<9} band={:<10} total={:<10} per_member={:.2}",
1108                band(&s),
1109                s.memory_bytes(),
1110                s.memory_bytes() as f64 / n as f64
1111            );
1112        }
1113        // The same counts, scattered over a range sixteen times as wide and
1114        // arriving out of order. A cheap multiplicative shuffle rather than a
1115        // real generator, because what this needs is holes and disorder and not
1116        // statistical quality.
1117        for n in [512usize, 1_000, 100_000, 1_000_000] {
1118            let mut s = Set::new();
1119            for i in 0..n {
1120                let v = (i as u64).wrapping_mul(0x9e37_79b9_7f4a_7c15) % (n as u64 * 16);
1121                s.add(v.to_string().as_bytes(), &limits);
1122            }
1123            println!(
1124                "sparse n={n:<9} band={:<10} total={:<10} members={:<9} per_member={:.2}",
1125                band(&s),
1126                s.memory_bytes(),
1127                s.len(),
1128                s.memory_bytes() as f64 / s.len() as f64
1129            );
1130        }
1131        // Sixteen byte members, so the payload is a round number and the
1132        // overhead is whatever is above it.
1133        for n in [128usize, 1_000, 100_000, 1_000_000] {
1134            let mut s = Set::new();
1135            let mut payload = 0usize;
1136            for i in 0..n {
1137                let m = format!("member:{i:09}");
1138                payload += m.len();
1139                s.add(m.as_bytes(), &limits);
1140            }
1141            let total = s.memory_bytes();
1142            let per = |b: usize| b as f64 / n as f64;
1143            println!(
1144                "bytes n={n:<9} band={:<10} total={total:<10} payload={payload:<9} per_member={:.2} over_per_member={:.2} slots={:.2} rows={:.2} names={:.2} name_slack={:.2}",
1145                band(&s),
1146                per(total),
1147                per(total - payload),
1148                per(s.slot_bytes()),
1149                per(s.row_bytes()),
1150                per(s.name_bytes()),
1151                per(s.name_bytes() - payload)
1152            );
1153        }
1154    }
1155
1156    /// Which of the four a set is in, spelled out rather than through
1157    /// [`Set::encoding`], which folds the two table bands into one word because
1158    /// that is what `OBJECT ENCODING` has to say.
1159    fn band(s: &Set) -> &'static str {
1160        match &s.body {
1161            Body::Ints(_) => "intset",
1162            Body::Packed(_) => "listpack",
1163            Body::Table(_) => "table",
1164            Body::Split(_) => "split",
1165        }
1166    }
1167
1168    /// Everything about the partitioned band that needs a real set past the real
1169    /// threshold, in one test, because building 262,145 members is the expensive
1170    /// part and there is no reason to pay for it four times.
1171    #[test]
1172    fn a_set_past_the_threshold_splits_without_the_client_being_able_to_tell() {
1173        let limits = Limits::DEFAULT;
1174        let mut s = Set::new();
1175        // One short of the threshold. Still one table: the check is strictly
1176        // greater, so the set that splits is the one that has outgrown a table
1177        // and not the one that is about to.
1178        for i in 0..PARTITION_AT {
1179            assert!(s.add(format!("m{i}").as_bytes(), &limits));
1180        }
1181        assert!(matches!(s.body, Body::Table(_)));
1182        assert_eq!(s.len(), PARTITION_AT);
1183
1184        // The member that tips it over.
1185        assert!(s.add(b"tipping", &limits));
1186        assert!(matches!(s.body, Body::Split(_)), "it should have split");
1187        assert_eq!(s.len(), PARTITION_AT + 1);
1188
1189        // And the client cannot tell. This is the whole point: Redis has three
1190        // set encodings and a fourth word here breaks every suite that reads it.
1191        assert_eq!(s.encoding(), Encoding::Hashtable);
1192        assert_eq!(s.encoding().name(), "hashtable");
1193
1194        // Every member survived the rehash, asked both ways.
1195        assert!(s.contains(b"tipping"));
1196        assert!(s.contains(b"m0"));
1197        assert!(s.contains(b"m262143"));
1198        assert!(!s.contains(b"m262144"));
1199        assert!(s.has(&Needle::new(b"m1000")));
1200        assert!(!s.has(&Needle::new(b"nothing")));
1201
1202        // A rewrite is still not an add.
1203        assert!(!s.add(b"m0", &limits));
1204        assert_eq!(s.len(), PARTITION_AT + 1);
1205
1206        // Removing goes back through the same partition it went into, and the
1207        // band never demotes however far the set shrinks.
1208        assert!(s.remove(b"tipping"));
1209        assert!(!s.remove(b"tipping"));
1210        assert_eq!(s.len(), PARTITION_AT);
1211        assert!(matches!(s.body, Body::Split(_)), "promotion is one way");
1212        assert_eq!(s.encoding(), Encoding::Hashtable);
1213
1214        // The draw `SPOP` and `SRANDMEMBER` run on reaches the last position,
1215        // which is the one that has to land in the highest non empty partition.
1216        let last = s.at(s.len() - 1).expect("inside the set").to_vec();
1217        assert!(s.contains(&last));
1218        assert!(s.at(s.len()).is_none());
1219        assert_eq!(s.remove_at(s.len() - 1), Some(last.clone()));
1220        assert!(!s.contains(&last));
1221        assert!(s.drop_at(0));
1222        assert_eq!(s.len(), PARTITION_AT - 2);
1223
1224        // And a full scan sees every member exactly once.
1225        let mut seen = 0usize;
1226        let mut cursor = Cursor::START;
1227        let mut rounds = 0;
1228        loop {
1229            cursor = s.scan(cursor, 1_000, |_| seen += 1);
1230            rounds += 1;
1231            assert!(rounds < 100_000, "the scan is not finishing");
1232            if cursor.is_end() {
1233                break;
1234            }
1235        }
1236        assert_eq!(seen, s.len());
1237    }
1238
1239    /// The crossing that actually happens in production: a client holding a
1240    /// cursor from before the split. It has to see every member that stayed, and
1241    /// repeats are what the `SCAN` guarantee gives up in exchange.
1242    #[test]
1243    fn a_scan_survives_the_set_splitting_underneath_it() {
1244        let limits = Limits::DEFAULT;
1245        let mut s = Set::new();
1246        for i in 0..PARTITION_AT {
1247            s.add(format!("m{i}").as_bytes(), &limits);
1248        }
1249        assert!(matches!(s.body, Body::Table(_)));
1250
1251        let mut seen = Vec::new();
1252        let cursor = s.scan(Cursor::START, 5_000, |m| seen.push(m.to_vec()));
1253        assert!(!cursor.is_end(), "the scan should have stopped part way");
1254
1255        s.add(b"tipping", &limits);
1256        assert!(matches!(s.body, Body::Split(_)));
1257
1258        let mut cursor = cursor;
1259        let mut rounds = 0;
1260        loop {
1261            cursor = s.scan(cursor, 5_000, |m| seen.push(m.to_vec()));
1262            rounds += 1;
1263            assert!(rounds < 100_000, "the scan is not finishing");
1264            if cursor.is_end() {
1265                break;
1266            }
1267        }
1268        seen.sort_unstable();
1269        seen.dedup();
1270        assert_eq!(
1271            seen.len(),
1272            PARTITION_AT + 1,
1273            "the split lost a member the client was entitled to"
1274        );
1275    }
1276
1277    #[test]
1278    fn a_hint_past_the_threshold_builds_the_band_up_front() {
1279        let limits = Limits::DEFAULT;
1280        // A caller loading a million members should not fill one table, cross the
1281        // threshold and then rehash the lot.
1282        let s = Set::with_hint(b"first", 1_000_000, &limits);
1283        assert!(matches!(s.body, Body::Split(_)));
1284        assert_eq!(s.encoding(), Encoding::Hashtable);
1285        assert!(s.is_empty());
1286
1287        // A hint at the threshold is still one table, matching the add path.
1288        let s = Set::with_hint(b"first", PARTITION_AT, &limits);
1289        assert!(matches!(s.body, Body::Table(_)));
1290
1291        // And a hint is only a hint: the band takes members like anything else.
1292        let mut s = Set::with_hint(b"a", 1_000_000, &limits);
1293        assert!(s.add(b"a", &limits));
1294        assert!(!s.add(b"a", &limits));
1295        assert!(s.contains(b"a"));
1296        assert_eq!(s.len(), 1);
1297        assert_eq!(s.at(0).map(|m| m.to_vec()), Some(b"a".to_vec()));
1298    }
1299
1300    fn members(s: &Set) -> Vec<String> {
1301        let mut v: Vec<String> = s
1302            .iter()
1303            .map(|m| String::from_utf8(m.to_vec()).expect("utf8 in these tests"))
1304            .collect();
1305        // Order is not part of the contract and the three representations do not
1306        // agree on it, so every assertion here is against a sorted list.
1307        v.sort();
1308        v
1309    }
1310
1311    #[test]
1312    fn a_new_set_is_an_empty_intset() {
1313        let s = Set::new();
1314        assert_eq!(s.encoding(), Encoding::Intset);
1315        assert_eq!(s.len(), 0);
1316        assert!(s.is_empty());
1317        assert!(!s.contains(b"1"));
1318        assert_eq!(s.at(0), None);
1319    }
1320
1321    #[test]
1322    fn integers_stay_an_intset_and_come_back_as_members() {
1323        let s = of(&["1", "2", "3"]);
1324        assert_eq!(s.encoding(), Encoding::Intset);
1325        assert_eq!(members(&s), ["1", "2", "3"]);
1326        assert!(s.contains(b"2"));
1327        assert!(!s.contains(b"4"));
1328        assert!(!s.contains(b"two"));
1329    }
1330
1331    #[test]
1332    fn adding_the_same_member_twice_says_so_in_all_three() {
1333        let mut ints = of(&["1", "2"]);
1334        assert!(!ints.add(b"2", &Limits::DEFAULT));
1335        assert_eq!(ints.len(), 2);
1336
1337        let mut packed = of(&["a", "b"]);
1338        assert_eq!(packed.encoding(), Encoding::Listpack);
1339        assert!(!packed.add(b"b", &Limits::DEFAULT));
1340        assert_eq!(packed.len(), 2);
1341
1342        let mut table = of(&["a", "b"]);
1343        table.become_table(0);
1344        assert!(!table.add(b"b", &Limits::DEFAULT));
1345        assert_eq!(table.len(), 2);
1346    }
1347
1348    #[test]
1349    fn a_string_turns_a_small_intset_into_a_listpack() {
1350        let mut s = of(&["1", "2", "3"]);
1351        assert!(s.add(b"hello", &Limits::DEFAULT));
1352        assert_eq!(s.encoding(), Encoding::Listpack);
1353        assert_eq!(members(&s), ["1", "2", "3", "hello"]);
1354        assert!(s.contains(b"1"), "the integers survived the rewrite");
1355        assert!(s.contains(b"hello"));
1356    }
1357
1358    #[test]
1359    fn a_string_turns_a_big_intset_straight_into_a_table() {
1360        // The asymmetric rule, and the one a natural implementation gets wrong.
1361        // Two hundred integers is a legal intset and is already past the
1362        // listpack ceiling, so this never passes through the listpack at all.
1363        let mut s = Set::new();
1364        for i in 0..200 {
1365            s.add(i.to_string().as_bytes(), &Limits::DEFAULT);
1366        }
1367        assert_eq!(s.encoding(), Encoding::Intset);
1368        assert_eq!(s.len(), 200);
1369
1370        assert!(s.add(b"hello", &Limits::DEFAULT));
1371        assert_eq!(s.encoding(), Encoding::Hashtable);
1372        assert_eq!(s.len(), 201);
1373        assert!(s.contains(b"199"));
1374        assert!(s.contains(b"hello"));
1375    }
1376
1377    #[test]
1378    fn an_intset_holds_five_hundred_and_twelve_and_converts_at_the_next_one() {
1379        let mut s = Set::new();
1380        for i in 0..512 {
1381            s.add(i.to_string().as_bytes(), &Limits::DEFAULT);
1382        }
1383        assert_eq!(s.encoding(), Encoding::Intset, "512 is still an intset");
1384        assert_eq!(s.len(), 512);
1385
1386        s.add(b"512", &Limits::DEFAULT);
1387        assert_eq!(s.encoding(), Encoding::Hashtable, "513 is not");
1388        assert_eq!(s.len(), 513);
1389        // And not a listpack on the way, because 513 is well past 128.
1390        for i in 0..513 {
1391            assert!(s.contains(i.to_string().as_bytes()), "{i} survived");
1392        }
1393    }
1394
1395    #[test]
1396    fn a_listpack_converts_at_a_hundred_and_twenty_eight_members() {
1397        let mut s = of(&["x"]);
1398        assert_eq!(s.encoding(), Encoding::Listpack);
1399        for i in 0..127 {
1400            s.add(format!("m{i}").as_bytes(), &Limits::DEFAULT);
1401        }
1402        assert_eq!(s.len(), 128);
1403        assert_eq!(s.encoding(), Encoding::Listpack, "128 is still a listpack");
1404
1405        s.add(b"one more", &Limits::DEFAULT);
1406        assert_eq!(s.len(), 129);
1407        assert_eq!(s.encoding(), Encoding::Hashtable);
1408        assert!(s.contains(b"x"));
1409        assert!(s.contains(b"m126"));
1410        assert!(s.contains(b"one more"));
1411    }
1412
1413    #[test]
1414    fn a_long_member_converts_a_listpack_whatever_the_count() {
1415        let mut s = of(&["a"]);
1416        let long = vec![b'z'; 65];
1417        assert!(s.add(&long, &Limits::DEFAULT));
1418        assert_eq!(s.encoding(), Encoding::Hashtable, "65 is past 64");
1419        assert_eq!(s.len(), 2);
1420        assert!(s.contains(&long));
1421
1422        // And exactly at the boundary it does not.
1423        let mut ok = of(&["a"]);
1424        ok.add(&[b'z'; 64], &Limits::DEFAULT);
1425        assert_eq!(ok.encoding(), Encoding::Listpack, "64 fits");
1426    }
1427
1428    #[test]
1429    fn a_long_member_sends_an_intset_to_a_table_and_not_a_listpack() {
1430        let mut s = of(&["1", "2"]);
1431        assert!(s.add(&[b'z'; 65], &Limits::DEFAULT));
1432        assert_eq!(s.encoding(), Encoding::Hashtable);
1433        assert_eq!(s.len(), 3);
1434    }
1435
1436    #[test]
1437    fn the_limits_are_configuration_and_moving_them_moves_the_encodings() {
1438        let tight = Limits {
1439            max_intset_entries: 2,
1440            max_listpack_entries: 2,
1441            max_listpack_value: 3,
1442        };
1443        let mut s = Set::new();
1444        s.add(b"1", &tight);
1445        s.add(b"2", &tight);
1446        assert_eq!(s.encoding(), Encoding::Intset);
1447        s.add(b"3", &tight);
1448        assert_eq!(s.encoding(), Encoding::Hashtable, "three is past two");
1449
1450        // And a member longer than three characters cannot go in a listpack.
1451        let mut t = Set::new();
1452        t.add(b"abc", &tight);
1453        assert_eq!(t.encoding(), Encoding::Listpack, "three characters fit");
1454        t.add(b"defg", &tight);
1455        assert_eq!(t.encoding(), Encoding::Hashtable, "four do not");
1456        assert!(t.contains(b"abc"));
1457        assert!(t.contains(b"defg"));
1458
1459        // Including as the first member, which is a table from the off rather
1460        // than a listpack that converts on the next thing to arrive. Redis gets
1461        // to the same place by the other road: it creates the listpack, tries
1462        // the add, and converts before the reply.
1463        let mut u = Set::new();
1464        u.add(b"abcd", &tight);
1465        assert_eq!(u.encoding(), Encoding::Hashtable);
1466        assert!(u.contains(b"abcd"));
1467    }
1468
1469    #[test]
1470    fn with_hint_picks_the_representation_up_front() {
1471        let d = &Limits::DEFAULT;
1472        assert_eq!(
1473            Set::with_hint(b"1", 10, d).encoding(),
1474            Encoding::Intset,
1475            "an integer and few enough of them"
1476        );
1477        assert_eq!(
1478            Set::with_hint(b"1", 1000, d).encoding(),
1479            Encoding::Hashtable,
1480            "an integer and too many"
1481        );
1482        assert_eq!(
1483            Set::with_hint(b"x", 10, d).encoding(),
1484            Encoding::Listpack,
1485            "not an integer and few enough"
1486        );
1487        assert_eq!(
1488            Set::with_hint(b"x", 1000, d).encoding(),
1489            Encoding::Hashtable,
1490            "not an integer and too many"
1491        );
1492    }
1493
1494    #[test]
1495    fn removing_works_in_all_three_and_never_demotes() {
1496        let mut ints = of(&["1", "2", "3"]);
1497        assert!(ints.remove(b"2"));
1498        assert!(!ints.remove(b"2"));
1499        assert!(!ints.remove(b"nope"), "not an integer, so not a member");
1500        assert_eq!(members(&ints), ["1", "3"]);
1501        assert_eq!(ints.encoding(), Encoding::Intset);
1502
1503        let mut packed = of(&["a", "b", "c"]);
1504        assert!(packed.remove(b"b"));
1505        assert!(!packed.remove(b"b"));
1506        assert_eq!(members(&packed), ["a", "c"]);
1507        assert_eq!(packed.encoding(), Encoding::Listpack);
1508
1509        let mut table = of(&["a", "b", "c"]);
1510        table.become_table(0);
1511        assert!(table.remove(b"b"));
1512        assert!(!table.remove(b"b"));
1513        assert_eq!(members(&table), ["a", "c"]);
1514        assert_eq!(
1515            table.encoding(),
1516            Encoding::Hashtable,
1517            "down to two members and still a table"
1518        );
1519    }
1520
1521    #[test]
1522    fn a_set_can_be_emptied_a_member_at_a_time() {
1523        for mut s in [of(&["1", "2", "3"]), of(&["a", "b", "c"])] {
1524            let all: Vec<Vec<u8>> = s.iter().map(|m| m.to_vec()).collect();
1525            for m in &all {
1526                assert!(s.remove(m));
1527            }
1528            assert!(s.is_empty());
1529            assert_eq!(s.at(0), None);
1530        }
1531    }
1532
1533    #[test]
1534    fn removing_by_position_hands_the_member_back() {
1535        // What `SPOP` runs on. Drawing position zero every time has to empty the
1536        // set rather than run off the end or repeat a member, in all three.
1537        let mut table = of(&["a", "b", "c", "d"]);
1538        table.become_table(0);
1539        for mut s in [
1540            of(&["10", "20", "30", "40"]),
1541            of(&["a", "b", "c", "d"]),
1542            table,
1543        ] {
1544            let mut got = Vec::new();
1545            while !s.is_empty() {
1546                got.push(String::from_utf8(s.remove_at(0).expect("not empty")).expect("utf8"));
1547            }
1548            got.sort();
1549            assert_eq!(got.len(), 4, "four members and no repeats");
1550            assert_eq!(s.len(), 0);
1551            assert_eq!(s.remove_at(0), None);
1552        }
1553    }
1554
1555    #[test]
1556    fn an_integer_member_is_the_same_member_however_it_is_written() {
1557        // An intset holds 42 as a number, so `SADD s 42` twice is one member.
1558        // `042` does not parse as an integer, so it is a different member and it
1559        // converts the set, which is what a real server does too.
1560        let mut s = of(&["42"]);
1561        assert!(!s.add(b"42", &Limits::DEFAULT));
1562        assert_eq!(s.len(), 1);
1563        assert!(s.add(b"042", &Limits::DEFAULT));
1564        assert_eq!(s.encoding(), Encoding::Listpack);
1565        assert_eq!(members(&s), ["042", "42"]);
1566        assert!(s.contains(b"42"));
1567        assert!(s.contains(b"042"));
1568    }
1569
1570    #[test]
1571    fn the_small_bands_answer_a_scan_in_one_go() {
1572        for s in [of(&["1", "2", "3"]), of(&["a", "b", "c"])] {
1573            let mut seen = Vec::new();
1574            // A count of one, which the table band would honour and these two
1575            // do not, and a cursor from nowhere, which these two ignore.
1576            let next = s.scan(Cursor::at(1, 0, 99), 1, |m| seen.push(m.to_vec()));
1577            assert!(next.is_end(), "{:?} split a scan up", s.encoding());
1578            assert_eq!(seen.len(), 3);
1579        }
1580    }
1581
1582    #[test]
1583    fn the_table_band_walks_a_scan_in_windows_and_misses_nothing() {
1584        let mut s = Set::new();
1585        for i in 0..300 {
1586            s.add(format!("m{i}").as_bytes(), &Limits::DEFAULT);
1587        }
1588        assert_eq!(s.encoding(), Encoding::Hashtable);
1589
1590        let mut seen = Vec::new();
1591        let mut c = Cursor::START;
1592        let mut turns = 0;
1593        loop {
1594            c = s.scan(c, 7, |m| seen.push(m.to_vec()));
1595            turns += 1;
1596            assert!(turns < 100, "the scan did not finish");
1597            if c.is_end() {
1598                break;
1599            }
1600        }
1601        assert!(
1602            turns > 1,
1603            "a window of seven over three hundred took one turn"
1604        );
1605        seen.sort();
1606        seen.dedup();
1607        assert_eq!(seen.len(), 300, "every member came back at least once");
1608    }
1609
1610    #[test]
1611    fn a_conversion_loses_no_member_at_any_of_the_three_boundaries() {
1612        // One walk over each path out of an intset and out of a listpack, each
1613        // checking every member is still findable after the rewrite rather than
1614        // only checking the count.
1615        let mut wide = Set::new();
1616        for v in [i64::MIN, -1, 0, 1, i64::MAX] {
1617            wide.add(v.to_string().as_bytes(), &Limits::DEFAULT);
1618        }
1619        wide.add(b"str", &Limits::DEFAULT);
1620        assert_eq!(wide.encoding(), Encoding::Listpack);
1621        for v in [i64::MIN, -1, 0, 1, i64::MAX] {
1622            assert!(wide.contains(v.to_string().as_bytes()), "{v} survived");
1623        }
1624
1625        wide.add(&[b'q'; 100], &Limits::DEFAULT);
1626        assert_eq!(wide.encoding(), Encoding::Hashtable);
1627        for v in [i64::MIN, -1, 0, 1, i64::MAX] {
1628            assert!(
1629                wide.contains(v.to_string().as_bytes()),
1630                "{v} survived twice"
1631            );
1632        }
1633        assert!(wide.contains(b"str"));
1634        assert_eq!(wide.len(), 7);
1635    }
1636}