probabilistic_machines 0.1.1

data/ext/probabilistic_machines_native/core/src/countmin.rs

@@ -0,0 +1,240 @@
+//! Frequency estimation for a stream nobody wants to store.
+
+use crate::hash::{hash64, mix64, u32_low_u8, u64_low_u16, usize_low_u8};
+use crate::{CodecError, TAG_COUNTMIN, read8};
+
+const HEADER_LEN: usize = 12;
+const COUNTER_LEN: usize = 4;
+
+/// Counts how often each item appeared, in `W * D` counters.
+///
+/// `W` is the row width, a power of two up to 65536; `D` is the number of rows.
+/// An estimate is never below the true count while [`saturated`](Self::saturated)
+/// is false, and overshoots it by more than
+/// `e * total / W` with probability at most `e^-D`, where `total` is the sum of
+/// every count added. 2048 columns over 5 rows hold the overshoot under 0.14%
+/// of the total with probability 0.993.
+///
+/// ```
+/// use probabilistic_machines::CountMinSketch;
+///
+/// let mut hits = CountMinSketch::<1024, 4>::new();
+/// hits.add_count(b"/login", 900);
+/// hits.add(b"/health");
+///
+/// assert!(hits.estimate(b"/login") >= 900);
+/// assert_eq!(hits.estimate(b"/never-seen"), 0);
+/// assert_eq!(hits.total(), 901);
+/// ```
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct CountMinSketch<const W: usize, const D: usize> {
+    rows: [[u32; W]; D],
+    saturated: bool,
+    total: u64,
+}
+
+impl<const W: usize, const D: usize> CountMinSketch<W, D> {
+    const COLUMN_MASK: usize = {
+        assert!(W.is_power_of_two(), "row width must be a power of two");
+        assert!(W >= 2, "row width must be at least 2");
+        assert!(W <= 1 << 16, "row width must be at most 65536");
+        assert!(D >= 1, "a sketch needs at least one row");
+        assert!(D <= 64, "a sketch needs at most 64 rows");
+        W - 1
+    };
+
+    /// Length of the [`encode`](Self::encode) frame.
+    pub const ENCODED_LEN: usize = HEADER_LEN + W * D * COUNTER_LEN;
+
+    #[must_use]
+    pub const fn new() -> Self {
+        Self {
+            rows: [[0; W]; D],
+            total: 0,
+            saturated: false,
+        }
+    }
+
+    /// Sum of every count added, saturating at [`u64::MAX`].
+    #[must_use]
+    pub const fn total(&self) -> u64 {
+        self.total
+    }
+
+    /// True once a counter has hit [`u32::MAX`], after which an estimate can
+    /// fall below the true count.
+    #[must_use]
+    pub const fn saturated(&self) -> bool {
+        self.saturated
+    }
+
+    pub fn add(&mut self, item: &[u8]) {
+        self.add_count(item, 1);
+    }
+
+    pub fn add_count(&mut self, item: &[u8], count: u32) {
+        self.add_hash_count(hash64(item), count);
+    }
+
+    /// Adds to an item the caller already hashed.
+    pub fn add_hash_count(&mut self, hash: u64, count: u32) {
+        let mut cursor = hash;
+        for row in &mut self.rows {
+            cursor = mix64(cursor);
+            let column = usize::from(u64_low_u16(cursor)) & Self::COLUMN_MASK;
+            let (sum, overflowed) = row[column].overflowing_add(count);
+            self.saturated |= overflowed;
+            row[column] = if overflowed { u32::MAX } else { sum };
+        }
+        self.total = self.total.saturating_add(u64::from(count));
+    }
+
+    /// Upper bound on how often `item` was added.
+    #[must_use]
+    pub fn estimate(&self, item: &[u8]) -> u32 {
+        self.estimate_hash(hash64(item))
+    }
+
+    /// Upper bound for an item the caller already hashed.
+    #[must_use]
+    pub fn estimate_hash(&self, hash: u64) -> u32 {
+        let mut cursor = hash;
+        let mut smallest = u32::MAX;
+        for row in &self.rows {
+            cursor = mix64(cursor);
+            let column = usize::from(u64_low_u16(cursor)) & Self::COLUMN_MASK;
+            if row[column] < smallest {
+                smallest = row[column];
+            }
+        }
+        smallest
+    }
+
+    /// Adds another sketch's counters to this one, which is exact for sketches
+    /// of the same shape: the result equals one sketch fed both streams.
+    pub fn merge(&mut self, other: &Self) {
+        let mut saturated = self.saturated || other.saturated;
+        for (mine, theirs) in self.rows.iter_mut().zip(&other.rows) {
+            for (counter, addend) in mine.iter_mut().zip(theirs) {
+                let (sum, overflowed) = counter.overflowing_add(*addend);
+                saturated |= overflowed;
+                *counter = if overflowed { u32::MAX } else { sum };
+            }
+        }
+        self.saturated = saturated;
+        self.total = self.total.saturating_add(other.total);
+    }
+
+    /// Writes the sketch into `out`, returning the number of bytes written.
+    ///
+    /// # Errors
+    ///
+    /// [`CodecError::Truncated`] when `out` is shorter than
+    /// [`ENCODED_LEN`](Self::ENCODED_LEN).
+    pub fn encode(&self, out: &mut [u8]) -> Result<usize, CodecError> {
+        let Some(frame) = out.get_mut(..Self::ENCODED_LEN) else {
+            return Err(CodecError::Truncated);
+        };
+        self.write_frame(frame);
+        Ok(Self::ENCODED_LEN)
+    }
+
+    /// The sketch's frame, exactly [`ENCODED_LEN`](Self::ENCODED_LEN) bytes.
+    #[cfg(feature = "alloc")]
+    #[must_use]
+    pub fn encoded(&self) -> alloc::boxed::Box<[u8]> {
+        let mut frame = alloc::vec![0_u8; Self::ENCODED_LEN].into_boxed_slice();
+        self.write_frame(&mut frame);
+        frame
+    }
+
+    /// Fills a frame of exactly [`ENCODED_LEN`](Self::ENCODED_LEN) bytes.
+    fn write_frame(&self, frame: &mut [u8]) {
+        frame[0] = TAG_COUNTMIN;
+        frame[1] = u32_low_u8(W.trailing_zeros());
+        frame[2] = usize_low_u8(D);
+        frame[3..11].copy_from_slice(&self.total.to_le_bytes());
+        frame[11] = u8::from(self.saturated);
+
+        let counters = self.rows.iter().flatten();
+        let (slots, _) = frame[HEADER_LEN..].as_chunks_mut::<COUNTER_LEN>();
+        for (slot, counter) in slots.iter_mut().zip(counters) {
+            *slot = counter.to_le_bytes();
+        }
+    }
+
+    /// Reads a sketch written by [`encode`](Self::encode).
+    ///
+    /// # Errors
+    ///
+    /// [`CodecError::Truncated`] when `bytes` is too short, [`CodecError::Tag`]
+    /// when it holds another structure, [`CodecError::Parameters`] when its
+    /// shape differs from `W` by `D` or its counters contradict its total.
+    pub fn decode(bytes: &[u8]) -> Result<Self, CodecError> {
+        let mut sketch = Self::new();
+        sketch.decode_into(bytes)?;
+        Ok(sketch)
+    }
+
+    /// Reads a frame over this sketch, for a caller that owns one already and
+    /// has no room on its stack for a second.
+    ///
+    /// The sketch is untouched when the frame is rejected.
+    ///
+    /// # Errors
+    ///
+    /// The same as [`decode`](Self::decode).
+    pub fn decode_into(&mut self, bytes: &[u8]) -> Result<(), CodecError> {
+        let Some(frame) = bytes.get(..Self::ENCODED_LEN) else {
+            return Err(CodecError::Truncated);
+        };
+        if frame[0] != TAG_COUNTMIN {
+            return Err(CodecError::Tag);
+        }
+        if frame[1] != u32_low_u8(W.trailing_zeros()) || frame[2] != usize_low_u8(D) {
+            return Err(CodecError::Parameters);
+        }
+
+        let total = u64::from_le_bytes(read8(&frame[3..11]));
+        let saturated = frame[11] != 0;
+        let (slots, _) = frame[HEADER_LEN..].as_chunks::<COUNTER_LEN>();
+        if !Self::rows_are_coherent(slots, total, saturated) {
+            return Err(CodecError::Parameters);
+        }
+
+        self.total = total;
+        self.saturated = saturated;
+        let counters = self.rows.iter_mut().flatten();
+        for (slot, counter) in slots.iter().zip(counters) {
+            *counter = u32::from_le_bytes(*slot);
+        }
+        Ok(())
+    }
+
+    /// Whether the counters carry the total the frame claims, as every sketch
+    /// built by adding items does: each row takes every count added, so it sums
+    /// to the total exactly. Only a counter stopped at [`u32::MAX`] can leave a
+    /// row short of it, and such a sketch says so in its saturation flag.
+    fn rows_are_coherent(slots: &[[u8; COUNTER_LEN]], total: u64, saturated: bool) -> bool {
+        slots.chunks(W).all(|row| {
+            let counted: u64 = row
+                .iter()
+                .map(|slot| u64::from(u32::from_le_bytes(*slot)))
+                .sum();
+            if saturated {
+                counted <= total
+            } else {
+                counted == total
+            }
+        })
+    }
+}
+
+impl<const W: usize, const D: usize> Default for CountMinSketch<W, D> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+#[cfg(test)]
+mod tests;

data/ext/probabilistic_machines_native/core/src/cuckoo.rs

@@ -0,0 +1,283 @@
+//! Approximate membership with deletions, which a Bloom filter cannot do.
+
+use crate::hash::{hash64, mix64, u32_low_u8, u64_low_u8, u64_low_u16};
+use crate::{CodecError, TAG_CUCKOO};
+
+const SLOTS: usize = 4;
+const EMPTY: u8 = 0;
+const MAX_KICKS: usize = 64;
+
+/// The filter has no room left for this item.
+///
+/// Cuckoo hashing gives up before the table is literally full; expect that
+/// past about 95% of [`CuckooFilter::CAPACITY`]. The filter is unchanged when
+/// an insert fails.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct Full;
+
+impl core::fmt::Display for Full {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        f.write_str("cuckoo filter full")
+    }
+}
+
+impl core::error::Error for Full {}
+
+/// Answers "have I seen this?" in `B * 4` bytes, and unlike a Bloom filter it
+/// can forget.
+///
+/// `B` is the bucket count, a power of two up to 65536; each bucket holds four
+/// 8-bit fingerprints, so the filter stores up to `4 * B` items in `4 * B`
+/// bytes. [`contains`](Self::contains) never says no to an item that was
+/// inserted; it says yes to an item that was not with probability at most
+/// `2 * 4 / 255`, about 3%.
+///
+/// [`remove`](Self::remove) must only be called for items that were inserted:
+/// deleting an item that merely collides with one takes the other item's
+/// fingerprint with it.
+///
+/// ```
+/// use probabilistic_machines::CuckooFilter;
+///
+/// let mut seen = CuckooFilter::<1024>::new();
+/// assert_eq!(seen.insert(b"device-9"), Ok(()));
+/// assert!(seen.contains(b"device-9"));
+///
+/// assert!(seen.remove(b"device-9"));
+/// assert!(!seen.contains(b"device-9"));
+/// ```
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct CuckooFilter<const B: usize> {
+    buckets: [[u8; SLOTS]; B],
+    items: usize,
+}
+
+impl<const B: usize> CuckooFilter<B> {
+    const BUCKET_MASK: usize = {
+        assert!(B.is_power_of_two(), "bucket count must be a power of two");
+        assert!(B >= 2, "bucket count must be at least 2");
+        assert!(B <= 1 << 16, "bucket count must be at most 65536");
+        B - 1
+    };
+
+    /// Fingerprint slots in the table.
+    pub const CAPACITY: usize = B * SLOTS;
+
+    /// Length of the [`encode`](Self::encode) frame.
+    pub const ENCODED_LEN: usize = 2 + B * SLOTS;
+
+    #[must_use]
+    pub const fn new() -> Self {
+        Self {
+            buckets: [[EMPTY; SLOTS]; B],
+            items: 0,
+        }
+    }
+
+    /// Items currently stored.
+    #[must_use]
+    pub const fn len(&self) -> usize {
+        self.items
+    }
+
+    #[must_use]
+    pub const fn is_empty(&self) -> bool {
+        self.items == 0
+    }
+
+    /// # Errors
+    ///
+    /// [`Full`] when the item's two buckets and 64 evictions all failed to
+    /// find a slot. The filter is left as it was.
+    pub fn insert(&mut self, item: &[u8]) -> Result<(), Full> {
+        self.insert_hash(hash64(item))
+    }
+
+    /// Inserts an item the caller already hashed.
+    ///
+    /// # Errors
+    ///
+    /// [`Full`] when the item's two buckets and 64 evictions all failed to
+    /// find a slot. The filter is left as it was.
+    pub fn insert_hash(&mut self, hash: u64) -> Result<(), Full> {
+        let fingerprint = fingerprint(hash);
+        let primary = Self::bucket(hash);
+        if self.place(primary, fingerprint) {
+            return Ok(());
+        }
+        let alternate = Self::alternate(primary, fingerprint);
+        if self.place(alternate, fingerprint) {
+            return Ok(());
+        }
+
+        let mut path = [0_usize; MAX_KICKS];
+        let mut index = alternate;
+        let mut carried = fingerprint;
+        let mut walk = mix64(hash);
+
+        for step in &mut path {
+            walk = mix64(walk);
+            let slot = usize::from(u64_low_u16(walk)) % SLOTS;
+            core::mem::swap(&mut carried, &mut self.buckets[index][slot]);
+            *step = index * SLOTS + slot;
+            index = Self::alternate(index, carried);
+            if self.place(index, carried) {
+                return Ok(());
+            }
+        }
+
+        for step in path.iter().rev() {
+            core::mem::swap(&mut carried, &mut self.buckets[step / SLOTS][step % SLOTS]);
+        }
+        Err(Full)
+    }
+
+    #[must_use]
+    pub fn contains(&self, item: &[u8]) -> bool {
+        self.contains_hash(hash64(item))
+    }
+
+    /// Queries an item the caller already hashed.
+    #[must_use]
+    pub fn contains_hash(&self, hash: u64) -> bool {
+        let fingerprint = fingerprint(hash);
+        let primary = Self::bucket(hash);
+        let alternate = Self::alternate(primary, fingerprint);
+        self.buckets[primary].contains(&fingerprint)
+            || self.buckets[alternate].contains(&fingerprint)
+    }
+
+    /// Drops one copy of an item, reporting whether a fingerprint was cleared.
+    pub fn remove(&mut self, item: &[u8]) -> bool {
+        self.remove_hash(hash64(item))
+    }
+
+    /// Drops an item the caller already hashed.
+    pub fn remove_hash(&mut self, hash: u64) -> bool {
+        let fingerprint = fingerprint(hash);
+        let primary = Self::bucket(hash);
+        let alternate = Self::alternate(primary, fingerprint);
+        self.clear(primary, fingerprint) || self.clear(alternate, fingerprint)
+    }
+
+    /// Writes the filter into `out`, returning the number of bytes written.
+    ///
+    /// # Errors
+    ///
+    /// [`CodecError::Truncated`] when `out` is shorter than
+    /// [`ENCODED_LEN`](Self::ENCODED_LEN).
+    pub fn encode(&self, out: &mut [u8]) -> Result<usize, CodecError> {
+        let Some(frame) = out.get_mut(..Self::ENCODED_LEN) else {
+            return Err(CodecError::Truncated);
+        };
+        self.write_frame(frame);
+        Ok(Self::ENCODED_LEN)
+    }
+
+    /// The filter's frame, exactly [`ENCODED_LEN`](Self::ENCODED_LEN) bytes.
+    #[cfg(feature = "alloc")]
+    #[must_use]
+    pub fn encoded(&self) -> alloc::boxed::Box<[u8]> {
+        let mut frame = alloc::vec![0_u8; Self::ENCODED_LEN].into_boxed_slice();
+        self.write_frame(&mut frame);
+        frame
+    }
+
+    /// Fills a frame of exactly [`ENCODED_LEN`](Self::ENCODED_LEN) bytes.
+    fn write_frame(&self, frame: &mut [u8]) {
+        frame[0] = TAG_CUCKOO;
+        frame[1] = u32_low_u8(B.trailing_zeros());
+
+        let fingerprints = self.buckets.iter().flatten();
+        for (slot, value) in frame[2..].iter_mut().zip(fingerprints) {
+            *slot = *value;
+        }
+    }
+
+    /// Reads a filter written by [`encode`](Self::encode).
+    ///
+    /// # Errors
+    ///
+    /// [`CodecError::Truncated`] when `bytes` is too short, [`CodecError::Tag`]
+    /// when it holds another structure, [`CodecError::Parameters`] when it was
+    /// written with a different bucket count.
+    pub fn decode(bytes: &[u8]) -> Result<Self, CodecError> {
+        let mut filter = Self::new();
+        filter.decode_into(bytes)?;
+        Ok(filter)
+    }
+
+    /// Reads a frame over this filter, for a caller that owns one already and
+    /// has no room on its stack for a second.
+    ///
+    /// The filter is untouched when the frame is rejected.
+    ///
+    /// # Errors
+    ///
+    /// The same as [`decode`](Self::decode).
+    pub fn decode_into(&mut self, bytes: &[u8]) -> Result<(), CodecError> {
+        let Some(frame) = bytes.get(..Self::ENCODED_LEN) else {
+            return Err(CodecError::Truncated);
+        };
+        if frame[0] != TAG_CUCKOO {
+            return Err(CodecError::Tag);
+        }
+        if frame[1] != u32_low_u8(B.trailing_zeros()) {
+            return Err(CodecError::Parameters);
+        }
+
+        self.items = 0;
+        let fingerprints = self.buckets.iter_mut().flatten();
+        for (value, byte) in fingerprints.zip(&frame[2..]) {
+            *value = *byte;
+            if *byte != EMPTY {
+                self.items += 1;
+            }
+        }
+        Ok(())
+    }
+
+    fn bucket(hash: u64) -> usize {
+        usize::from(u64_low_u16(hash)) & Self::BUCKET_MASK
+    }
+
+    fn alternate(index: usize, fingerprint: u8) -> usize {
+        let twist = usize::from(u64_low_u16(mix64(u64::from(fingerprint)))) & Self::BUCKET_MASK;
+        index ^ twist
+    }
+
+    fn place(&mut self, index: usize, fingerprint: u8) -> bool {
+        for slot in &mut self.buckets[index] {
+            if *slot == EMPTY {
+                *slot = fingerprint;
+                self.items += 1;
+                return true;
+            }
+        }
+        false
+    }
+
+    fn clear(&mut self, index: usize, fingerprint: u8) -> bool {
+        for slot in &mut self.buckets[index] {
+            if *slot == fingerprint {
+                *slot = EMPTY;
+                self.items -= 1;
+                return true;
+            }
+        }
+        false
+    }
+}
+
+impl<const B: usize> Default for CuckooFilter<B> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+fn fingerprint(hash: u64) -> u8 {
+    u64_low_u8(mix64(hash)).max(1)
+}
+
+#[cfg(test)]
+mod tests;

data/ext/probabilistic_machines_native/core/src/hash.rs

@@ -0,0 +1,88 @@
+//! The crate's only hash. Sketches index by bit slices of a 64-bit value, so
+//! they need avalanche behaviour, not cryptography.
+
+const PRIME: u64 = 0x9E37_79B9_7F4A_7C15;
+
+/// `splitmix64`'s finalizer: a bijection on [`u64`].
+///
+/// Every output bit depends on every input bit, and iterating it yields
+/// decorrelated values, which is how the sketches derive one index per row
+/// from a single hash.
+#[must_use]
+pub const fn mix64(value: u64) -> u64 {
+    let mut z = value.wrapping_add(PRIME);
+    z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
+    z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
+    z ^ (z >> 31)
+}
+
+/// Hashes bytes with the crate's default seed.
+#[must_use]
+pub const fn hash64(bytes: &[u8]) -> u64 {
+    hash64_with_seed(bytes, 0)
+}
+
+/// Hashes bytes, folding eight at a time through [`mix64`].
+///
+/// Two seeds give two independent hash families, which is what a filter needs
+/// when it wants fingerprints uncorrelated with bucket indices.
+#[must_use]
+pub const fn hash64_with_seed(bytes: &[u8], seed: u64) -> u64 {
+    let mut accumulator = seed ^ PRIME;
+    let mut consumed: u64 = 0;
+    let mut cursor = 0;
+
+    while cursor + 8 <= bytes.len() {
+        let lane = u64::from_le_bytes([
+            bytes[cursor],
+            bytes[cursor + 1],
+            bytes[cursor + 2],
+            bytes[cursor + 3],
+            bytes[cursor + 4],
+            bytes[cursor + 5],
+            bytes[cursor + 6],
+            bytes[cursor + 7],
+        ]);
+        accumulator = (accumulator ^ mix64(lane))
+            .rotate_left(27)
+            .wrapping_mul(PRIME);
+        cursor += 8;
+        consumed += 8;
+    }
+
+    let mut tail: u64 = 0;
+    let mut shift = 0;
+    while cursor < bytes.len() {
+        tail |= (bytes[cursor] as u64) << shift;
+        shift += 8;
+        cursor += 1;
+        consumed += 1;
+    }
+
+    mix64(mix64(accumulator ^ consumed.wrapping_mul(PRIME)) ^ mix64(tail))
+}
+
+pub(crate) const fn u64_low_u16(value: u64) -> u16 {
+    let [low, high, ..] = value.to_le_bytes();
+    u16::from_le_bytes([low, high])
+}
+
+pub(crate) const fn u64_low_u8(value: u64) -> u8 {
+    value.to_le_bytes()[0]
+}
+
+pub(crate) const fn u32_low_u8(value: u32) -> u8 {
+    value.to_le_bytes()[0]
+}
+
+pub(crate) const fn usize_low_u8(value: usize) -> u8 {
+    value.to_le_bytes()[0]
+}
+
+pub(crate) const fn usize_low_u16(value: usize) -> u16 {
+    let [low, high, ..] = value.to_le_bytes();
+    u16::from_le_bytes([low, high])
+}
+
+#[cfg(test)]
+mod tests;