Skip to main content

yo_kv/
intset.rs

1//! A set of integers as sorted packed arrays, which is Redis's intset in runs.
2//!
3//! A set whose members all parse as integers is held as the integers themselves,
4//! sorted, in the narrowest width that covers the widest of them, with no hash
5//! table and no per member allocation anywhere. One run looks exactly like a
6//! Redis intset, because it is one:
7//!
8//! ```text
9//! +----------+----------+---------+---------+-----+
10//! | u32 width| u32 count| member 0| member 1| ... |
11//! +----------+----------+---------+---------+-----+
12//!   2, 4 or 8   how many   sorted ascending, width bytes each
13//! ```
14//!
15//! Eight bytes of header and then nothing but members. At two byte width that is
16//! two bytes an element with no overhead at all, which is the number G8 asks for
17//! from a set of integers, and it is why this exists as a third representation
18//! rather than everything small going in a listpack.
19//!
20//! Two byte width is not only for sets of small numbers. Past one run a run
21//! stores its members as distances from a base of its own rather than as
22//! themselves, so a set of billions is still two bytes a member. That is the
23//! frame of reference, and the `Run` type is where it is explained.
24//!
25//! Measured, on a set of five hundred and twelve small integers, that is 2.0
26//! bytes a member against the listpack's 3.0 and the element table's 24.0. The
27//! header is the only thing between it and exactly two, and it is amortised away
28//! by about sixty members.
29//!
30//! Both header fields are little endian whatever the machine is, because Redis
31//! writes them that way: `intrev32ifbe` is a no-op on a little endian host and a
32//! byte swap on a big endian one, so the bytes on the wire and in the file are
33//! little endian from either. Getting that backwards would produce a file a real
34//! server cannot read, on the one class of machine nobody tests on.
35//!
36//! # Why there is more than one array
37//!
38//! Redis gives up on the intset at five hundred and twelve members and rehashes
39//! the set into a dictionary, which measured here at 24.60 to 30.92 bytes a
40//! member against the intset's 4.00. That is a twelvefold jump in memory for a
41//! set that got one member bigger, and it is worth asking what forced it.
42//!
43//! What forced it is that Redis has one array. An insert into the middle of one
44//! sorted array memmoves the tail, so a million member set moves half a megabyte
45//! per `SADD`, and no ceiling on the memory saves you from that. The conversion
46//! is a fix for the memmove and the memory is what it costs.
47//!
48//! That is an argument against having one array. It is not an argument for
49//! giving up two bytes a member. So the members here live in a list of runs,
50//! each one a complete intset in Redis's own layout, holding disjoint ranges of
51//! values in ascending order. A run is capped at `RUN_MAX` members, so the
52//! memmove an insert pays is bounded by the run and not by the set: a thousand
53//! bytes at two byte width, whether the set holds a thousand members or a
54//! hundred million. Membership is a binary search over the run maxima to pick
55//! the run and a binary search inside it, so it stays logarithmic in the whole
56//! set with both searches in cache.
57//!
58//! `RUN_MAX` is five hundred and twelve on purpose. A set that a default
59//! configured Redis would still call an intset is exactly one run here, so its
60//! bytes are still one array and [`Intset::as_bytes`] still answers with a blob
61//! a real server can read. The runs only appear past the point where Redis has
62//! stopped having an intset at all.
63//!
64//! # Finding the member at a position
65//!
66//! `SRANDMEMBER` and `SPOP` need the member at an index, which one array answers
67//! by multiplying and a list of runs does not. Adding up run lengths would be
68//! linear in the number of runs, which is 3906 of them at a million members, and
69//! `SRANDMEMBER key 100` would walk that four thousand entry array a hundred
70//! times.
71//!
72//! So the run lengths are kept in a Fenwick tree, which answers "which run holds
73//! position `k`, and how far into it" in a walk down the tree rather than a walk
74//! along the runs. An add or a remove that leaves the run structure alone is one
75//! more walk down the tree, and a split or a merge rebuilds it, which is linear
76//! in the number of runs and happens once per couple of hundred writes. Measured
77//! at 11.2 ns for the member at a position on a set of a million.
78//!
79//! # What the runs cost, measured
80//!
81//! The whole change is a memory argument, so the memory row is the one to read
82//! first. Per member, on an all integer set, before the runs against after, from
83//! `measure_bytes_per_member`:
84//!
85//! ```text
86//!   members     one array          the runs     and the frame
87//!       512     4.00  intset       2.09         2.11
88//!     1,000    24.60  hashtable    2.22         2.25
89//!   100,000    30.92  hashtable    3.57         2.26
90//! 1,000,000    29.19  hashtable    4.11         2.21
91//! ```
92//!
93//! The four bytes a member the larger sizes used to cost was not slack, it was
94//! the width: values up to a million need four byte slots, so four bytes was the
95//! floor and the overhead above it was eight hundredths of a byte. Getting under
96//! it meant making the width smaller rather than the overhead, which is what the
97//! frame of reference does, and it is the whole of the last column.
98//!
99//! The frame costs eight bytes a run, which is the base sitting in the `Run`
100//! next to the buffer, and that is the two hundredths of a byte the first two
101//! rows go backwards by. It is a bad trade on a set of small integers, which had
102//! no width to save, and it pays for itself several hundred times over on
103//! anything bigger.
104//!
105//! Filled in scattered order rather than ascending the answer is much the same,
106//! 2.11, 2.24, 2.23 and 2.28, against 4.33 at a million before the frame. That
107//! matters more than the ascending row does, because a run whose members arrive
108//! out of order is the one that has to widen, and it is the shape a real
109//! keyspace has.
110//!
111//! What it costs in time, from `intset_runs` in `benches/intset.rs`:
112//!
113//! ```text
114//!                    4,096     100,000     1,000,000
115//!   contains hit   11.4 ns     13.0 ns       14.9 ns
116//!   contains miss   9.3 ns     10.3 ns       12.7 ns
117//!   member at k     3.5 ns      7.2 ns       10.4 ns
118//!   runs                15         390          3906
119//! ```
120//!
121//! Against the same benchmark before the frame that is 3 to 5 percent slower at
122//! 4,096 and 3 to 12 percent quicker at 100,000 and a million. The slower end is
123//! the subtract the frame adds. The quicker end is the set being half the size
124//! it was, which at 4,096 buys nothing because both fit in cache anyway, and at
125//! a million buys more than the subtract costs.
126//!
127//! End to end through [`crate::Set`], a set of a million integers filled in
128//! scattered order went from 30.71 bytes a member to 2.28, and `SADD` went from
129//! 72.6 ns to 49.8. Membership went the other way, 13.6 ns to about 15, which is
130//! the price of the two searches and is what the memory bought.
131//!
132//! That last number was 40.5 ns at first, which would not have been a trade
133//! worth making, and the fix is the maxima array: picking the run by asking
134//! each one for its own largest member is a pointer chase into a separate heap
135//! buffer at every step of the binary search, and the same search over a
136//! contiguous array of the maxima is not.
137//!
138//! # Why sorted, and what it costs
139//!
140//! Membership is a binary search, which is nine steps at the 512 member ceiling
141//! against the element table's one probe. The reason to accept that is that 512
142//! members is at most four kilobytes, so the search stays in cache and the steps
143//! are not nine cache misses.
144//!
145//! That paragraph used to be an argument with no measurement behind it, which in
146//! this project is a warning sign: L6 put a positional probe at 70 ns and it
147//! measured 13, and K11's crossover does not exist. `benches/intset.rs` settled
148//! it. Minimum per iteration on an M3 laptop, membership against a member that
149//! is there, at the sizes either side of the ceiling:
150//!
151//! ```text
152//!   members     intset     listpack     element table
153//!         8     4.6 ns       6.2 ns            7.7 ns
154//!        64     6.6 ns      29.5 ns           10.2 ns
155//!       128     7.7 ns      60.7 ns           10.2 ns
156//!       512    10.4 ns     239.3 ns            9.0 ns
157//! ```
158//!
159//! So the search is affordable, and the number that makes the case is not the
160//! one against the table. Doubling the set three times costs the intset about
161//! 3 ns in total, which is what a search that stays in cache looks like. What
162//! the intset is actually replacing below the ceiling is the listpack, and there
163//! it is eight times quicker at 128 members and pulling away, because a listpack
164//! walks and this does not.
165//!
166//! The crossover with the element table lands almost exactly on Redis's ceiling.
167//! At 128 the intset wins by a quarter, at 512 the table wins by a seventh. That
168//! is a better outcome than the argument deserved, and it was not predicted here:
169//! the guess was that the search would be affordable, not that the constant Redis
170//! picked in 2011 would sit on the crossover.
171//!
172//! Sorted also means an insert memmoves the tail, and that turns out not to
173//! matter at these sizes. A scattered fill, where every add lands in the middle,
174//! measured 6.47 ns a member at 128 against an ascending fill's 6.46, and the
175//! two only separate at 512 where scattered costs 5.26 against 4.44. Four
176//! kilobytes is not a memmove worth avoiding, which is the whole reason a run is
177//! allowed to be that big. The reason the ascending case is still worth having,
178//! and worth a test, is a shape argument and not a timing one: a fill in
179//! ascending order hits the "greater than the last member" test in front of the
180//! search, so it never searches and never moves anything, and
181//! `an_ascending_fill_never_moves_anything` asserts that rather than timing it.
182//!
183//! # Widening
184//!
185//! Adding a member too wide for the current width rewrites every member of that
186//! run into the new width. That happens at most twice in a run's life, 2 to 4 and
187//! 4 to 8, and the new member is known to sit at one end before the rewrite
188//! starts, because being too wide is exactly what it means to be outside the
189//! range of everything already there. Negative goes to the front and positive to
190//! the back.
191//!
192//! Width is per run and not per set, which is a small win Redis cannot have. A
193//! set holding a million small integers and one huge one keeps every run but the
194//! last at two byte width, where one array would have rewritten all million
195//! members to eight.
196//!
197//! Removing never narrows the width back. Redis does not either, and a set that
198//! narrowed on the way down would rewrite itself on every second operation for a
199//! workload that adds and removes around a boundary.
200
201/// Why an intset from somewhere else was refused.
202#[derive(Debug, Clone, Copy, PartialEq, Eq)]
203pub enum Malformed {
204    /// Shorter than the eight byte header.
205    Short,
206    /// The width is not 2, 4 or 8.
207    Width,
208    /// The count and the width do not account for the bytes that arrived.
209    Length,
210    /// The members are not in ascending order, or one appears twice.
211    Order,
212}
213
214/// The widths a member can be stored in, which are the widths of the three
215/// signed integer types Redis uses and nothing else.
216const W16: u32 = 2;
217const W32: u32 = 4;
218const W64: u32 = 8;
219
220/// Width, then count.
221const HEADER: usize = 8;
222
223/// Members a run holds before it splits in two.
224///
225/// Five hundred and twelve, which is Redis's default `set-max-intset-entries`,
226/// and matching it is deliberate rather than a coincidence. Every set a default
227/// configured server would call an intset is one run here, so it is still one
228/// blob in Redis's own layout, and the runs only start once Redis has given up
229/// on the encoding entirely.
230///
231/// It is also about where the measurements put the ceiling on a memmove that is
232/// still free. A scattered insert into a 512 member run costs 5.26 ns a member
233/// against an ascending one's 4.44, so the tail move is under a nanosecond at
234/// this size and does not need to be smaller.
235const RUN_MAX: usize = 512;
236
237/// Members a run falls to before it is folded into a neighbour.
238///
239/// A quarter of the ceiling. Runs that split leave two halves at half the
240/// ceiling, so the gap between this and that is what stops a set sitting on a
241/// boundary from splitting and merging the same run on alternate writes.
242const RUN_MIN: usize = RUN_MAX / 4;
243
244/// Members a run's buffer grows by when it fills.
245///
246/// A `Vec` grows by doubling, which is right for a buffer whose final size
247/// nobody knows and wrong for one that is never allowed past [`RUN_MAX`]
248/// members. A run one member over a power of two would hold twice the bytes it
249/// needs, and the bytes are the entire point of this representation. Growing in
250/// fixed steps leaves at most this many members of slack whatever the run's
251/// size, which is 64 bytes at two byte width against a run of a few hundred.
252const STEP: usize = 32;
253
254/// A sorted packed set of integers, in one or more runs.
255#[derive(Debug, Clone)]
256pub struct Intset {
257    /// The runs, in ascending order of the values they hold, with disjoint
258    /// ranges. Never empty: a set with no members is one empty run, so that
259    /// every lookup has a run to land in without a special case.
260    runs: Vec<Run>,
261    /// Members across every run.
262    total: usize,
263    /// The largest member of each run, so that picking the run is a search over
264    /// one contiguous array.
265    ///
266    /// It is a copy of something the runs already know, and it earns its eight
267    /// bytes a run several times over. Asking each run for its own largest
268    /// member walks a binary search over a list of pointers into separate heap
269    /// buffers, which is a cache miss a step and was measured at 40.5 ns for a
270    /// membership test on a set of a million. The same search over this array
271    /// touches a few kilobytes that stay in L2.
272    ///
273    /// An empty run holds [`i64::MAX`], so it sorts last and every value lands
274    /// in it, which is what an empty set needs and is the only time a run is
275    /// empty at all.
276    maxima: Vec<i64>,
277    /// Run lengths as a Fenwick tree, one indexed, so that the member at a
278    /// position is found without adding up run lengths. See [`Intset::select`].
279    fen: Vec<u32>,
280}
281
282impl Intset {
283    /// An empty set at the narrowest width.
284    #[must_use]
285    pub fn new() -> Intset {
286        Intset {
287            runs: vec![Run::new()],
288            total: 0,
289            maxima: vec![i64::MAX],
290            fen: vec![0, 0],
291        }
292    }
293
294    /// An empty set with room for `n` members at the narrowest width.
295    ///
296    /// Only a hint. A member that needs a wider slot still widens the run it
297    /// lands in, and the reservation is then short, which costs one growth and
298    /// no correctness.
299    #[must_use]
300    pub fn with_capacity(n: usize) -> Intset {
301        let mut s = Intset::new();
302        s.runs[0].reserve_members(n.min(RUN_MAX));
303        if n > RUN_MAX {
304            // A split leaves two runs of half the ceiling, so that is what the
305            // expected count divides by rather than the ceiling itself.
306            s.runs.reserve(n / (RUN_MAX / 2));
307            s.maxima.reserve(n / (RUN_MAX / 2));
308        }
309        s
310    }
311
312    /// Read a blob written by us or by a real server.
313    ///
314    /// The order check is the one worth having. A truncated blob is caught by
315    /// the length arithmetic, but a blob whose members are out of order reads
316    /// as a perfectly valid set that silently answers no to members it holds,
317    /// because every search here assumes the order.
318    pub fn from_bytes(bytes: &[u8]) -> Result<Intset, Malformed> {
319        let run = Run::from_bytes(bytes)?;
320        let total = run.len();
321        let mut s = Intset {
322            runs: vec![run],
323            total,
324            maxima: Vec::new(),
325            fen: Vec::new(),
326        };
327        s.maxima.push(top_of(&s.runs[0]));
328        s.rebuild_ranks();
329        Ok(s)
330    }
331
332    /// The blob, header included, ready to write to a file, when there is one.
333    ///
334    /// `None` once the set has split, because Redis's format is one array and
335    /// carries nothing that could say otherwise. Nothing is lost by that: the
336    /// split happens at `RUN_MAX` members, which is where a default configured
337    /// server has already stopped storing the set as an intset, so a set that
338    /// could have been written as one still is.
339    #[inline]
340    #[must_use]
341    pub fn as_bytes(&self) -> Option<&[u8]> {
342        match self.runs.as_slice() {
343            // The base check is a backstop rather than a case that comes up on
344            // the way here. A run only takes a base once the set has more than
345            // one of them, and a set drained back down to one run gives its
346            // frame up on the way, so a one run set with a base should not
347            // exist. This is cheaper than being sure of that.
348            [run] if run.base == 0 => Some(run.as_bytes()),
349            _ => None,
350        }
351    }
352
353    /// How many members.
354    #[inline]
355    #[must_use]
356    pub const fn len(&self) -> usize {
357        self.total
358    }
359
360    /// Whether there are none.
361    #[inline]
362    #[must_use]
363    pub const fn is_empty(&self) -> bool {
364        self.total == 0
365    }
366
367    /// How many runs the members are spread over, which is one until the set
368    /// passes `RUN_MAX` members.
369    #[inline]
370    #[must_use]
371    pub fn runs(&self) -> usize {
372        self.runs.len()
373    }
374
375    /// Bytes a member occupies in the widest run, which is 2, 4 or 8.
376    ///
377    /// The widest and not one number for the set, because width is per run here.
378    /// This is what a caller asking "how wide did this set have to get" means,
379    /// and no search uses it.
380    ///
381    /// It is the width of the stored offset and not of the value, so a set of
382    /// integers around a billion reports two once it has split into runs. That
383    /// is the point of the frame of reference the runs are packed against.
384    #[must_use]
385    pub fn width(&self) -> usize {
386        self.runs
387            .iter()
388            .map(Run::width)
389            .max()
390            .unwrap_or(W16 as usize)
391    }
392
393    /// The bytes the members occupy, which is what `MEMORY USAGE` counts.
394    #[inline]
395    #[must_use]
396    pub fn byte_len(&self) -> usize {
397        self.runs.iter().map(Run::byte_len).sum()
398    }
399
400    /// Bytes held, including whatever the vectors have reserved and not used.
401    #[must_use]
402    pub fn memory_bytes(&self) -> usize {
403        let runs: usize = self.runs.iter().map(Run::memory_bytes).sum();
404        runs + self.runs.capacity() * size_of::<Run>()
405            + self.maxima.capacity() * size_of::<i64>()
406            + self.fen.capacity() * size_of::<u32>()
407    }
408
409    /// The member at `index`, counting from the smallest.
410    ///
411    /// # Panics
412    ///
413    /// If `index` is not under [`Intset::len`]. Every caller here has already
414    /// bounded it, and a draw for `SRANDMEMBER` bounds it by construction.
415    #[inline]
416    #[must_use]
417    pub fn at(&self, index: usize) -> i64 {
418        assert!(index < self.total, "index {index} is past the set");
419        let (run, offset) = self.select(index);
420        self.runs[run].at(offset)
421    }
422
423    /// The member at `index`, or `None` past the end.
424    #[inline]
425    #[must_use]
426    pub fn get(&self, index: usize) -> Option<i64> {
427        (index < self.total).then(|| self.at(index))
428    }
429
430    /// The smallest member, or `None` if there are none.
431    #[inline]
432    #[must_use]
433    pub fn min(&self) -> Option<i64> {
434        self.runs.first().and_then(Run::min)
435    }
436
437    /// The largest member, or `None` if there are none.
438    #[inline]
439    #[must_use]
440    pub fn max(&self) -> Option<i64> {
441        self.runs.last().and_then(Run::max)
442    }
443
444    /// Whether `v` is a member.
445    #[inline]
446    #[must_use]
447    pub fn contains(&self, v: i64) -> bool {
448        self.runs[self.run_for(v)].contains(v)
449    }
450
451    /// Every member, smallest first.
452    pub fn iter(&self) -> impl Iterator<Item = i64> + '_ {
453        self.runs.iter().flat_map(Run::iter)
454    }
455
456    /// A cursor on the smallest member, for a merge. See [`Walk`].
457    #[inline]
458    #[must_use]
459    pub fn walk(&self) -> Walk<'_> {
460        Walk::new(self)
461    }
462
463    /// Add `v`. Answers whether it was not already there.
464    pub fn add(&mut self, v: i64) -> bool {
465        let i = self.run_for(v);
466        // A one run set never moves its base, so its bytes stay a Redis intset
467        // and `as_bytes` stays a borrow. See [`Run`].
468        let rebase = self.runs.len() > 1;
469        if !self.runs[i].add(v, rebase) {
470            return false;
471        }
472        self.total += 1;
473        if self.runs[i].len() > RUN_MAX {
474            self.split(i);
475        } else {
476            self.maxima[i] = top_of(&self.runs[i]);
477            self.bump(i, 1);
478        }
479        true
480    }
481
482    /// Remove `v`. Answers whether it was there.
483    pub fn remove(&mut self, v: i64) -> bool {
484        let i = self.run_for(v);
485        if !self.runs[i].remove(v) {
486            return false;
487        }
488        self.total -= 1;
489        self.maxima[i] = top_of(&self.runs[i]);
490        if self.runs.len() > 1 && self.runs[i].len() < RUN_MIN {
491            self.shrink(i);
492        } else {
493            self.bump(i, -1);
494        }
495        true
496    }
497
498    /// Which run holds `v`, or would hold it.
499    ///
500    /// The runs cover disjoint ranges in ascending order, so the first one whose
501    /// largest member is not under `v` is the only one that can hold it. A value
502    /// above every run belongs at the end of the last one, which is what the
503    /// clamp says, and it is the ascending fill: every add lands in the last run
504    /// and appends inside it.
505    #[inline]
506    fn run_for(&self, v: i64) -> usize {
507        let i = self.maxima.partition_point(|&m| m < v);
508        i.min(self.runs.len() - 1)
509    }
510
511    /// Cut run `i` in half, because it has passed [`RUN_MAX`].
512    ///
513    /// The upper half moves into a new run built at whatever width its own
514    /// members need, which is how a set of small integers with one huge member
515    /// keeps most of its runs at two bytes.
516    fn split(&mut self, i: usize) {
517        let n = self.runs[i].len();
518        let half = n / 2;
519        let src = &self.runs[i];
520        // The members are ascending, so the two ends bound the frame of
521        // everything between them and there is nothing to scan.
522        let (base, w) = frame(src.at(half), src.at(n - 1));
523        let mut hi = Run::with_base(base, w, n - half);
524        for k in half..n {
525            hi.push_back(src.at(k));
526        }
527        self.runs[i].truncate(half);
528        // Both halves cover a narrower range than the run they came out of, and
529        // the upper one was built knowing that. This is where the lower one
530        // finds out, and it is the whole reason a split is where the frames get
531        // tight: an ascending fill splits every run exactly once.
532        self.runs[i].rebase();
533        // The lower half keeps the buffer the whole run had, which is twice
534        // what it now holds, and an ascending fill splits every run exactly
535        // once and then never touches the lower half again. Left alone that is
536        // two bytes a member of pure slack on the commonest fill there is, so
537        // the buffer is handed back here and the next insert into it reserves a
538        // step like any other.
539        self.runs[i].tighten();
540        self.runs.insert(i + 1, hi);
541        // The maxima are patched rather than recomputed. Recomputing them means
542        // asking every run for its own last member, which is a pointer chase
543        // per run into a separate heap buffer, and it measured at a sixth of
544        // the cost of a whole scattered fill of a million members. Moving eight
545        // bytes a run along an array is nothing next to that.
546        self.maxima[i] = top_of(&self.runs[i]);
547        let top = top_of(&self.runs[i + 1]);
548        self.maxima.insert(i + 1, top);
549        self.rebuild_ranks();
550    }
551
552    /// Fold run `i` into a neighbour, because it has fallen under [`RUN_MIN`].
553    ///
554    /// An empty run simply goes. Otherwise it merges with whichever neighbour
555    /// the two of them fit inside one run, preferring the one on the left so
556    /// that a set being drained from the front collapses rather than leaving a
557    /// trail of short runs. Two neighbours that are both too full to take it is
558    /// not a problem to solve: the run stays short and costs one entry in the
559    /// tree.
560    fn shrink(&mut self, i: usize) {
561        if self.runs[i].is_empty() {
562            self.runs.remove(i);
563            self.maxima.remove(i);
564            self.rebuild_ranks();
565            return;
566        }
567        let fits = |a: usize, b: usize| self.runs[a].len() + self.runs[b].len() <= RUN_MAX;
568        let (lo, hi) = if i > 0 && fits(i - 1, i) {
569            (i - 1, i)
570        } else if i + 1 < self.runs.len() && fits(i, i + 1) {
571            (i, i + 1)
572        } else {
573            self.bump(i, -1);
574            return;
575        };
576        let src = self.runs.remove(hi);
577        self.maxima.remove(hi);
578        self.runs[lo].append(&src);
579        // A set drained back down to one run gives up its frame, so that it is
580        // a Redis intset again and [`Intset::as_bytes`] is a borrow rather than
581        // a rebuild. See [`Run`] for why one run never carries a base.
582        if self.runs.len() == 1 {
583            self.runs[0].unframe();
584        }
585        self.maxima[lo] = top_of(&self.runs[lo]);
586        self.rebuild_ranks();
587    }
588
589    /// Which run holds position `k`, and how far into it.
590    ///
591    /// The standard Fenwick descent: walk the powers of two downward, taking a
592    /// step whenever the members it covers are all still behind `k`. What is
593    /// left over when the steps run out is the offset inside the run.
594    fn select(&self, k: usize) -> (usize, usize) {
595        let n = self.runs.len();
596        let mut pos = 0usize;
597        let mut rem = k;
598        let mut step = 1usize << (usize::BITS - 1 - n.leading_zeros());
599        while step > 0 {
600            let next = pos + step;
601            if next <= n {
602                let covered = self.fen[next] as usize;
603                if covered <= rem {
604                    pos = next;
605                    rem -= covered;
606                }
607            }
608            step >>= 1;
609        }
610        (pos, rem)
611    }
612
613    /// Tell the tree that run `i` gained or lost one member.
614    fn bump(&mut self, i: usize, delta: i32) {
615        let n = self.runs.len();
616        let mut at = i + 1;
617        while at <= n {
618            if delta > 0 {
619                self.fen[at] += 1;
620            } else {
621                self.fen[at] -= 1;
622            }
623            at += at & at.wrapping_neg();
624        }
625    }
626
627    /// Rebuild the tree from the run lengths.
628    ///
629    /// What a split or a merge needs, because both of them renumber every run
630    /// after the one they touched and a Fenwick tree is not a thing you patch
631    /// in the middle. It is linear in the number of runs, over one array that
632    /// is read and written straight through, and it happens once per couple of
633    /// hundred writes.
634    fn rebuild_ranks(&mut self) {
635        let n = self.runs.len();
636        self.fen.clear();
637        self.fen.resize(n + 1, 0);
638        for i in 1..=n {
639            let len = u32::try_from(self.runs[i - 1].len()).expect("a run is under RUN_MAX");
640            self.fen[i] += len;
641            let parent = i + (i & i.wrapping_neg());
642            if parent <= n {
643                let carry = self.fen[i];
644                self.fen[parent] += carry;
645            }
646        }
647    }
648}
649
650impl Default for Intset {
651    fn default() -> Intset {
652        Intset::new()
653    }
654}
655
656/// A cursor over the members that only ever moves forward.
657///
658/// [`Intset::iter`] is enough to read a set out, and it is not enough to merge
659/// two of them, because a merge needs to skip. Intersecting a set of ten with a
660/// set of a million should touch ten members of the big one and not a million,
661/// and that is [`Walk::seek`], which jumps to the first member at or past a
662/// value instead of stepping to it.
663///
664/// Forward only, and that is the whole reason it is worth having. A cursor that
665/// could go backwards would have to binary search the entire set on every seek.
666/// This one searches from where it already is, so a merge that walks two sets in
667/// lockstep pays one comparison a member in the common case and only searches
668/// when it actually skipped something.
669///
670/// Stepping is a pointer step and nothing else, which is what makes a merge a
671/// different order of cost from a probe. `setops.rs` explains what that buys and
672/// has the numbers.
673#[derive(Debug, Clone, Copy)]
674pub struct Walk<'a> {
675    set: &'a Intset,
676    /// Which run. Equal to the run count once the cursor is past the end.
677    run: usize,
678    /// How far into that run. Always under the run's length except when the
679    /// cursor is past the end, where the pair is `(runs.len(), 0)`.
680    off: usize,
681}
682
683impl<'a> Walk<'a> {
684    /// A cursor on the smallest member.
685    fn new(set: &'a Intset) -> Walk<'a> {
686        let mut w = Walk {
687            set,
688            run: 0,
689            off: 0,
690        };
691        w.settle();
692        w
693    }
694
695    /// The member the cursor is on, or `None` past the end.
696    #[inline]
697    #[must_use]
698    pub fn peek(&self) -> Option<i64> {
699        (self.run < self.set.runs.len()).then(|| self.set.runs[self.run].at(self.off))
700    }
701
702    /// Move to the next member.
703    #[inline]
704    pub fn bump(&mut self) {
705        self.off += 1;
706        self.settle();
707    }
708
709    /// Move to the first member that is not under `v`, without going backwards.
710    ///
711    /// A seek to a value the cursor is already at or past does nothing, which is
712    /// what makes this safe to call in a loop that does not know whether it has
713    /// moved.
714    pub fn seek(&mut self, v: i64) {
715        match self.peek() {
716            Some(cur) if cur < v => {}
717            // Already there, or there is nothing left to seek to.
718            _ => return,
719        }
720        // The runs hold disjoint ranges in ascending order, so the run is the
721        // first one at or after this one whose largest member is not under `v`.
722        // Searching from `run + 1` rather than from the start is what keeps a
723        // seek near the cursor cheap: a merge that steps through both sets
724        // together never leaves its current run.
725        if self.set.maxima[self.run] < v {
726            let after = &self.set.maxima[self.run + 1..];
727            let hop = after.partition_point(|&m| m < v);
728            self.run += 1 + hop;
729            if self.run >= self.set.runs.len() {
730                self.run = self.set.runs.len();
731                self.off = 0;
732                return;
733            }
734            self.off = 0;
735        }
736        self.off = self.set.runs[self.run].lower_bound(v, self.off);
737        self.settle();
738    }
739
740    /// Step off the end of a run onto the next one.
741    ///
742    /// A loop rather than a test because an empty set is one empty run, and that
743    /// is the only time two runs in a row have nothing to land on.
744    #[inline]
745    fn settle(&mut self) {
746        while self.run < self.set.runs.len() && self.off >= self.set.runs[self.run].len() {
747            self.run += 1;
748            self.off = 0;
749        }
750    }
751}
752
753/// Two sets are equal when they hold the same members.
754///
755/// Written out rather than derived, because where the run boundaries fell is an
756/// artefact of the order the members arrived in and not something a caller has
757/// any business seeing. A set filled ascending and the same set filled scattered
758/// are the same set.
759impl PartialEq for Intset {
760    fn eq(&self, other: &Intset) -> bool {
761        self.total == other.total && self.iter().eq(other.iter())
762    }
763}
764
765impl Eq for Intset {}
766
767/// One run: a complete intset in Redis's own layout, offset from a base.
768///
769/// The base is what takes a run of large integers down to two bytes a member.
770/// A run holds at most [`RUN_MAX`] members out of a set that may hold millions,
771/// so the values inside one run are close together whatever the set as a whole
772/// spans: a million members scattered over sixteen million values leave every
773/// run covering a few thousand of them. Stored as themselves those need four
774/// bytes each, and stored as their distance from the middle of the run's own
775/// range they need two.
776///
777/// The base is the middle of that range rather than the bottom of it, which is
778/// worth a sentence because it is not obvious. The stored offsets are read back
779/// through the same signed readers Redis uses, so a base at the bottom would
780/// only ever use the positive half of the width and hold a span of thirty two
781/// thousand at two bytes. Centred, the offsets run either side of zero and the
782/// same two bytes hold a span of sixty five thousand. It also leaves room on
783/// both sides for the members still to arrive rather than only above.
784///
785/// A base of zero is a run that is byte for byte a Redis intset, and a run only
786/// takes a base once the set it belongs to has more than one of them. That is
787/// deliberate: a set a default configured server would still call an intset
788/// stays one array in Redis's own layout, so [`Intset::as_bytes`] is still a
789/// borrow rather than a rebuild, and the frame only appears past the point
790/// where Redis has stopped having an intset at all.
791#[derive(Debug, Clone, PartialEq, Eq)]
792struct Run {
793    /// The header and the members, in Redis's own layout, so that handing this
794    /// to an RDB writer is a copy when the base is zero.
795    bytes: Vec<u8>,
796    /// What every stored member is measured from.
797    base: i64,
798}
799
800impl Run {
801    /// An empty run at the narrowest width.
802    fn new() -> Run {
803        Run::with_base(0, W16 as usize, 0)
804    }
805
806    /// An empty run against `base`, `w` bytes a member, with room for `n`.
807    fn with_base(base: i64, w: usize, n: usize) -> Run {
808        let mut bytes = Vec::with_capacity(HEADER + n * w);
809        bytes.extend_from_slice(&(w as u32).to_le_bytes());
810        bytes.extend_from_slice(&0u32.to_le_bytes());
811        Run { bytes, base }
812    }
813
814    /// Read a blob written by us or by a real server. See [`Intset::from_bytes`].
815    fn from_bytes(bytes: &[u8]) -> Result<Run, Malformed> {
816        if bytes.len() < HEADER {
817            return Err(Malformed::Short);
818        }
819        let width = u32::from_le_bytes(bytes[0..4].try_into().expect("four bytes"));
820        if width != W16 && width != W32 && width != W64 {
821            return Err(Malformed::Width);
822        }
823        let count = u32::from_le_bytes(bytes[4..8].try_into().expect("four bytes")) as usize;
824        let want = count
825            .checked_mul(width as usize)
826            .and_then(|n| n.checked_add(HEADER))
827            .ok_or(Malformed::Length)?;
828        if bytes.len() != want {
829            return Err(Malformed::Length);
830        }
831        let s = Run {
832            bytes: bytes.to_vec(),
833            base: 0,
834        };
835        for i in 1..count {
836            if s.at(i - 1) >= s.at(i) {
837                return Err(Malformed::Order);
838            }
839        }
840        Ok(s)
841    }
842
843    /// The blob, header included.
844    #[inline]
845    fn as_bytes(&self) -> &[u8] {
846        &self.bytes
847    }
848
849    /// How many members.
850    #[inline]
851    fn len(&self) -> usize {
852        u32::from_le_bytes(self.bytes[4..8].try_into().expect("four bytes")) as usize
853    }
854
855    /// Whether there are none.
856    #[inline]
857    fn is_empty(&self) -> bool {
858        self.len() == 0
859    }
860
861    /// Bytes a member occupies, which is 2, 4 or 8.
862    #[inline]
863    fn width(&self) -> usize {
864        u32::from_le_bytes(self.bytes[0..4].try_into().expect("four bytes")) as usize
865    }
866
867    /// The blob's length.
868    #[inline]
869    fn byte_len(&self) -> usize {
870        self.bytes.len()
871    }
872
873    /// Bytes held, including whatever the vector has reserved and not used.
874    #[inline]
875    fn memory_bytes(&self) -> usize {
876        self.bytes.capacity()
877    }
878
879    /// The member at `index`, counting from the smallest.
880    #[inline]
881    fn at(&self, index: usize) -> i64 {
882        self.base + self.raw(index)
883    }
884
885    /// What is stored at `index`, which is the member less the base.
886    #[inline]
887    fn raw(&self, index: usize) -> i64 {
888        self.raw_w(index, self.width())
889    }
890
891    /// [`Run::raw`] for a caller that already knows the width.
892    ///
893    /// The width lives in the buffer, so reading it is a load, and a binary
894    /// search that reads it at every step reads the same four bytes nine times.
895    /// Out here it is read once and the search compares stored offsets against
896    /// a stored offset rather than adding the base back nine times.
897    #[inline]
898    fn raw_w(&self, index: usize, w: usize) -> i64 {
899        let at = HEADER + index * w;
900        let raw = &self.bytes[at..at + w];
901        match w {
902            2 => i64::from(i16::from_le_bytes(raw.try_into().expect("two bytes"))),
903            4 => i64::from(i32::from_le_bytes(raw.try_into().expect("four bytes"))),
904            _ => i64::from_le_bytes(raw.try_into().expect("eight bytes")),
905        }
906    }
907
908    /// Whether `v` is inside the frame this run is packed against.
909    ///
910    /// A value outside it is not a member, because every member is inside it,
911    /// and saying so costs a subtract and a compare instead of a search.
912    #[inline]
913    fn framed(&self, v: i64) -> bool {
914        v.checked_sub(self.base)
915            .is_some_and(|off| width_of(off) as usize <= self.width())
916    }
917
918    /// The smallest member, or `None` if there are none.
919    #[inline]
920    fn min(&self) -> Option<i64> {
921        (!self.is_empty()).then(|| self.at(0))
922    }
923
924    /// The largest member, or `None` if there are none.
925    #[inline]
926    fn max(&self) -> Option<i64> {
927        self.len().checked_sub(1).map(|last| self.at(last))
928    }
929
930    /// Whether `v` is a member.
931    #[inline]
932    fn contains(&self, v: i64) -> bool {
933        self.framed(v) && self.search(v).is_ok()
934    }
935
936    /// Every member, smallest first.
937    fn iter(&self) -> impl Iterator<Item = i64> + '_ {
938        (0..self.len()).map(|i| self.at(i))
939    }
940
941    /// Room for `n` more members without a growth.
942    fn reserve_members(&mut self, n: usize) {
943        self.bytes.reserve_exact(n * self.width());
944    }
945
946    /// Add `v`. Answers whether it was not already there.
947    ///
948    /// `rebase` is whether this run is allowed to move its base, which it is
949    /// only once the set has more than one run. See [`Run`].
950    fn add(&mut self, v: i64, rebase: bool) -> bool {
951        if !self.framed(v) {
952            self.refit_and_add(v, rebase);
953            return true;
954        }
955        match self.search(v) {
956            Ok(_) => false,
957            Err(at) => {
958                self.insert_at(at, v);
959                true
960            }
961        }
962    }
963
964    /// Put `v` on the end, where it is already known to belong.
965    ///
966    /// Only used when one run is being built out of another, so the caller has
967    /// the members in ascending order and the width is already wide enough.
968    fn push_back(&mut self, v: i64) {
969        let w = self.width();
970        let at = self.bytes.len();
971        self.grow_by(w);
972        write_at(&mut self.bytes, at, w, v - self.base);
973        self.set_len(self.len() + 1);
974    }
975
976    /// Put every member of `other` on the end, where they all belong.
977    fn append(&mut self, other: &Run) {
978        self.reserve_members(other.len());
979        for v in other.iter() {
980            // Through `add` and not `push_back`, because `other` may hold
981            // members outside this run's frame and refitting is `add`'s job.
982            // Every one of them is past the last member, so the range test in
983            // front of the search answers and nothing moves.
984            self.add(v, true);
985        }
986    }
987
988    /// Drop everything from `n` onward.
989    fn truncate(&mut self, n: usize) {
990        self.bytes.truncate(HEADER + n * self.width());
991        self.set_len(n);
992    }
993
994    /// Hand back whatever the buffer is holding and not using.
995    fn tighten(&mut self) {
996        self.bytes.shrink_to_fit();
997    }
998
999    /// Remove `v`. Answers whether it was there.
1000    fn remove(&mut self, v: i64) -> bool {
1001        if !self.framed(v) {
1002            return false;
1003        }
1004        let Ok(at) = self.search(v) else {
1005            return false;
1006        };
1007        let w = self.width();
1008        let from = HEADER + at * w;
1009        self.bytes.drain(from..from + w);
1010        self.set_len(self.len() - 1);
1011        true
1012    }
1013
1014    /// The first member at or past `v`, searching only from `from`.
1015    ///
1016    /// [`Walk::seek`]'s inner half, and the reason it takes a lower bound rather
1017    /// than reusing [`Run::search`]: a cursor never goes backwards, so
1018    /// everything before where it already is has been ruled out and searching it
1019    /// again is work with a known answer. On a merge that steps through two sets
1020    /// together the range is one or two members wide.
1021    fn lower_bound(&self, v: i64, from: usize) -> usize {
1022        let w = self.width();
1023        let off = self.offset_of(v);
1024        let (mut lo, mut hi) = (from, self.len());
1025        while lo < hi {
1026            let mid = lo.midpoint(hi);
1027            if self.raw_w(mid, w) < off {
1028                lo = mid + 1;
1029            } else {
1030                hi = mid;
1031            }
1032        }
1033        lo
1034    }
1035
1036    /// `v` as this run would store it, saturating rather than wrapping.
1037    ///
1038    /// A search compares stored offsets against a stored offset, so the base
1039    /// comes off the value it is looking for once instead of going back onto
1040    /// every member the search touches. Saturating is right here and not just
1041    /// convenient: a value too far from the base to subtract at all is a value
1042    /// past every member on that side, and the saturated offset is past every
1043    /// stored offset on that side too, so the search lands where it should.
1044    #[inline]
1045    fn offset_of(&self, v: i64) -> i64 {
1046        v.checked_sub(self.base)
1047            .unwrap_or(if v < self.base { i64::MIN } else { i64::MAX })
1048    }
1049
1050    /// Where `v` is, or where it would go.
1051    ///
1052    /// The two range tests in front of the binary search are Redis's and they
1053    /// are not an optimisation of the search, they are what makes an ascending
1054    /// fill linear: every add lands past the last member, answers in two loads,
1055    /// and appends with nothing to move.
1056    fn search(&self, v: i64) -> Result<usize, usize> {
1057        let n = self.len();
1058        if n == 0 {
1059            return Err(0);
1060        }
1061        let w = self.width();
1062        let off = self.offset_of(v);
1063        if off > self.raw_w(n - 1, w) {
1064            return Err(n);
1065        }
1066        if off < self.raw_w(0, w) {
1067            return Err(0);
1068        }
1069        let (mut lo, mut hi) = (0usize, n - 1);
1070        while lo <= hi {
1071            let mid = lo.midpoint(hi);
1072            let cur = self.raw_w(mid, w);
1073            if off > cur {
1074                lo = mid + 1;
1075            } else if off < cur {
1076                // `mid` is at least one here, because `v` is not under the
1077                // first member and so cannot be under member zero.
1078                hi = mid - 1;
1079            } else {
1080                return Ok(mid);
1081            }
1082        }
1083        Err(lo)
1084    }
1085
1086    /// Rewrite every member against a frame that holds `v` too, and add `v`.
1087    ///
1088    /// `v` is outside the frame everything here is packed against, so it is not
1089    /// a member and it is not in the middle: it is under the smallest or over
1090    /// the largest, and either way there is no search.
1091    ///
1092    /// The members go out to a fixed buffer and come back rather than being
1093    /// shuffled where they lie. The old code could move them in place because a
1094    /// widen only ever moved a member to a higher offset, so back to front was
1095    /// safe. A refit can narrow as well as widen, and it can shift the members
1096    /// up by one at the same time, and there is no single direction that is
1097    /// safe for all of those. A run is capped at [`RUN_MAX`] members so the
1098    /// buffer is a known four kilobytes, and this runs once per couple of
1099    /// hundred inserts and never on the ascending fill.
1100    fn refit_and_add(&mut self, v: i64, rebase: bool) {
1101        let n = self.len();
1102        let mut held = [0i64; RUN_MAX + 2];
1103        for (i, slot) in held.iter_mut().enumerate().take(n) {
1104            *slot = self.at(i);
1105        }
1106        let ahead = usize::from(n > 0 && v < held[0]);
1107        if ahead == 1 {
1108            held.copy_within(0..n, 1);
1109        }
1110        held[if ahead == 1 { 0 } else { n }] = v;
1111        self.repack(&held[..n + 1], rebase);
1112    }
1113
1114    /// Repack against the tightest frame for the members that are here.
1115    ///
1116    /// Called after a split, where both halves cover a narrower range than the
1117    /// run they came out of and neither of them knows it yet.
1118    fn rebase(&mut self) {
1119        let n = self.len();
1120        let mut held = [0i64; RUN_MAX + 2];
1121        for (i, slot) in held.iter_mut().enumerate().take(n) {
1122            *slot = self.at(i);
1123        }
1124        self.repack(&held[..n], true);
1125    }
1126
1127    /// Give up the frame and store the members as themselves.
1128    ///
1129    /// Called when a set shrinks back to one run, which is the one shape that
1130    /// has to stay byte for byte a Redis intset. It costs whatever the wider
1131    /// width costs, on a set small enough that the difference is a few hundred
1132    /// bytes, and it buys back a borrow on every save.
1133    fn unframe(&mut self) {
1134        if self.base == 0 {
1135            return;
1136        }
1137        let n = self.len();
1138        let mut held = [0i64; RUN_MAX + 2];
1139        for (i, slot) in held.iter_mut().enumerate().take(n) {
1140            *slot = self.at(i);
1141        }
1142        self.repack(&held[..n], false);
1143    }
1144
1145    /// Write `members` out against the tightest frame that holds them.
1146    fn repack(&mut self, members: &[i64], rebase: bool) {
1147        let (base, w) = match members {
1148            [] => (0, W16 as usize),
1149            [only] => (if rebase { *only } else { 0 }, {
1150                let off = if rebase { 0 } else { *only };
1151                width_of(off) as usize
1152            }),
1153            [lo, .., hi] if rebase => frame(*lo, *hi),
1154            [lo, .., hi] => (0, width_of(*lo).max(width_of(*hi)) as usize),
1155        };
1156        self.base = base;
1157        self.bytes.resize(HEADER + members.len() * w, 0);
1158        self.bytes[0..4].copy_from_slice(&(w as u32).to_le_bytes());
1159        for (i, &v) in members.iter().enumerate() {
1160            write_at(&mut self.bytes, HEADER + i * w, w, v - base);
1161        }
1162        self.set_len(members.len());
1163    }
1164
1165    /// Open a slot at `at` and put `v` in it.
1166    fn insert_at(&mut self, at: usize, v: i64) {
1167        let w = self.width();
1168        let from = HEADER + at * w;
1169        let old = self.bytes.len();
1170        self.grow_by(w);
1171        self.bytes.copy_within(from..old, from + w);
1172        write_at(&mut self.bytes, from, w, v - self.base);
1173        self.set_len(self.len() + 1);
1174    }
1175
1176    /// Make the blob `w` bytes longer without letting the vector double.
1177    ///
1178    /// [`STEP`] says why this is not just a `resize`. The reserve is skipped
1179    /// when the capacity already covers it, so a run built with room for what is
1180    /// about to go in it never calls the allocator at all.
1181    #[inline]
1182    fn grow_by(&mut self, w: usize) {
1183        let want = self.bytes.len() + w;
1184        if want > self.bytes.capacity() {
1185            // `yo_alloc::for_the_data` and not a fix. A run that has taken its
1186            // ten thousandth member has grown along the way, and this is the
1187            // only place in the intset that ever asks the allocator for
1188            // anything.
1189            yo_alloc::for_the_data(|| self.bytes.reserve_exact(STEP * w));
1190        }
1191        self.bytes.resize(want, 0);
1192    }
1193
1194    #[inline]
1195    fn set_len(&mut self, n: usize) {
1196        let n = u32::try_from(n).expect("a run never reaches four billion members");
1197        self.bytes[4..8].copy_from_slice(&n.to_le_bytes());
1198    }
1199}
1200
1201/// A run's largest member, or [`i64::MAX`] when it has none.
1202///
1203/// The sentinel is what puts an empty run last in the maxima and sends every
1204/// value into it, which is the empty set and nothing else.
1205#[inline]
1206fn top_of(r: &Run) -> i64 {
1207    r.max().unwrap_or(i64::MAX)
1208}
1209
1210/// The base and width that hold every value from `lo` to `hi`.
1211///
1212/// The base is the middle of the range and not the bottom of it, because the
1213/// offsets are read back through signed readers: from the bottom they would only
1214/// use the positive half of the width and two bytes would hold a span of thirty
1215/// two thousand, and from the middle they use both halves and two bytes hold
1216/// sixty five thousand.
1217///
1218/// A range too wide to subtract at all is a run holding both ends of the
1219/// sixty four bit line, which no frame helps with, so it gets no base and the
1220/// widest width.
1221#[inline]
1222fn frame(lo: i64, hi: i64) -> (i64, usize) {
1223    let Some(span) = hi.checked_sub(lo) else {
1224        return (0, W64 as usize);
1225    };
1226    let base = lo + span / 2;
1227    let w = width_of(lo - base).max(width_of(hi - base));
1228    (base, w as usize)
1229}
1230
1231/// The narrowest width that holds `v`.
1232#[inline]
1233const fn width_of(v: i64) -> u32 {
1234    if v < i32::MIN as i64 || v > i32::MAX as i64 {
1235        W64
1236    } else if v < i16::MIN as i64 || v > i16::MAX as i64 {
1237        W32
1238    } else {
1239        W16
1240    }
1241}
1242
1243/// Write `v` at `at` in `w` bytes, little endian.
1244#[inline]
1245fn write_at(bytes: &mut [u8], at: usize, w: usize, v: i64) {
1246    match w {
1247        2 => bytes[at..at + 2].copy_from_slice(&(v as i16).to_le_bytes()),
1248        4 => bytes[at..at + 4].copy_from_slice(&(v as i32).to_le_bytes()),
1249        _ => bytes[at..at + 8].copy_from_slice(&v.to_le_bytes()),
1250    }
1251}
1252
1253#[cfg(test)]
1254mod tests {
1255    use super::*;
1256
1257    fn of(vals: &[i64]) -> Intset {
1258        let mut s = Intset::new();
1259        for &v in vals {
1260            assert!(s.add(v), "{v} was supposed to be new");
1261        }
1262        s
1263    }
1264
1265    fn members(s: &Intset) -> Vec<i64> {
1266        s.iter().collect()
1267    }
1268
1269    /// The frame is the whole memory argument, so this is the row that says it
1270    /// worked. A billion apart is far outside two byte range and the set still
1271    /// stores its members in two bytes each, because no one run spans more than
1272    /// a few thousand of them.
1273    #[test]
1274    fn a_set_of_large_integers_still_stores_them_in_two_bytes() {
1275        // Still several runs at the Miri size, which is what the two byte width
1276        // has to survive here.
1277        let n: i64 = if cfg!(miri) { 3_000 } else { 10_000 };
1278        let mut s = Intset::new();
1279        for i in 0..n {
1280            s.add(1_000_000_000 + i * 3);
1281        }
1282        assert_eq!(s.width(), W16 as usize, "every run is two bytes a member");
1283        assert!(
1284            s.byte_len() < n as usize * 2 + s.runs() * 16,
1285            "{} bytes for {n} members over {} runs",
1286            s.byte_len(),
1287            s.runs()
1288        );
1289        for i in 0..n {
1290            assert!(s.contains(1_000_000_000 + i * 3), "member {i}");
1291            assert!(!s.contains(1_000_000_000 + i * 3 + 1), "gap after {i}");
1292        }
1293        assert_eq!(s.len(), n as usize);
1294    }
1295
1296    /// A member arriving under a run's smallest moves the frame down rather
1297    /// than widening it, which is the direction the old widen path never had to
1298    /// think about.
1299    #[test]
1300    fn a_member_under_the_frame_moves_it_instead_of_widening_it() {
1301        let mut s = Intset::new();
1302        // Past the ceiling, so the runs and the frames exist at all.
1303        for i in 0..2_000i64 {
1304            s.add(500_000 + i * 100);
1305        }
1306        let before = s.width();
1307        for i in 0..50i64 {
1308            assert!(s.add(500_000 - 1 - i), "{i} is new and under everything");
1309        }
1310        assert_eq!(s.width(), before, "still two bytes a member");
1311        assert_eq!(s.min(), Some(500_000 - 50));
1312        assert_eq!(s.len(), 2_050);
1313        for i in 0..50i64 {
1314            assert!(s.contains(500_000 - 1 - i));
1315        }
1316    }
1317
1318    /// Negative members, which is where a centred base and a signed reader
1319    /// could disagree with each other and nothing else would notice.
1320    #[test]
1321    fn the_frame_holds_negative_members_too() {
1322        let n: i64 = if cfg!(miri) { 1_500 } else { 3_000 };
1323        let mut s = Intset::new();
1324        for i in 0..n {
1325            s.add(-2_000_000_000 + i * 7);
1326        }
1327        assert_eq!(s.width(), W16 as usize);
1328        assert_eq!(s.min(), Some(-2_000_000_000));
1329        assert_eq!(s.max(), Some(-2_000_000_000 + (n - 1) * 7));
1330        for i in 0..n {
1331            assert!(s.contains(-2_000_000_000 + i * 7), "member {i}");
1332        }
1333        assert_eq!(members(&s).len(), n as usize);
1334    }
1335
1336    /// A run holding both ends of the sixty four bit line, which no frame helps
1337    /// with and which the subtraction cannot even be done on.
1338    #[test]
1339    fn a_span_too_wide_to_subtract_gets_no_frame() {
1340        assert_eq!(frame(i64::MIN, i64::MAX), (0, W64 as usize));
1341        let mut s = Intset::new();
1342        for i in 0..600i64 {
1343            s.add(i);
1344        }
1345        s.add(i64::MIN);
1346        s.add(i64::MAX);
1347        assert_eq!(s.width(), W64 as usize, "the widest run holds both ends");
1348        assert!(s.contains(i64::MIN) && s.contains(i64::MAX) && s.contains(300));
1349        assert_eq!(s.len(), 602);
1350    }
1351
1352    /// A set small enough for a real server to call it an intset hands over the
1353    /// same bytes it always did, whatever its members are, because a one run set
1354    /// never takes a base.
1355    #[test]
1356    fn a_one_run_set_is_still_a_redis_intset() {
1357        let s = of(&[1_000_000_000, 1_000_000_001, 2_000_000_000]);
1358        let bytes = s.as_bytes().expect("one run");
1359        assert_eq!(
1360            Intset::from_bytes(bytes).expect("a real server could read this"),
1361            s
1362        );
1363        assert_eq!(s.width(), W32 as usize, "no base, so the values decide");
1364    }
1365
1366    #[test]
1367    fn a_set_drained_back_to_one_run_is_a_redis_intset_again() {
1368        // Enough to be several runs either way, since one run is the thing this
1369        // has to drain back down to.
1370        let n: i64 = if cfg!(miri) { 1_500 } else { 4_000 };
1371        let mut s = Intset::new();
1372        for i in 0..n {
1373            s.add(1_000_000_000 + i * 3);
1374        }
1375        assert!(s.runs.len() > 1, "several runs to start with");
1376        assert!(s.as_bytes().is_none(), "framed, so not a Redis intset");
1377        for i in 100..n {
1378            s.remove(1_000_000_000 + i * 3);
1379        }
1380        sound(&s);
1381        assert_eq!(s.runs.len(), 1, "the merges took it back to one run");
1382        let bytes = s.as_bytes().expect("one run, so the frame is gone");
1383        assert_eq!(
1384            Intset::from_bytes(bytes).expect("a real server could read this"),
1385            s
1386        );
1387        assert_eq!(s.len(), 100);
1388    }
1389
1390    /// Everything the runs have to keep true, checked in one place so that a
1391    /// test only has to call this rather than remember all four.
1392    fn sound(s: &Intset) {
1393        assert!(!s.runs.is_empty(), "there is always a run to land in");
1394        assert_eq!(s.maxima.len(), s.runs.len(), "one maximum per run");
1395        let mut seen = 0usize;
1396        let mut last: Option<i64> = None;
1397        for (i, r) in s.runs.iter().enumerate() {
1398            assert!(
1399                !r.is_empty() || s.runs.len() == 1,
1400                "run {i} is empty and is not the only one"
1401            );
1402            assert!(r.len() <= RUN_MAX, "run {i} holds {} members", r.len());
1403            // A stale maximum is the one thing that sends a lookup to the wrong
1404            // run, and it fails silently: the member is simply not found.
1405            assert_eq!(s.maxima[i], top_of(r), "the maximum of run {i} is stale");
1406            for v in r.iter() {
1407                if let Some(prev) = last {
1408                    assert!(prev < v, "{prev} then {v} is not ascending");
1409                }
1410                last = Some(v);
1411            }
1412            seen += r.len();
1413        }
1414        assert_eq!(seen, s.len(), "the runs and the count disagree");
1415        // The tree has to agree with the runs at every position, which is the
1416        // one thing a wrong `bump` breaks silently.
1417        let mut at = 0usize;
1418        for (i, r) in s.runs.iter().enumerate() {
1419            for k in 0..r.len() {
1420                assert_eq!(s.select(at), (i, k), "position {at}");
1421                at += 1;
1422            }
1423        }
1424    }
1425
1426    #[test]
1427    fn an_empty_set_is_eight_bytes_and_holds_nothing() {
1428        let s = Intset::new();
1429        assert_eq!(s.len(), 0);
1430        assert!(s.is_empty());
1431        assert_eq!(s.width(), 2);
1432        assert_eq!(s.byte_len(), 8);
1433        assert_eq!(s.min(), None);
1434        assert_eq!(s.max(), None);
1435        assert!(!s.contains(0));
1436        assert_eq!(s.as_bytes(), Some(&[2, 0, 0, 0, 0, 0, 0, 0][..]));
1437    }
1438
1439    #[test]
1440    fn members_come_back_sorted_however_they_went_in() {
1441        let s = of(&[5, -3, 100, 0, -70, 42]);
1442        assert_eq!(members(&s), [-70, -3, 0, 5, 42, 100]);
1443        assert_eq!(s.min(), Some(-70));
1444        assert_eq!(s.max(), Some(100));
1445        assert_eq!(s.len(), 6);
1446        sound(&s);
1447    }
1448
1449    #[test]
1450    fn adding_the_same_member_twice_says_so_and_changes_nothing() {
1451        let mut s = of(&[1, 2, 3]);
1452        assert!(!s.add(2));
1453        assert_eq!(members(&s), [1, 2, 3]);
1454        assert_eq!(s.byte_len(), 8 + 3 * 2);
1455    }
1456
1457    #[test]
1458    fn a_small_set_of_integers_costs_two_bytes_each() {
1459        // G8's number for a set of integers, and the reason this representation
1460        // exists next to the listpack rather than instead of it.
1461        let s = of(&(0..512).collect::<Vec<i64>>());
1462        assert_eq!(s.runs(), 1, "512 is still one run");
1463        assert_eq!(s.width(), 2);
1464        assert_eq!(s.byte_len(), 8 + 512 * 2);
1465        assert_eq!((s.byte_len() - 8) / s.len(), 2);
1466    }
1467
1468    #[test]
1469    fn the_width_follows_the_widest_member_and_never_comes_back_down() {
1470        let mut s = of(&[1, 2, 3]);
1471        assert_eq!(s.width(), 2);
1472
1473        s.add(100_000);
1474        assert_eq!(s.width(), 4, "past an i16");
1475        assert_eq!(members(&s), [1, 2, 3, 100_000]);
1476
1477        s.add(-5_000_000_000);
1478        assert_eq!(s.width(), 8, "past an i32");
1479        assert_eq!(members(&s), [-5_000_000_000, 1, 2, 3, 100_000]);
1480
1481        assert!(s.remove(-5_000_000_000));
1482        assert!(s.remove(100_000));
1483        assert_eq!(s.width(), 8, "removing does not narrow it back");
1484        assert_eq!(members(&s), [1, 2, 3]);
1485    }
1486
1487    #[test]
1488    fn widening_puts_a_negative_at_the_front_and_a_positive_at_the_back() {
1489        // The whole of `widen_and_add` turns on this: the new member is outside
1490        // the range of what is there, so it needs no search, and getting the end
1491        // wrong writes it over a member instead of next to one.
1492        let mut up = of(&[-2, -1, 0, 1, 2]);
1493        up.add(70_000);
1494        assert_eq!(members(&up), [-2, -1, 0, 1, 2, 70_000]);
1495
1496        let mut down = of(&[-2, -1, 0, 1, 2]);
1497        down.add(-70_000);
1498        assert_eq!(members(&down), [-70_000, -2, -1, 0, 1, 2]);
1499    }
1500
1501    #[test]
1502    fn widening_an_empty_set_still_works() {
1503        let mut s = Intset::new();
1504        assert!(s.add(i64::MIN));
1505        assert_eq!(s.width(), 8);
1506        assert_eq!(members(&s), [i64::MIN]);
1507    }
1508
1509    #[test]
1510    fn the_extremes_of_every_width_land_in_the_width_they_belong_to() {
1511        assert_eq!(width_of(0), 2);
1512        assert_eq!(width_of(i64::from(i16::MAX)), 2);
1513        assert_eq!(width_of(i64::from(i16::MIN)), 2);
1514        assert_eq!(width_of(i64::from(i16::MAX) + 1), 4);
1515        assert_eq!(width_of(i64::from(i16::MIN) - 1), 4);
1516        assert_eq!(width_of(i64::from(i32::MAX)), 4);
1517        assert_eq!(width_of(i64::from(i32::MIN)), 4);
1518        assert_eq!(width_of(i64::from(i32::MAX) + 1), 8);
1519        assert_eq!(width_of(i64::from(i32::MIN) - 1), 8);
1520        assert_eq!(width_of(i64::MAX), 8);
1521        assert_eq!(width_of(i64::MIN), 8);
1522
1523        let s = of(&[i64::MIN, i64::MAX, 0]);
1524        assert_eq!(members(&s), [i64::MIN, 0, i64::MAX]);
1525        assert!(s.contains(i64::MIN));
1526        assert!(s.contains(i64::MAX));
1527    }
1528
1529    #[test]
1530    fn a_member_too_wide_for_the_set_is_not_in_it() {
1531        // Not merely absent, unrepresentable, and answering that without a
1532        // search is the point.
1533        let s = of(&[1, 2, 3]);
1534        assert!(!s.contains(100_000));
1535        assert!(!s.contains(i64::MAX));
1536    }
1537
1538    #[test]
1539    fn removing_takes_out_the_right_one_and_only_that_one() {
1540        let mut s = of(&[10, 20, 30, 40, 50]);
1541        assert!(s.remove(30));
1542        assert_eq!(members(&s), [10, 20, 40, 50]);
1543        assert!(!s.remove(30), "gone already");
1544        assert!(s.remove(10), "the first");
1545        assert_eq!(members(&s), [20, 40, 50]);
1546        assert!(s.remove(50), "the last");
1547        assert_eq!(members(&s), [20, 40]);
1548        assert_eq!(s.byte_len(), 8 + 2 * 2, "and the blob shrank each time");
1549    }
1550
1551    #[test]
1552    fn a_set_can_be_emptied_and_used_again() {
1553        let mut s = of(&[1, 2, 3]);
1554        for v in [1, 2, 3] {
1555            assert!(s.remove(v));
1556        }
1557        assert!(s.is_empty());
1558        assert_eq!(s.byte_len(), 8);
1559        assert!(s.add(9));
1560        assert_eq!(members(&s), [9]);
1561        sound(&s);
1562    }
1563
1564    #[test]
1565    fn every_member_of_a_big_set_is_found_and_no_stranger_is() {
1566        // Enough members to make the binary search do real work, in an order
1567        // that is neither ascending nor descending so the two range tests in
1568        // front of it are not what is being exercised.
1569        let mut s = Intset::new();
1570        for i in 0..1000i64 {
1571            assert!(s.add((i * 7919) % 1000 * 2));
1572        }
1573        assert_eq!(s.len(), 1000);
1574        for i in 0..1000i64 {
1575            assert!(s.contains(i * 2), "{} is a member", i * 2);
1576            assert!(!s.contains(i * 2 + 1), "{} is not", i * 2 + 1);
1577        }
1578        assert_eq!(members(&s), (0..1000i64).map(|i| i * 2).collect::<Vec<_>>());
1579        sound(&s);
1580    }
1581
1582    #[test]
1583    fn a_blob_survives_a_round_trip_through_bytes() {
1584        for vals in [
1585            &[][..],
1586            &[0],
1587            &[1, 2, 3],
1588            &[-70_000, 5, 70_000],
1589            &[i64::MIN, 0, i64::MAX],
1590        ] {
1591            let s = of(vals);
1592            let back = Intset::from_bytes(s.as_bytes().expect("one run")).expect("we wrote it");
1593            assert_eq!(back, s);
1594            assert_eq!(members(&back), members(&s));
1595        }
1596    }
1597
1598    #[test]
1599    fn a_blob_that_is_wrong_is_refused_rather_than_believed() {
1600        assert_eq!(Intset::from_bytes(&[]), Err(Malformed::Short));
1601        assert_eq!(
1602            Intset::from_bytes(&[2, 0, 0, 0, 0, 0, 0]),
1603            Err(Malformed::Short)
1604        );
1605
1606        let good = |vals: &[i64]| of(vals).as_bytes().expect("one run").to_vec();
1607
1608        let mut bad = good(&[1, 2, 3]);
1609        bad[0] = 3;
1610        assert_eq!(Intset::from_bytes(&bad), Err(Malformed::Width));
1611
1612        let mut short = good(&[1, 2, 3]);
1613        short.pop();
1614        assert_eq!(Intset::from_bytes(&short), Err(Malformed::Length));
1615
1616        let mut over = good(&[1, 2, 3]);
1617        over[4] = 9;
1618        assert_eq!(Intset::from_bytes(&over), Err(Malformed::Length));
1619
1620        // The one that would otherwise be believed: valid arithmetic, members
1621        // out of order, and every search after that quietly wrong.
1622        let mut jumbled = good(&[1, 2, 3]);
1623        jumbled[8..10].copy_from_slice(&9i16.to_le_bytes());
1624        assert_eq!(Intset::from_bytes(&jumbled), Err(Malformed::Order));
1625
1626        let mut twice = good(&[1, 2, 3]);
1627        twice[10..12].copy_from_slice(&1i16.to_le_bytes());
1628        assert_eq!(Intset::from_bytes(&twice), Err(Malformed::Order));
1629    }
1630
1631    #[test]
1632    fn the_header_is_little_endian_on_every_machine() {
1633        // Redis writes it little endian from a big endian host too, so a blob
1634        // this code produces has to be readable by a real server whatever it is
1635        // running on. Written out as bytes rather than as a round trip, because
1636        // a round trip through our own reader agrees with itself either way.
1637        let s = of(&[1, 258]);
1638        assert_eq!(
1639            s.as_bytes(),
1640            Some(
1641                &[
1642                    2, 0, 0, 0, // width, u32 little endian
1643                    2, 0, 0, 0, // count, u32 little endian
1644                    1, 0, // 1 as an i16 little endian
1645                    2, 1, // 258 as an i16 little endian
1646                ][..]
1647            )
1648        );
1649    }
1650
1651    #[test]
1652    fn an_ascending_fill_never_moves_anything() {
1653        // Not a timing claim, a shape claim: `search` answers past the end for
1654        // every one of these, which is the branch that makes the fill linear.
1655        let mut r = Run::new();
1656        for i in 0..100i64 {
1657            assert_eq!(r.search(i), Err(i as usize), "{i} appends");
1658            r.add(i, false);
1659        }
1660        assert_eq!(r.len(), 100);
1661    }
1662
1663    #[test]
1664    fn a_set_splits_at_the_ceiling_and_the_client_cannot_tell() {
1665        let mut s = Intset::new();
1666        for i in 0..RUN_MAX as i64 {
1667            s.add(i);
1668        }
1669        assert_eq!(s.runs(), 1, "at the ceiling it is still one array");
1670        assert!(s.as_bytes().is_some());
1671
1672        s.add(RUN_MAX as i64);
1673        assert_eq!(s.runs(), 2, "one past it splits");
1674        assert_eq!(s.as_bytes(), None, "and there is no single blob any more");
1675        assert_eq!(s.len(), RUN_MAX + 1);
1676        assert_eq!(
1677            members(&s),
1678            (0..=RUN_MAX as i64).collect::<Vec<_>>(),
1679            "and every member is still there in order"
1680        );
1681        sound(&s);
1682    }
1683
1684    #[test]
1685    fn a_scattered_fill_past_the_ceiling_stays_sorted_and_whole() {
1686        // Scattered, so the splits land in the middle of runs rather than at
1687        // the end of the last one, which is the case an ascending fill never
1688        // reaches.
1689        // Under Miri it is a fifth of the members and a fifth of the runs to
1690        // expect, since what makes the case is splitting many times over and
1691        // not the number of times.
1692        let (n, least) = if cfg!(miri) {
1693            (4_000i64, 6)
1694        } else {
1695            (20_000, 30)
1696        };
1697        let mut s = Intset::new();
1698        for i in 0..n {
1699            assert!(s.add((i * 7919) % n), "{i}");
1700        }
1701        assert_eq!(s.len(), n as usize);
1702        assert!(s.runs() > least, "it really did split, {} runs", s.runs());
1703        sound(&s);
1704        for i in 0..n {
1705            assert!(s.contains(i), "{i} is a member");
1706            assert_eq!(s.at(i as usize), i, "position {i}");
1707        }
1708        assert!(!s.contains(n));
1709        assert!(!s.contains(-1));
1710    }
1711
1712    #[test]
1713    fn draining_a_split_set_folds_the_runs_back_together() {
1714        // A run holds 512, so the Miri size still splits three times and the
1715        // merge still has a neighbour on both sides to choose between.
1716        let (n, least) = if cfg!(miri) {
1717            (1_600i64, 2)
1718        } else {
1719            (5_000, 5)
1720        };
1721        let mut s = Intset::new();
1722        for i in 0..n {
1723            s.add(i);
1724        }
1725        let split = s.runs();
1726        assert!(split > least, "{split} runs to start with");
1727        // Out from the middle, so runs empty out in the middle of the list and
1728        // the merge has a neighbour on both sides to choose between.
1729        for i in (0..n).map(|i| (i * 7919) % n) {
1730            assert!(s.remove(i), "{i}");
1731        }
1732        assert!(s.is_empty());
1733        assert_eq!(s.runs(), 1, "back to one run, not {} empty ones", s.runs());
1734        sound(&s);
1735        assert!(s.add(1));
1736        assert_eq!(members(&s), [1]);
1737    }
1738
1739    #[test]
1740    fn adds_and_removes_in_any_order_leave_the_runs_sound() {
1741        // The one that catches a wrong `bump` or a merge that loses a member,
1742        // by mirroring the whole thing against a `BTreeSet` and checking the
1743        // invariants after every write.
1744        use std::collections::BTreeSet;
1745        let mut s = Intset::new();
1746        let mut want = BTreeSet::new();
1747        // Three writes for every value in the space either way, so the mix of
1748        // adds that land and removes that find something is the same mix.
1749        let (steps, space) = if cfg!(miri) {
1750            (4_500, 1_500)
1751        } else {
1752            (12_000, 4_000)
1753        };
1754        let mut x = 12_345i64;
1755        for step in 0..steps {
1756            // A cheap deterministic spread, so the run boundaries move around
1757            // rather than the whole thing filling in one direction.
1758            x = x.wrapping_mul(6_364_136_223_846_793_005).wrapping_add(1);
1759            let v = (x >> 33) % space;
1760            if step % 3 == 2 {
1761                assert_eq!(s.remove(v), want.remove(&v), "removing {v} at {step}");
1762            } else {
1763                assert_eq!(s.add(v), want.insert(v), "adding {v} at {step}");
1764            }
1765            assert_eq!(s.len(), want.len(), "at {step}");
1766        }
1767        sound(&s);
1768        assert_eq!(members(&s), want.iter().copied().collect::<Vec<_>>());
1769    }
1770
1771    #[test]
1772    fn a_run_only_widens_the_members_it_holds() {
1773        // One array would have rewritten every member to eight bytes. Here only
1774        // the run the big member lands in pays for it, which is the one thing
1775        // this layout gives that Redis's cannot.
1776        let (n, least) = if cfg!(miri) {
1777            (2_000i64, 2)
1778        } else {
1779            (5_000, 5)
1780        };
1781        let mut s = Intset::new();
1782        for i in 0..n {
1783            s.add(i);
1784        }
1785        s.add(i64::MAX);
1786        assert_eq!(s.width(), 8, "the widest run is eight");
1787        let narrow = s.runs.iter().filter(|r| r.width() == 2).count();
1788        assert!(narrow > least, "only {narrow} runs stayed narrow");
1789        assert_eq!(s.max(), Some(i64::MAX));
1790        sound(&s);
1791    }
1792
1793    #[test]
1794    #[cfg_attr(
1795        miri,
1796        ignore = "a cost spread over a population is a claim about the population"
1797    )]
1798    fn a_run_never_holds_much_more_than_it_uses() {
1799        // The whole point of the representation. A `Vec` that doubled would put
1800        // this near four bytes a member at two byte width, and `STEP` is what
1801        // stops it.
1802        let mut s = Intset::new();
1803        for i in 0..100_000i64 {
1804            s.add(i);
1805        }
1806        let per = s.memory_bytes() as f64 / s.len() as f64;
1807        // Four byte members, because a hundred thousand is past an i16, plus
1808        // the run headers and the run list and the tree.
1809        assert!(per < 4.6, "{per:.2} bytes a member");
1810    }
1811
1812    #[test]
1813    fn the_member_at_a_position_is_the_same_one_a_walk_would_reach() {
1814        // `at` goes down the tree and `iter` goes along the runs, and the two
1815        // of them agreeing at every position either side of a run boundary is
1816        // what makes `SRANDMEMBER` on a split set draw uniformly.
1817        let n: usize = if cfg!(miri) { 1_500 } else { 3_000 };
1818        let mut s = Intset::new();
1819        for i in 0..n as i64 {
1820            s.add(i * 3);
1821        }
1822        let walked: Vec<i64> = s.iter().collect();
1823        assert_eq!(walked.len(), n);
1824        for (i, &v) in walked.iter().enumerate() {
1825            assert_eq!(s.at(i), v, "position {i}");
1826        }
1827        assert_eq!(s.get(n), None, "past the end");
1828    }
1829}