probabilistic_machines 0.1.1

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

@@ -0,0 +1,325 @@
+//! Which keys dominate a stream, in `K` counters instead of one per key.
+
+use crate::hash::usize_low_u16;
+use crate::{CodecError, TAG_TOPK, read8};
+
+const HEADER_LEN: usize = 13;
+const RECORD_LEN: usize = 24;
+/// Where a record's count and overestimate begin, after its key.
+const COUNT_AT: usize = 8;
+const OVERESTIMATE_AT: usize = 16;
+
+/// A tracked key, its counter, and how much of that counter it may have
+/// inherited from keys evicted before it or from a merged-in stream that
+/// tracked it no longer.
+///
+/// The true count lies in `count - overestimate ..= count`.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct Hitter<T> {
+    pub key: T,
+    pub count: u64,
+    pub overestimate: u64,
+}
+
+/// Space-Saving over `K` counters: whatever is frequent ends up tracked.
+///
+/// After `n` observations every key that occurred more than `n / K` times is in
+/// the table, and no tracked key's counter is below its true count. Lookups and
+/// insertions scan all `K` counters, which is the trade for needing no heap and
+/// no key ordering: keep `K` in the dozens-to-hundreds range.
+///
+/// ```
+/// use probabilistic_machines::HeavyHitters;
+///
+/// let mut traffic = HeavyHitters::<u32, 8>::new();
+/// for request in 0..1_000_u32 {
+///     traffic.add(request % 97);
+/// }
+/// for _ in 0..5_000 {
+///     traffic.add(7);
+/// }
+///
+/// assert_eq!(traffic.observed(), 6_000);
+/// assert!(traffic.count(&7) >= 5_000);
+/// ```
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct HeavyHitters<T, const K: usize> {
+    slots: [Option<Hitter<T>>; K],
+    observed: u64,
+}
+
+impl<T: Copy + Eq, const K: usize> HeavyHitters<T, K> {
+    #[must_use]
+    pub const fn new() -> Self {
+        Self {
+            slots: [const { None }; K],
+            observed: 0,
+        }
+    }
+
+    /// Observations fed in, saturating at [`u64::MAX`].
+    #[must_use]
+    pub const fn observed(&self) -> u64 {
+        self.observed
+    }
+
+    /// Keys currently tracked, at most `K`.
+    #[must_use]
+    pub fn len(&self) -> usize {
+        self.slots.iter().flatten().count()
+    }
+
+    #[must_use]
+    pub fn is_empty(&self) -> bool {
+        self.slots.iter().all(Option::is_none)
+    }
+
+    /// Every tracked key, in no particular order.
+    pub fn hitters(&self) -> impl Iterator<Item = &Hitter<T>> {
+        self.slots.iter().flatten()
+    }
+
+    /// Upper bound on how often `key` occurred, 0 when it is not tracked.
+    #[must_use]
+    pub fn count(&self, key: &T) -> u64 {
+        self.hitters()
+            .find(|hitter| hitter.key == *key)
+            .map_or(0, |hitter| hitter.count)
+    }
+
+    /// The most frequent tracked key.
+    #[must_use]
+    pub fn top(&self) -> Option<&Hitter<T>> {
+        self.hitters().max_by_key(|hitter| hitter.count)
+    }
+
+    pub fn add(&mut self, key: T) {
+        self.add_count(key, 1);
+    }
+
+    pub fn add_count(&mut self, key: T, count: u64) {
+        self.observed = self.observed.saturating_add(count);
+        self.absorb(key, count, 0);
+    }
+
+    /// Folds another table in. Counts of keys tracked by both are exact sums.
+    /// A key only one side tracked may have occurred in the other stream too,
+    /// up to that side's eviction floor, so it takes the floor on both its
+    /// bounds: the counter still covers the true count and the interval
+    /// `count - overestimate ..= count` still holds it.
+    pub fn merge(&mut self, other: &Self) {
+        let floor = other.floor();
+        for tracked in self.slots.iter_mut().flatten() {
+            if other.hitters().all(|hitter| hitter.key != tracked.key) {
+                tracked.count = tracked.count.saturating_add(floor);
+                tracked.overestimate = tracked.overestimate.saturating_add(floor);
+            }
+        }
+        for hitter in other.hitters() {
+            self.absorb(hitter.key, hitter.count, hitter.overestimate);
+        }
+        self.observed = self.observed.saturating_add(other.observed);
+    }
+
+    /// How often an untracked key may have occurred: nothing while a counter is
+    /// still free, otherwise the weakest counter.
+    fn floor(&self) -> u64 {
+        if self.slots.iter().any(Option::is_none) {
+            return 0;
+        }
+        self.hitters().map(|hitter| hitter.count).min().unwrap_or(0)
+    }
+
+    fn absorb(&mut self, key: T, count: u64, overestimate: u64) {
+        if let Some(tracked) = self
+            .slots
+            .iter_mut()
+            .flatten()
+            .find(|hitter| hitter.key == key)
+        {
+            tracked.count = tracked.count.saturating_add(count);
+            tracked.overestimate = tracked.overestimate.saturating_add(overestimate);
+            return;
+        }
+
+        if let Some(free) = self.slots.iter_mut().find(|slot| slot.is_none()) {
+            *free = Some(Hitter {
+                key,
+                count,
+                overestimate,
+            });
+            return;
+        }
+
+        let Some(weakest) = self
+            .slots
+            .iter_mut()
+            .flatten()
+            .min_by_key(|hitter| hitter.count)
+        else {
+            return;
+        };
+        let floor = weakest.count;
+        *weakest = Hitter {
+            key,
+            count: floor.saturating_add(count),
+            overestimate: floor.saturating_add(overestimate),
+        };
+    }
+}
+
+impl<const K: usize> HeavyHitters<u64, K> {
+    /// `K` as the frame's shape field.
+    const COUNTERS: [u8; 2] = {
+        assert!(K >= 1, "a table needs at least one counter");
+        assert!(K < 1 << 16, "a table needs fewer than 65536 counters");
+        usize_low_u16(K).to_le_bytes()
+    };
+
+    /// Length of the [`encode`](Self::encode) frame.
+    pub const ENCODED_LEN: usize = HEADER_LEN + K * RECORD_LEN;
+
+    /// Writes the table into `out`, returning the number of bytes written.
+    ///
+    /// The frame is the tag, `K`, the observation count, the number of tracked
+    /// keys and `K` records of key, count and overestimate. Records past the
+    /// tracked count are zero.
+    ///
+    /// # 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 table'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_TOPK;
+        frame[1..3].copy_from_slice(&Self::COUNTERS);
+        frame[3..11].copy_from_slice(&self.observed.to_le_bytes());
+        frame[11..13].copy_from_slice(&usize_low_u16(self.len()).to_le_bytes());
+        frame[HEADER_LEN..].fill(0);
+
+        let (slots, _) = frame[HEADER_LEN..].as_chunks_mut::<RECORD_LEN>();
+        for (slot, hitter) in slots.iter_mut().zip(self.hitters()) {
+            slot[..COUNT_AT].copy_from_slice(&hitter.key.to_le_bytes());
+            slot[COUNT_AT..OVERESTIMATE_AT].copy_from_slice(&hitter.count.to_le_bytes());
+            slot[OVERESTIMATE_AT..].copy_from_slice(&hitter.overestimate.to_le_bytes());
+        }
+    }
+
+    /// Reads a table 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 counter count, claims more tracked keys than
+    /// `K`, or holds records no stream could have left.
+    pub fn decode(bytes: &[u8]) -> Result<Self, CodecError> {
+        let mut table = Self::new();
+        table.decode_into(bytes)?;
+        Ok(table)
+    }
+
+    /// Reads a frame over this table, for a caller that owns one already and
+    /// has no room on its stack for a second.
+    ///
+    /// The table 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_TOPK {
+            return Err(CodecError::Tag);
+        }
+        if frame[1..3] != Self::COUNTERS {
+            return Err(CodecError::Parameters);
+        }
+        let tracked = usize::from(u16::from_le_bytes([frame[11], frame[12]]));
+        if tracked > K {
+            return Err(CodecError::Parameters);
+        }
+
+        let observed = u64::from_le_bytes(read8(&frame[3..11]));
+        let (slots, _) = frame[HEADER_LEN..].as_chunks::<RECORD_LEN>();
+        let records = &slots[..tracked];
+        if !records_are_coherent(records, observed) {
+            return Err(CodecError::Parameters);
+        }
+
+        self.observed = observed;
+        for (slot, record) in self.slots.iter_mut().zip(records) {
+            *slot = Some(hitter(record));
+        }
+        self.slots[tracked..].fill(None);
+        Ok(())
+    }
+}
+
+/// A record as the hitter it stands for.
+fn hitter(record: &[u8; RECORD_LEN]) -> Hitter<u64> {
+    Hitter {
+        key: u64::from_le_bytes(read8(&record[..COUNT_AT])),
+        count: u64::from_le_bytes(read8(&record[COUNT_AT..OVERESTIMATE_AT])),
+        overestimate: u64::from_le_bytes(read8(&record[OVERESTIMATE_AT..])),
+    }
+}
+
+/// Whether the records are ones Space-Saving could have tracked: a counter
+/// stands for at least `count - overestimate` occurrences, which cannot outrun
+/// the observations they came from, an interval runs upwards, and a key holds
+/// one counter.
+///
+/// The counter itself may pass the observation count: a merge lifts a key the
+/// other table evicted to that table's eviction floor and widens the
+/// overestimate with it, so the two tables' floors can stack above the two
+/// streams' length while the occurrences the interval claims stay inside it.
+fn records_are_coherent(records: &[[u8; RECORD_LEN]], observed: u64) -> bool {
+    records.iter().enumerate().all(|(index, record)| {
+        let tracked = hitter(record);
+        tracked.overestimate <= tracked.count
+            && tracked.count - tracked.overestimate <= observed
+            && records[..index]
+                .iter()
+                .all(|earlier| hitter(earlier).key != tracked.key)
+    })
+}
+
+impl<T: Copy + Eq, const K: usize> Default for HeavyHitters<T, K> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+#[cfg(feature = "alloc")]
+impl<T: Copy + Eq, const K: usize> HeavyHitters<T, K> {
+    /// Tracked keys, most frequent first.
+    #[must_use]
+    pub fn ranked(&self) -> alloc::vec::Vec<Hitter<T>> {
+        let mut ranked: alloc::vec::Vec<Hitter<T>> = self.hitters().copied().collect();
+        ranked.sort_unstable_by_key(|hitter| core::cmp::Reverse(hitter.count));
+        ranked
+    }
+}
+
+#[cfg(test)]
+mod tests;

data/ext/probabilistic_machines_native/extconf.rb

@@ -0,0 +1,3 @@
+# frozen_string_literal: true
+
+require_relative 'ffi/extconf'

data/ext/probabilistic_machines_native/ffi/Cargo.toml

@@ -0,0 +1,20 @@
+[package]
+name = "probabilistic_machines_native"
+version = "0.1.1"
+edition.workspace = true
+rust-version.workspace = true
+authors = ["Abdelkader Boudih <oss@seuros.com>"]
+license = "MIT"
+description = "Ruby bindings for the probabilistic-machines sketches"
+publish = false
+
+[lib]
+crate-type = ["cdylib"]
+
+[lints]
+workspace = true
+
+[dependencies]
+probabilistic-machines = { workspace = true, features = ["std"] }
+magnus = { version = ">= 0.9.3", default-features = false, features = ["bytes", "io"] }
+bytes = "1"

data/ext/probabilistic_machines_native/ffi/extconf.rb

@@ -0,0 +1,21 @@
+# frozen_string_literal: true
+
+require 'mkmf'
+require 'pathname'
+require 'rb_sys/mkmf'
+
+create_rust_makefile('probabilistic_machines_native/probabilistic_machines_native') do |r|
+  r.ext_dir = Pathname(__dir__).relative_path_from(Pathname(Dir.pwd)).to_s
+  r.profile = ENV.fetch('RB_SYS_CARGO_PROFILE', :release).to_sym
+end
+
+# rake-compiler builds from tmp/, where the generated relative manifest path no
+# longer resolves.
+makefile = File.join(Dir.pwd, 'Makefile')
+File.write(
+  makefile,
+  File.read(makefile).gsub(
+    /^RB_SYS_CARGO_MANIFEST_DIR \?=.*$/,
+    "RB_SYS_CARGO_MANIFEST_DIR ?= #{File.expand_path(__dir__)}"
+  )
+)

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

@@ -0,0 +1,188 @@
+//! `ProbabilisticMachines::CountMinSketch`: how often each item was seen.
+
+use std::cell::RefCell;
+
+use bytes::Bytes;
+use magnus::scan_args::{get_kwargs, scan_args};
+use magnus::typed_data::Obj;
+use magnus::{Error, RHash, RModule, RString, Ruby, Value, function, method, prelude::*};
+use probabilistic_machines::{CodecError, CountMinSketch};
+
+use crate::{Whole, codec_error, mismatch, optional_kwarg, unsupported};
+
+/// Row widths [`Sketch`] is monomorphised for.
+///
+/// Capped so the widest sketch is 128 KiB. The core returns by value, so a
+/// wider shape could not be built on Ruby's 256 KiB fiber stack without an
+/// in-place constructor, and that needs `unsafe`.
+const WIDTHS: [usize; 2] = [1024, 4096];
+/// Row counts [`Sketch`] is monomorphised for.
+const DEPTHS: [usize; 2] = [4, 8];
+
+#[derive(Clone)]
+enum Sketch {
+    W1024D4(Box<CountMinSketch<1024, 4>>),
+    W1024D8(Box<CountMinSketch<1024, 8>>),
+    W4096D4(Box<CountMinSketch<4096, 4>>),
+    W4096D8(Box<CountMinSketch<4096, 8>>),
+}
+
+/// Runs `body` against whichever variant is held.
+macro_rules! sketch {
+    ($sketch:expr, |$inner:ident| $body:expr) => {
+        match $sketch {
+            Sketch::W1024D4($inner) => $body,
+            Sketch::W1024D8($inner) => $body,
+            Sketch::W4096D4($inner) => $body,
+            Sketch::W4096D8($inner) => $body,
+        }
+    };
+}
+
+#[magnus::wrap(
+    class = "ProbabilisticMachines::CountMinSketch",
+    free_immediately,
+    size
+)]
+struct Wrapper(RefCell<Sketch>);
+
+const fn width_of<const W: usize, const D: usize>(_: &CountMinSketch<W, D>) -> usize {
+    W
+}
+
+const fn depth_of<const W: usize, const D: usize>(_: &CountMinSketch<W, D>) -> usize {
+    D
+}
+
+/// Each shape is constructed behind its own call frame. Inlined, a debug build
+/// sizes one frame for every arm of the match below at once.
+#[inline(never)]
+fn boxed_new<const W: usize, const D: usize>() -> Box<CountMinSketch<W, D>> {
+    Box::new(CountMinSketch::new())
+}
+
+#[inline(never)]
+fn boxed_decode<const W: usize, const D: usize>(
+    ruby: &Ruby,
+    frame: &[u8],
+) -> Result<Box<CountMinSketch<W, D>>, Error> {
+    let mut sketch = boxed_new::<W, D>();
+    sketch
+        .decode_into(frame)
+        .map_err(|error| codec_error(ruby, error))?;
+    Ok(sketch)
+}
+
+impl Wrapper {
+    fn new(ruby: &Ruby, args: &[Value]) -> Result<Self, Error> {
+        let args = scan_args::<(), (), (), (), RHash, ()>(args)?;
+        let (Whole(width), Whole(depth)) = get_kwargs::<_, (Whole<usize>, Whole<usize>), (), ()>(
+            args.keywords,
+            &["width", "depth"],
+            &[],
+        )?
+        .required;
+        if !DEPTHS.contains(&depth) {
+            return Err(unsupported(ruby, "depth", depth, &DEPTHS));
+        }
+        let sketch = match (width, depth) {
+            (1024, 4) => Sketch::W1024D4(boxed_new::<1024, 4>()),
+            (1024, 8) => Sketch::W1024D8(boxed_new::<1024, 8>()),
+            (4096, 4) => Sketch::W4096D4(boxed_new::<4096, 4>()),
+            (4096, 8) => Sketch::W4096D8(boxed_new::<4096, 8>()),
+            _ => return Err(unsupported(ruby, "width", width, &WIDTHS)),
+        };
+        Ok(Self(RefCell::new(sketch)))
+    }
+
+    fn load(ruby: &Ruby, blob: Bytes) -> Result<Self, Error> {
+        let Some(shape) = blob.get(1..3) else {
+            return Err(codec_error(ruby, CodecError::Truncated));
+        };
+        let sketch = match (shape[0], shape[1]) {
+            (10, 4) => Sketch::W1024D4(boxed_decode::<1024, 4>(ruby, &blob)?),
+            (10, 8) => Sketch::W1024D8(boxed_decode::<1024, 8>(ruby, &blob)?),
+            (12, 4) => Sketch::W4096D4(boxed_decode::<4096, 4>(ruby, &blob)?),
+            (12, 8) => Sketch::W4096D8(boxed_decode::<4096, 8>(ruby, &blob)?),
+            _ => {
+                return Err(Error::new(
+                    ruby.exception_arg_error(),
+                    format!("blob was written with a shape outside {WIDTHS:?} x {DEPTHS:?}"),
+                ));
+            }
+        };
+        Ok(Self(RefCell::new(sketch)))
+    }
+
+    fn width(&self) -> usize {
+        sketch!(&*self.0.borrow(), |sketch| width_of(&**sketch))
+    }
+
+    fn depth(&self) -> usize {
+        sketch!(&*self.0.borrow(), |sketch| depth_of(&**sketch))
+    }
+
+    fn add(rb_self: &Self, args: &[Value]) -> Result<(), Error> {
+        let args = scan_args::<(Bytes,), (), (), (), RHash, ()>(args)?;
+        let (item,) = args.required;
+        let count =
+            optional_kwarg::<Whole<u32>>(args.keywords, "count")?.map_or(1, |Whole(count)| count);
+        sketch!(&mut *rb_self.0.borrow_mut(), |sketch| sketch
+            .add_count(&item, count));
+        Ok(())
+    }
+
+    fn count(&self, item: Bytes) -> u32 {
+        sketch!(&*self.0.borrow(), |sketch| sketch.estimate(&item))
+    }
+
+    fn saturated(&self) -> bool {
+        sketch!(&*self.0.borrow(), |sketch| sketch.saturated())
+    }
+
+    fn total(&self) -> u64 {
+        sketch!(&*self.0.borrow(), |sketch| sketch.total())
+    }
+
+    fn merge(ruby: &Ruby, rb_self: Obj<Self>, other: &Self) -> Result<Obj<Self>, Error> {
+        let source = other.0.borrow().clone();
+        match (&mut *rb_self.0.borrow_mut(), &source) {
+            (Sketch::W1024D4(target), Sketch::W1024D4(source)) => target.merge(source),
+            (Sketch::W1024D8(target), Sketch::W1024D8(source)) => target.merge(source),
+            (Sketch::W4096D4(target), Sketch::W4096D4(source)) => target.merge(source),
+            (Sketch::W4096D8(target), Sketch::W4096D8(source)) => target.merge(source),
+            _ => return Err(mismatch(ruby, "sketch")),
+        }
+        Ok(rb_self)
+    }
+
+    fn dump(ruby: &Ruby, rb_self: &Self) -> RString {
+        let frame = sketch!(&*rb_self.0.borrow(), |sketch| sketch.encoded());
+        ruby.str_from_slice(&frame)
+    }
+
+    fn inspect(&self) -> String {
+        format!(
+            "#<ProbabilisticMachines::CountMinSketch width={} depth={} total={}>",
+            self.width(),
+            self.depth(),
+            self.total()
+        )
+    }
+}
+
+pub fn define(ruby: &Ruby, module: RModule) -> Result<(), Error> {
+    let class = module.define_class("CountMinSketch", ruby.class_object())?;
+    class.define_singleton_method("new", function!(Wrapper::new, -1))?;
+    class.define_singleton_method("load", function!(Wrapper::load, 1))?;
+    class.define_method("width", method!(Wrapper::width, 0))?;
+    class.define_method("depth", method!(Wrapper::depth, 0))?;
+    class.define_method("add", method!(Wrapper::add, -1))?;
+    class.define_method("count", method!(Wrapper::count, 1))?;
+    class.define_method("total", method!(Wrapper::total, 0))?;
+    class.define_method("saturated?", method!(Wrapper::saturated, 0))?;
+    class.define_method("merge", method!(Wrapper::merge, 1))?;
+    class.define_method("dump", method!(Wrapper::dump, 0))?;
+    class.define_method("inspect", method!(Wrapper::inspect, 0))?;
+    Ok(())
+}