Skip to main content

yo_kv/
hll.rs

1//! The HyperLogLog sketch, byte for byte the one Redis writes.
2//!
3//! A HyperLogLog in Redis is a string, the same as a bitmap is, and the bytes of
4//! that string are a documented format rather than an internal detail. A client
5//! can `GET` a sketch out of Redis, `SET` it into us, and `PFCOUNT` it here, and
6//! it has to answer the same number. That is the whole reason this file copies
7//! the layout instead of picking a better one, and it is why every constant
8//! below is Redis's constant.
9//!
10//! # The layout
11//!
12//! Sixteen header bytes, then the registers. The header is the four magic bytes
13//! `HYLL`, an encoding byte, three unused bytes, and eight bytes of cached
14//! cardinality, little endian, with the top bit of the last one meaning the
15//! cache is stale. Every write sets that bit and [`Keyspace::pfcount`] clears it
16//! again, which is why a read of a sketch is really a write.
17//!
18//! [`Keyspace::pfcount`]: crate::Keyspace::pfcount
19//!
20//! There are 16384 registers of six bits each, so the dense form is 12288 bytes
21//! of registers and 12304 bytes altogether. The six bit fields are packed low
22//! bits first, which is the opposite way round from the bit a `SETBIT` names,
23//! and the two are unrelated: this packing is internal to the sketch.
24//!
25//! The sparse form is a run length encoding of the same registers, and a sketch
26//! stays in it until it either needs a register larger than 32 or would grow
27//! past [`SPARSE_MAX`] bytes. Three opcodes: `ZERO` is one byte and up to 64
28//! empty registers, `XZERO` is two bytes and up to 16384 of them, and `VAL` is
29//! one byte holding a value from 1 to 32 repeated up to four times. An empty
30//! sketch is one `XZERO` covering all 16384 registers, which is 18 bytes, and
31//! that is what `PFADD k` with no elements leaves behind.
32//!
33//! # The parts that had to be copied exactly
34//!
35//! The hash is MurmurHash64A with the seed `0xadc83b19`, and the register index
36//! is the low fourteen bits of it. Checked against a running 8.10.1: `a` lands
37//! in register 12711 with a count of 2, `b` in 15780 and `c` in 8436, and those
38//! are the registers a real server has after `PFADD h a b c`.
39//!
40//! The estimator is Ertl's, the one Redis moved to in 5.0, out of "New
41//! cardinality estimation algorithms for HyperLogLog sketches". It is not the
42//! original bias corrected estimator and it is not LogLog-Beta, and using either
43//! of those would answer a different number for the same bytes.
44//!
45//! [`set`] is the part with the most rules in it and all of them are Redis's:
46//! which opcode a run splits into, when a split is written as `ZERO` rather than
47//! `XZERO`, and the five opcode window that merges neighbouring `VAL` runs
48//! afterwards. Getting any of those wrong still gives a working sketch that
49//! counts correctly and does not give the same bytes, and the same bytes are the
50//! point.
51
52use yo_common::{Code, Error, Result};
53
54/// The number of bits of the hash that pick a register.
55pub const P: u32 = 14;
56/// How many registers a sketch has.
57pub const REGISTERS: usize = 1 << P;
58/// How many bits of the hash are left to count zeros in.
59pub const Q: u32 = 64 - P;
60/// How many bits a dense register takes.
61const BITS: usize = 6;
62/// The largest value a register can hold, which is what six bits reach.
63const REGISTER_MAX: u32 = 63;
64/// The header, in bytes.
65pub const HDR: usize = 16;
66/// The length of a dense sketch, header included.
67pub const DENSE: usize = HDR + REGISTERS * BITS / 8;
68/// How large a sparse sketch is allowed to get before it turns dense.
69///
70/// Redis calls this `hll-sparse-max-bytes` and defaults it to 3000. Above it the
71/// run length encoding stops paying for itself, both in space and in the walk
72/// every read does over it.
73pub const SPARSE_MAX: usize = 3000;
74
75/// The four bytes every sketch starts with.
76const MAGIC: [u8; 4] = *b"HYLL";
77/// The encoding byte of a dense sketch.
78const DENSE_TAG: u8 = 0;
79/// The encoding byte of a sparse sketch.
80const SPARSE_TAG: u8 = 1;
81/// The seed Redis hashes with, and the reason our registers are its registers.
82const SEED: u64 = 0xadc8_3b19;
83
84/// The top bits of an `XZERO` opcode.
85const XZERO_BIT: u8 = 0x40;
86/// The top bit of a `VAL` opcode.
87const VAL_BIT: u8 = 0x80;
88/// The longest run a one byte `ZERO` opcode can hold.
89const ZERO_MAX: usize = 64;
90/// The largest value a `VAL` opcode can hold, and the promotion trigger.
91const VAL_MAX: u8 = 32;
92/// How many times a `VAL` opcode can repeat its value.
93const VAL_MAX_LEN: usize = 4;
94
95/// What Redis says about a string that is not a sketch.
96///
97/// The word `WRONGTYPE` is in the message rather than only in the prefix on a
98/// real server, and the sentence has a full stop on the end, which the ordinary
99/// wrong type error does not.
100const NOT_HLL: &str = "Key is not a valid HyperLogLog string value.";
101/// What Redis says about a sketch whose opcodes do not add up.
102const CORRUPT: &str = "Corrupted HLL object detected";
103
104/// Which of the two representations a sketch is in.
105#[derive(Debug, Clone, Copy, PartialEq, Eq)]
106pub enum Encoding {
107    /// Every register written out, 12304 bytes whatever the cardinality.
108    Dense,
109    /// The registers run length encoded, which is smaller while most of them
110    /// are still empty.
111    Sparse,
112}
113
114impl Encoding {
115    /// The word `PFDEBUG ENCODING` answers with.
116    #[must_use]
117    pub const fn name(self) -> &'static str {
118        match self {
119            Encoding::Dense => "dense",
120            Encoding::Sparse => "sparse",
121        }
122    }
123}
124
125/// The error a command answers for a string that is not a sketch.
126#[must_use]
127pub fn not_hll() -> Error {
128    Error::new(Code::WrongType, NOT_HLL)
129}
130
131/// The error a command answers for a sketch whose opcodes do not add up.
132///
133/// The prefix a real server writes is neither `ERR` nor `WRONGTYPE`, it is
134/// `INVALIDOBJ`, and [`Code::Corrupt`] is what the wire layer turns into that.
135/// Nothing else on the command path answers that code, which is what makes the
136/// mapping safe to make there rather than here.
137#[must_use]
138pub fn corrupt() -> Error {
139    Error::new(Code::Corrupt, CORRUPT)
140}
141
142/// MurmurHash64A, the one Redis hashes elements with.
143///
144/// Not a hash anybody would choose today, and not one we use anywhere else. It
145/// is here because the register an element lands in is part of the file format:
146/// hash an element differently and the sketch is still a valid sketch and no
147/// longer the same sketch a real server would have written.
148#[must_use]
149pub fn hash(ele: &[u8]) -> u64 {
150    const M: u64 = 0xc6a4_a793_5bd1_e995;
151    const R: u32 = 47;
152    let mut h = SEED ^ (ele.len() as u64).wrapping_mul(M);
153    let (blocks, tail) = ele.as_chunks::<8>();
154    for block in blocks {
155        let mut k = u64::from_le_bytes(*block);
156        k = k.wrapping_mul(M);
157        k ^= k >> R;
158        k = k.wrapping_mul(M);
159        h ^= k;
160        h = h.wrapping_mul(M);
161    }
162    if !tail.is_empty() {
163        for (i, &b) in tail.iter().enumerate() {
164            h ^= u64::from(b) << (8 * i);
165        }
166        h = h.wrapping_mul(M);
167    }
168    h ^= h >> R;
169    h = h.wrapping_mul(M);
170    h ^= h >> R;
171    h
172}
173
174/// Which register an element belongs to, and what it wants written there.
175///
176/// The low fourteen bits pick the register and the rest is scanned for its first
177/// set bit, counting from one. A bit is forced in at position `Q` so that the
178/// scan always terminates, which caps the answer at `Q + 1` and is why 51 is the
179/// largest value a register ever holds.
180#[must_use]
181pub fn place(ele: &[u8]) -> (usize, u8) {
182    let h = hash(ele);
183    let index = (h & (REGISTERS as u64 - 1)) as usize;
184    let rest = (h >> P) | (1 << Q);
185    (index, rest.trailing_zeros() as u8 + 1)
186}
187
188/// Write an empty sketch, which is the header and one `XZERO` over everything.
189pub fn empty(out: &mut Vec<u8>) {
190    out.clear();
191    out.extend_from_slice(&MAGIC);
192    out.push(SPARSE_TAG);
193    out.extend_from_slice(&[0; 3]);
194    out.extend_from_slice(&[0; 8]);
195    let mut left = REGISTERS;
196    while left > 0 {
197        let run = left.min(1 << P);
198        out.extend_from_slice(&xzero_bytes(run));
199        left -= run;
200    }
201}
202
203/// Check that a string really is a sketch, and say which kind.
204///
205/// The rules are Redis's and the order matters as little as it looks: a string
206/// shorter than the header, a bad magic, an encoding byte that is neither of the
207/// two, and a dense sketch whose length is not exactly [`DENSE`] all answer the
208/// same sentence.
209pub fn check(bytes: &[u8]) -> Result<Encoding> {
210    if bytes.len() < HDR || bytes[..4] != MAGIC {
211        return Err(not_hll());
212    }
213    match bytes[4] {
214        DENSE_TAG if bytes.len() == DENSE => Ok(Encoding::Dense),
215        SPARSE_TAG => Ok(Encoding::Sparse),
216        _ => Err(not_hll()),
217    }
218}
219
220/// The cached cardinality, or `None` when a write has invalidated it.
221#[must_use]
222pub fn cached(bytes: &[u8]) -> Option<u64> {
223    let card = u64::from_le_bytes(bytes[8..16].try_into().expect("eight bytes"));
224    (card >> 63 == 0).then_some(card)
225}
226
227/// Write a freshly computed cardinality into the header, marking it good.
228pub fn cache(bytes: &mut [u8], n: u64) {
229    bytes[8..16].copy_from_slice(&(n & !(1 << 63)).to_le_bytes());
230}
231
232/// Mark the cached cardinality stale, which every write does.
233pub fn invalidate(bytes: &mut [u8]) {
234    bytes[15] |= 0x80;
235}
236
237/// Read one dense register.
238///
239/// The six bit fields are packed low bits first, so a field either sits inside
240/// one byte or straddles two. Redis always reads the second byte and relies on
241/// the string being null terminated to make the last register safe; we stop at
242/// the end of the slice instead, which is the same zero.
243#[must_use]
244#[inline]
245pub fn dense_get(regs: &[u8], index: usize) -> u8 {
246    let bit = index * BITS;
247    let (byte, shift) = (bit / 8, (bit % 8) as u32);
248    let low = u32::from(regs[byte]);
249    let high = regs.get(byte + 1).map_or(0, |&b| u32::from(b));
250    (((low >> shift) | (high << (8 - shift))) & REGISTER_MAX) as u8
251}
252
253/// Write one dense register, answering whether it changed.
254///
255/// The second byte is only touched when the field really straddles two, which is
256/// what keeps a write to the last register inside the slice.
257#[inline]
258pub fn dense_set(regs: &mut [u8], index: usize, val: u8) -> bool {
259    if dense_get(regs, index) >= val {
260        return false;
261    }
262    let bit = index * BITS;
263    let (byte, shift) = (bit / 8, (bit % 8) as u32);
264    let v = u32::from(val);
265    regs[byte] = ((u32::from(regs[byte]) & !(REGISTER_MAX << shift)) | (v << shift)) as u8;
266    if shift > 2 {
267        let rest = 8 - shift;
268        let high = &mut regs[byte + 1];
269        *high = ((u32::from(*high) & !(REGISTER_MAX >> rest)) | (v >> rest)) as u8;
270    }
271    true
272}
273
274/// Whether the byte at `p` starts a `ZERO` opcode.
275const fn is_zero(b: u8) -> bool {
276    b & 0xc0 == 0
277}
278
279/// Whether the byte at `p` starts an `XZERO` opcode.
280const fn is_xzero(b: u8) -> bool {
281    b & 0xc0 == XZERO_BIT
282}
283
284/// Whether the byte at `p` is a `VAL` opcode.
285const fn is_val(b: u8) -> bool {
286    b & VAL_BIT != 0
287}
288
289/// How many registers a `ZERO` opcode covers.
290const fn zero_len(b: u8) -> usize {
291    (b & 0x3f) as usize + 1
292}
293
294/// How many registers an `XZERO` opcode covers.
295const fn xzero_len(a: u8, b: u8) -> usize {
296    (((a & 0x3f) as usize) << 8 | b as usize) + 1
297}
298
299/// The value a `VAL` opcode holds.
300const fn val_value(b: u8) -> u8 {
301    ((b >> 2) & 0x1f) + 1
302}
303
304/// How many registers a `VAL` opcode covers.
305const fn val_len(b: u8) -> usize {
306    (b & 3) as usize + 1
307}
308
309/// A `VAL` opcode holding `val` repeated `len` times.
310const fn val_byte(val: u8, len: usize) -> u8 {
311    ((val - 1) << 2) | (len as u8 - 1) | VAL_BIT
312}
313
314/// A `ZERO` opcode covering `len` registers.
315const fn zero_byte(len: usize) -> u8 {
316    (len - 1) as u8
317}
318
319/// An `XZERO` opcode covering `len` registers.
320const fn xzero_bytes(len: usize) -> [u8; 2] {
321    let n = len - 1;
322    [((n >> 8) as u8) | XZERO_BIT, (n & 0xff) as u8]
323}
324
325/// How many registers the opcode at `at` covers, and how many bytes it is.
326fn opcode(sparse: &[u8], at: usize) -> Option<(usize, usize)> {
327    let b = *sparse.get(at)?;
328    if is_zero(b) {
329        Some((zero_len(b), 1))
330    } else if is_xzero(b) {
331        Some((xzero_len(b, *sparse.get(at + 1)?), 2))
332    } else {
333        Some((val_len(b), 1))
334    }
335}
336
337/// Walk a sparse body, handing each run to `each` as a value and a length.
338///
339/// `false` when the runs do not add up to exactly [`REGISTERS`], which is the
340/// only corruption any of the readers here can detect and the one Redis checks
341/// for as well.
342fn walk(sparse: &[u8], mut each: impl FnMut(u8, usize, usize)) -> bool {
343    let mut at = 0;
344    let mut index = 0;
345    while at < sparse.len() {
346        let b = sparse[at];
347        if is_val(b) {
348            let len = val_len(b);
349            if index + len > REGISTERS {
350                return false;
351            }
352            each(val_value(b), index, len);
353            index += len;
354            at += 1;
355        } else if is_zero(b) {
356            index += zero_len(b);
357            at += 1;
358        } else {
359            let Some(&next) = sparse.get(at + 1) else {
360                return false;
361            };
362            index += xzero_len(b, next);
363            at += 2;
364        }
365    }
366    index == REGISTERS
367}
368
369/// Turn a sparse sketch into a dense one in place.
370///
371/// `false` for a body whose opcodes do not add up, which leaves the buffer as it
372/// was. The registers go through the stack rather than through a second buffer,
373/// which is sixteen kibibytes and the same thing Redis does when it merges.
374pub fn to_dense(buf: &mut Vec<u8>) -> bool {
375    if buf[4] == DENSE_TAG {
376        return true;
377    }
378    let mut regs = [0u8; REGISTERS];
379    if !walk(&buf[HDR..], |val, at, len| {
380        regs[at..at + len].fill(val);
381    }) {
382        return false;
383    }
384    buf.truncate(HDR);
385    buf.resize(DENSE, 0);
386    buf[4] = DENSE_TAG;
387    let body = &mut buf[HDR..];
388    for (i, &val) in regs.iter().enumerate() {
389        if val != 0 {
390            dense_set(body, i, val);
391        }
392    }
393    true
394}
395
396/// Raise one register to `val`, answering whether anything changed.
397///
398/// `None` for a corrupted body. A sparse sketch turns dense here when the value
399/// will not fit in a `VAL` opcode or when the rewrite would push it past
400/// [`SPARSE_MAX`], and the caller does not have to know which happened.
401pub fn set(buf: &mut Vec<u8>, index: usize, val: u8) -> Option<bool> {
402    if buf[4] == DENSE_TAG {
403        let changed = dense_set(&mut buf[HDR..], index, val);
404        if changed {
405            invalidate(buf);
406        }
407        return Some(changed);
408    }
409    if val > VAL_MAX {
410        return promote(buf, index, val);
411    }
412
413    // Find the opcode covering the register, keeping the one before it, which is
414    // where the merge pass at the end starts from.
415    let (mut at, mut first, mut prev, mut span) = (HDR, 0usize, None, 0usize);
416    while at < buf.len() {
417        let (covers, bytes) = opcode(buf, at)?;
418        span = covers;
419        if index < first + span {
420            break;
421        }
422        prev = Some(at);
423        at += bytes;
424        first += span;
425    }
426    if span == 0 || at >= buf.len() {
427        return None;
428    }
429
430    let here = buf[at];
431    let (zero, xzero, run) = if is_val(here) {
432        (false, false, val_len(here))
433    } else if is_zero(here) {
434        (true, false, zero_len(here))
435    } else {
436        (false, true, xzero_len(here, *buf.get(at + 1)?))
437    };
438
439    // Two shapes need no rewriting at all. A run already holding a value this
440    // large is the common case once a sketch has any size to it, and a single
441    // register run is written over where it lies whatever it held.
442    if is_val(here) {
443        if val_value(here) >= val {
444            return Some(false);
445        }
446        if run == 1 {
447            buf[at] = val_byte(val, 1);
448            return Some(finish(buf, prev));
449        }
450    }
451    if zero && run == 1 {
452        buf[at] = val_byte(val, 1);
453        return Some(finish(buf, prev));
454    }
455
456    // Everything else splits the run into up to three opcodes, which is five
457    // bytes in the worst case: an `XZERO` on each side of a one register `VAL`.
458    let mut seq = [0u8; 5];
459    let mut n = 0;
460    let last = first + span - 1;
461    let gap = |seq: &mut [u8; 5], n: &mut usize, len: usize| {
462        if len > ZERO_MAX {
463            seq[*n..*n + 2].copy_from_slice(&xzero_bytes(len));
464            *n += 2;
465        } else {
466            seq[*n] = zero_byte(len);
467            *n += 1;
468        }
469    };
470    if zero || xzero {
471        if index != first {
472            gap(&mut seq, &mut n, index - first);
473        }
474        seq[n] = val_byte(val, 1);
475        n += 1;
476        if index != last {
477            gap(&mut seq, &mut n, last - index);
478        }
479    } else {
480        let had = val_value(here);
481        if index != first {
482            seq[n] = val_byte(had, index - first);
483            n += 1;
484        }
485        seq[n] = val_byte(val, 1);
486        n += 1;
487        if index != last {
488            seq[n] = val_byte(had, last - index);
489            n += 1;
490        }
491    }
492
493    // Put the new opcodes where the old one was. Growing past the sparse limit
494    // is what turns a sketch dense in the ordinary case, long before any single
495    // register needs a value larger than 32.
496    let old = if xzero { 2 } else { 1 };
497    let end = buf.len();
498    if n > old && end + (n - old) > SPARSE_MAX {
499        return promote(buf, index, val);
500    }
501    if n > old {
502        buf.resize(end + (n - old), 0);
503        buf.copy_within(at + old..end, at + n);
504    } else if n < old {
505        buf.copy_within(at + old..end, at + n);
506        buf.truncate(end - (old - n));
507    }
508    buf[at..at + n].copy_from_slice(&seq[..n]);
509    Some(finish(buf, prev))
510}
511
512/// Turn the sketch dense and write the register into it.
513fn promote(buf: &mut Vec<u8>, index: usize, val: u8) -> Option<bool> {
514    if !to_dense(buf) {
515        return None;
516    }
517    let changed = dense_set(&mut buf[HDR..], index, val);
518    invalidate(buf);
519    Some(changed)
520}
521
522/// Tidy up after a write and mark the cache stale.
523///
524/// The tidying is Redis's five opcode window: a split can leave two `VAL`
525/// opcodes holding the same value next to each other, and joining them back up
526/// is what stops a sketch from growing a byte every time a register is written
527/// twice. Five is Redis's number and it is enough, because a single write can
528/// only ever produce three new opcodes.
529fn finish(buf: &mut Vec<u8>, prev: Option<usize>) -> bool {
530    let mut at = prev.unwrap_or(HDR);
531    let mut left = 5;
532    while at < buf.len() && left > 0 {
533        left -= 1;
534        let b = buf[at];
535        if is_xzero(b) {
536            at += 2;
537            continue;
538        }
539        if is_zero(b) {
540            at += 1;
541            continue;
542        }
543        if let Some(&next) = buf.get(at + 1)
544            && is_val(next)
545            && val_value(b) == val_value(next)
546        {
547            let len = val_len(b) + val_len(next);
548            if len <= VAL_MAX_LEN {
549                buf[at + 1] = val_byte(val_value(b), len);
550                let end = buf.len();
551                buf.copy_within(at + 1..end, at);
552                buf.truncate(end - 1);
553                // Try the merged opcode against its new neighbour before moving
554                // on, which is what lets four ones become one run.
555                continue;
556            }
557        }
558        at += 1;
559    }
560    invalidate(buf);
561    true
562}
563
564/// How many registers hold each value, which is all the estimator needs.
565///
566/// `None` for a body whose opcodes do not add up.
567fn histogram(bytes: &[u8], enc: Encoding) -> Option<[u32; 64]> {
568    let mut hist = [0u32; 64];
569    match enc {
570        Encoding::Dense => {
571            let regs = &bytes[HDR..];
572            for i in 0..REGISTERS {
573                hist[dense_get(regs, i) as usize] += 1;
574            }
575        }
576        Encoding::Sparse => {
577            let mut seen = 0;
578            if !walk(&bytes[HDR..], |val, _, len| {
579                hist[val as usize] += len as u32;
580                seen += len as u32;
581            }) {
582                return None;
583            }
584            hist[0] = REGISTERS as u32 - seen;
585        }
586    }
587    Some(hist)
588}
589
590/// Every register of a sketch, which is what a merge and `PFDEBUG GETREG` want.
591///
592/// `false` for a body whose opcodes do not add up. The registers are raised to
593/// what the sketch holds rather than overwritten, so merging several sketches is
594/// calling this once per sketch over the same array.
595pub fn merge(max: &mut [u8; REGISTERS], bytes: &[u8], enc: Encoding) -> bool {
596    match enc {
597        Encoding::Dense => {
598            let regs = &bytes[HDR..];
599            for (i, slot) in max.iter_mut().enumerate() {
600                *slot = (*slot).max(dense_get(regs, i));
601            }
602            true
603        }
604        Encoding::Sparse => walk(&bytes[HDR..], |val, at, len| {
605            for slot in &mut max[at..at + len] {
606                *slot = (*slot).max(val);
607            }
608        }),
609    }
610}
611
612/// Ertl's tau, the correction for the registers that are already saturated.
613fn tau(mut x: f64) -> f64 {
614    if x == 0.0 || x == 1.0 {
615        return 0.0;
616    }
617    let mut y = 1.0;
618    let mut z = 1.0 - x;
619    loop {
620        x = x.sqrt();
621        let was = z;
622        y *= 0.5;
623        z -= (1.0 - x).powi(2) * y;
624        if was == z {
625            return z / 3.0;
626        }
627    }
628}
629
630/// Ertl's sigma, the correction for the registers that are still empty.
631fn sigma(mut x: f64) -> f64 {
632    if x == 1.0 {
633        return f64::INFINITY;
634    }
635    let mut y = 1.0;
636    let mut z = x;
637    loop {
638        x *= x;
639        let was = z;
640        z += x * y;
641        y += y;
642        if was == z {
643            return z;
644        }
645    }
646}
647
648/// The cardinality a register histogram implies.
649///
650/// This is Ertl's estimator out of "New cardinality estimation algorithms for
651/// HyperLogLog sketches", which is what Redis has used since 5.0. It replaced
652/// the original bias corrected estimator and the switch to linear counting at
653/// small cardinalities, and it is a single expression with no thresholds in it.
654/// Writing a different estimator here would answer a different number for a
655/// sketch a real server wrote, which is the one thing this file must not do.
656#[must_use]
657pub fn estimate(hist: &[u32; 64]) -> u64 {
658    /// One over twice the natural log of two, which is the limit of the alpha
659    /// correction as the register count grows. Redis writes it out as a literal
660    /// rather than computing it and so do we, so the last bit agrees.
661    const ALPHA_INF: f64 = 0.721_347_520_444_481_7;
662    let m = REGISTERS as f64;
663    let mut z = m * tau((m - f64::from(hist[Q as usize + 1])) / m);
664    for j in (1..=Q as usize).rev() {
665        z += f64::from(hist[j]);
666        z *= 0.5;
667    }
668    z += m * sigma(f64::from(hist[0]) / m);
669    (ALPHA_INF * m * m / z).round() as u64
670}
671
672/// The cardinality of one sketch, ignoring whatever the header has cached.
673pub fn count(bytes: &[u8], enc: Encoding) -> Result<u64> {
674    match histogram(bytes, enc) {
675        Some(hist) => Ok(estimate(&hist)),
676        None => Err(corrupt()),
677    }
678}
679
680/// The sparse opcodes written out, which is what `PFDEBUG DECODE` answers.
681///
682/// The spelling is a real server's, checked against 8.10.1 on a sketch built by
683/// hand to have one of each: lowercase `z` for a `ZERO` run, uppercase `Z` for
684/// an `XZERO` one, and `v` for a value and how many times it repeats, the three
685/// separated by single spaces. It goes into a byte buffer rather than a `String`
686/// so that the caller can hand it the one it already has.
687pub fn decode(bytes: &[u8], out: &mut Vec<u8>) {
688    use std::io::Write;
689    let sparse = &bytes[HDR..];
690    let mut at = 0;
691    while at < sparse.len() {
692        if !out.is_empty() {
693            out.push(b' ');
694        }
695        let b = sparse[at];
696        if is_val(b) {
697            let _ = write!(out, "v:{},{}", val_value(b), val_len(b));
698            at += 1;
699        } else if is_zero(b) {
700            let _ = write!(out, "z:{}", zero_len(b));
701            at += 1;
702        } else {
703            let Some(&next) = sparse.get(at + 1) else {
704                return;
705            };
706            let _ = write!(out, "Z:{}", xzero_len(b, next));
707            at += 2;
708        }
709    }
710}
711
712#[cfg(test)]
713mod tests {
714    use super::*;
715    use crate::many;
716
717    /// The hash is the file format, so it is pinned against a real server.
718    ///
719    /// These three registers were read out of a running 8.10.1 with `PFDEBUG
720    /// GETREG` after `PFADD h a b c`, and they are the whole reason this file
721    /// carries a hash function nothing else in the tree uses.
722    #[test]
723    fn an_element_lands_where_a_real_server_puts_it() {
724        assert_eq!(place(b"a"), (12711, 2));
725        assert_eq!(place(b"b"), (15780, 1));
726        assert_eq!(place(b"c"), (8436, 1));
727    }
728
729    /// The bytes of an empty sketch, and of one with three elements in it.
730    ///
731    /// Both were read off a real server with `GET`. Both have the cache marked
732    /// stale, and on the empty one that is not this function's doing: Redis
733    /// creates the sketch with a valid cache of zero and `PFADD` invalidates it
734    /// on the way out, even when it added nothing. The bytes a client can see
735    /// are the ones that matter and they have the top bit set.
736    #[test]
737    fn a_sketch_is_the_bytes_a_real_server_writes() {
738        let mut buf = Vec::new();
739        empty(&mut buf);
740        assert_eq!(buf.len(), 18);
741        assert_eq!(&buf[..], b"HYLL\x01\0\0\0\0\0\0\0\0\0\0\0\x7f\xff");
742        invalidate(&mut buf);
743        assert_eq!(&buf[..], b"HYLL\x01\0\0\0\0\0\0\0\0\0\0\x80\x7f\xff");
744
745        for ele in [&b"a"[..], b"b", b"c"] {
746            let (index, val) = place(ele);
747            assert_eq!(set(&mut buf, index, val), Some(true));
748        }
749        assert_eq!(
750            &buf[..],
751            b"HYLL\x01\0\0\0\0\0\0\0\0\0\0\x80\x60\xf3\x80\x50\xb1\x84\x4b\xfb\x80\x42\x5a"
752        );
753    }
754
755    /// A second write of the same element changes nothing at all.
756    #[test]
757    fn writing_the_same_element_twice_is_not_a_change() {
758        let mut buf = Vec::new();
759        empty(&mut buf);
760        let (index, val) = place(b"a");
761        assert_eq!(set(&mut buf, index, val), Some(true));
762        let before = buf.clone();
763        assert_eq!(set(&mut buf, index, val), Some(false));
764        assert_eq!(buf, before);
765    }
766
767    /// The dense packing, against the obvious slow version of itself.
768    #[test]
769    fn a_dense_register_is_six_bits_packed_from_the_bottom() {
770        let mut regs = vec![0u8; REGISTERS * BITS / 8];
771        let mut want = vec![0u8; REGISTERS];
772        for (i, slot) in want.iter_mut().enumerate() {
773            *slot = ((i * 7 + 1) % 52) as u8;
774        }
775        // Written in a shuffled order so that a write that spilled into its
776        // neighbour would be caught rather than overwritten afterwards.
777        for step in [1usize, 3, 5] {
778            let mut i = 0;
779            while i < REGISTERS {
780                let val = want[i];
781                if val > dense_get(&regs, i) {
782                    assert!(dense_set(&mut regs, i, val));
783                }
784                i += step;
785            }
786        }
787        for (i, &val) in want.iter().enumerate() {
788            assert_eq!(dense_get(&regs, i), val, "register {i}");
789        }
790        // A write that would lower a register is refused, the way a sketch needs.
791        assert!(!dense_set(&mut regs, 5, 0));
792    }
793
794    /// A sparse sketch and the dense one it turns into hold the same registers.
795    #[test]
796    fn turning_dense_keeps_every_register() {
797        let mut buf = Vec::new();
798        empty(&mut buf);
799        let mut want = [0u8; REGISTERS];
800        for i in 0..400 {
801            let ele = format!("e:{i}");
802            let (index, val) = place(ele.as_bytes());
803            set(&mut buf, index, val).expect("a write");
804            want[index] = want[index].max(val);
805        }
806        let sparse = count(&buf, Encoding::Sparse).expect("a count");
807
808        assert!(to_dense(&mut buf));
809        assert_eq!(buf.len(), DENSE);
810        assert_eq!(check(&buf).expect("a sketch"), Encoding::Dense);
811        for (i, &val) in want.iter().enumerate() {
812            assert_eq!(dense_get(&buf[HDR..], i), val, "register {i}");
813        }
814        assert_eq!(count(&buf, Encoding::Dense).expect("a count"), sparse);
815    }
816
817    /// Enough elements to push a sketch over the sparse limit on its own.
818    #[test]
819    #[cfg_attr(
820        miri,
821        ignore = "fewer elements do not outgrow the sparse form, which is the claim"
822    )]
823    fn a_sketch_turns_dense_when_it_outgrows_the_sparse_form() {
824        let mut buf = Vec::new();
825        empty(&mut buf);
826        for i in 0..2000 {
827            let ele = format!("e:{i}");
828            let (index, val) = place(ele.as_bytes());
829            set(&mut buf, index, val).expect("a write");
830            assert!(buf.len() <= SPARSE_MAX || buf.len() == DENSE);
831        }
832        assert_eq!(check(&buf).expect("a sketch"), Encoding::Dense);
833    }
834
835    /// A value larger than a `VAL` opcode can hold turns the sketch dense.
836    #[test]
837    fn a_large_register_turns_the_sketch_dense() {
838        let mut buf = Vec::new();
839        empty(&mut buf);
840        assert_eq!(set(&mut buf, 100, VAL_MAX), Some(true));
841        assert_eq!(check(&buf).expect("a sketch"), Encoding::Sparse);
842        assert_eq!(set(&mut buf, 200, VAL_MAX + 1), Some(true));
843        assert_eq!(check(&buf).expect("a sketch"), Encoding::Dense);
844        assert_eq!(dense_get(&buf[HDR..], 100), VAL_MAX);
845        assert_eq!(dense_get(&buf[HDR..], 200), VAL_MAX + 1);
846    }
847
848    /// Neighbouring runs of the same value are joined back up.
849    ///
850    /// Without the merge pass a sketch grows an opcode for every register
851    /// written, and four ones in a row here would be four bytes instead of one.
852    #[test]
853    fn neighbouring_runs_of_the_same_value_are_joined() {
854        let mut buf = Vec::new();
855        empty(&mut buf);
856        for i in 0..4 {
857            set(&mut buf, 100 + i, 1).expect("a write");
858        }
859        let mut decoded = Vec::new();
860        decode(&buf, &mut decoded);
861        assert_eq!(decoded, b"Z:100 v:1,4 Z:16280");
862    }
863
864    /// One sketch with all three opcodes in it, against a real server.
865    ///
866    /// The bytes were written by hand, `SET` into 8.10.1, and read back through
867    /// `PFDEBUG DECODE` and `PFCOUNT`. It is the only case that pins the
868    /// lowercase `z`, since a sketch built by adding elements rarely has a gap
869    /// short enough to need one.
870    #[test]
871    fn all_three_opcodes_decode_the_way_a_real_server_prints_them() {
872        let mut buf = Vec::new();
873        empty(&mut buf);
874        buf.truncate(HDR);
875        buf.extend_from_slice(&xzero_bytes(100));
876        buf.push(val_byte(1, 4));
877        buf.push(zero_byte(10));
878        buf.push(val_byte(3, 2));
879        buf.extend_from_slice(&xzero_bytes(REGISTERS - 100 - 4 - 10 - 2));
880
881        let mut decoded = Vec::new();
882        decode(&buf, &mut decoded);
883        assert_eq!(decoded, b"Z:100 v:1,4 z:10 v:3,2 Z:16268");
884        assert_eq!(count(&buf, Encoding::Sparse).expect("a count"), 6);
885    }
886
887    /// The counted answer against the sketch itself, over a range of sizes.
888    ///
889    /// A HyperLogLog is allowed to be wrong and this checks it is wrong by less
890    /// than the two percent the parameters promise, which is what would catch a
891    /// register being written in the wrong place.
892    #[test]
893    #[cfg_attr(
894        miri,
895        ignore = "an error bound only means something at the sizes that produce it"
896    )]
897    fn the_estimate_is_close_to_the_truth() {
898        for n in [10usize, 100, 1000, 10_000, 100_000] {
899            let mut buf = Vec::new();
900            empty(&mut buf);
901            for i in 0..n {
902                let ele = format!("element:{i}");
903                let (index, val) = place(ele.as_bytes());
904                set(&mut buf, index, val).expect("a write");
905            }
906            let enc = check(&buf).expect("a sketch");
907            let got = count(&buf, enc).expect("a count") as f64;
908            let off = (got - n as f64).abs() / n as f64;
909            assert!(off < 0.02, "{n} counted as {got}");
910        }
911    }
912
913    /// The numbers a real 8.10.1 answered for the same elements.
914    ///
915    /// A sketch that counted correctly and not identically would be useless for
916    /// the thing this is for, which is a client moving sketches between servers.
917    #[test]
918    #[cfg_attr(
919        miri,
920        ignore = "the counts are the claim, so there is no smaller version of it"
921    )]
922    fn the_estimate_is_the_number_a_real_server_gives() {
923        for (n, want) in [(100usize, 100u64), (1000, 995), (10_000, 10_077)] {
924            let mut buf = Vec::new();
925            empty(&mut buf);
926            for i in 0..n {
927                let ele = format!("e:{i}");
928                let (index, val) = place(ele.as_bytes());
929                set(&mut buf, index, val).expect("a write");
930            }
931            let enc = check(&buf).expect("a sketch");
932            assert_eq!(count(&buf, enc).expect("a count"), want, "{n} elements");
933        }
934    }
935
936    /// The two sizes the milestone gate names, on the elements that reach them.
937    #[test]
938    #[cfg_attr(
939        miri,
940        ignore = "the sizes are the claim, and only these counts reach them"
941    )]
942    fn the_two_sizes_are_the_ones_a_real_server_has() {
943        let build = |n: usize| {
944            let mut buf = Vec::new();
945            empty(&mut buf);
946            for i in 0..n {
947                let ele = format!("e:{i}");
948                let (index, val) = place(ele.as_bytes());
949                set(&mut buf, index, val).expect("a write");
950            }
951            buf
952        };
953        assert_eq!(build(1000).len(), 1880);
954        assert_eq!(build(10_000).len(), DENSE);
955        assert_eq!(DENSE, 12304);
956        const { assert!(1880 <= SPARSE_MAX) };
957    }
958
959    /// What is refused, and what a stale cache looks like.
960    #[test]
961    fn a_string_that_is_not_a_sketch_is_refused() {
962        assert!(check(b"").is_err());
963        assert!(check(b"HYLL").is_err());
964        assert!(check(b"NOPE\x01\0\0\0\0\0\0\0\0\0\0\0\x7f\xff").is_err());
965        assert!(check(b"HYLL\x02\0\0\0\0\0\0\0\0\0\0\0\x7f\xff").is_err());
966        // A dense sketch has to be exactly the right length.
967        assert!(check(b"HYLL\0\0\0\0\0\0\0\0\0\0\0\0\x7f\xff").is_err());
968
969        let mut buf = Vec::new();
970        empty(&mut buf);
971        assert_eq!(cached(&buf), Some(0));
972        cache(&mut buf, 12345);
973        assert_eq!(cached(&buf), Some(12345));
974        invalidate(&mut buf);
975        assert_eq!(cached(&buf), None);
976    }
977
978    /// Runs that do not cover all 16384 registers are a corrupted sketch.
979    #[test]
980    fn a_body_that_does_not_add_up_is_corrupt() {
981        let mut buf = Vec::new();
982        empty(&mut buf);
983        buf.truncate(buf.len() - 1);
984        assert!(count(&buf, Encoding::Sparse).is_err());
985        let mut short = Vec::new();
986        empty(&mut short);
987        short.pop();
988        short.pop();
989        assert!(count(&short, Encoding::Sparse).is_err());
990    }
991
992    /// Merging takes the larger of each register, whichever form it is in.
993    #[test]
994    fn merging_takes_the_larger_of_every_register() {
995        let build = |from: usize, to: usize| {
996            let mut buf = Vec::new();
997            empty(&mut buf);
998            for i in from..to {
999                let ele = format!("e:{i}");
1000                let (index, val) = place(ele.as_bytes());
1001                set(&mut buf, index, val).expect("a write");
1002            }
1003            buf
1004        };
1005        // Two runs that overlap in the middle, so the merge has registers only
1006        // one side has and registers both sides disagree about.
1007        let (mid, end) = (many(400usize), many(900usize));
1008        let a = build(0, end - mid);
1009        let b = build(mid, end);
1010        let mut max = [0u8; REGISTERS];
1011        assert!(merge(&mut max, &a, Encoding::Sparse));
1012        assert!(merge(&mut max, &b, Encoding::Sparse));
1013
1014        let both = build(0, end);
1015        let mut want = [0u8; REGISTERS];
1016        assert!(merge(&mut want, &both, check(&both).expect("a sketch")));
1017        assert_eq!(max, want);
1018    }
1019}