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}