feat(heap): estimated_memory_bytes for HHHeap and WeightedFrequency

CHANGELOG.md

@@ -23,6 +23,12 @@ signals a backwards-compatible change.
   into the rescoring instead of copying them, and scores them with one
   buffer, about halving its cost at capacity 100; the retained candidates
   and their scores are unchanged.
+  `WeightedFrequency::estimated_memory_bytes` and
+  `HHHeap::estimated_memory_bytes` give the bytes a state of a given shape
+  holds once its top-k heap is full of keys of a given shape, so a memory
+  model can price one before building it. A retained candidate's identity
+  tuple is allocated at its exact length (it held four slots for a
+  one-value tuple, 72 bytes per candidate more than needed).
 
 - **`Vector3D<T>`**, filled in from a stub. A `rows x cols` grid whose every
   cell is a contiguous `depth`-length record, addressed by the same

src/common/heap.rs

@@ -167,6 +167,30 @@ impl HHHeap {
         self.k
     }
 
+    /// Bytes a heap built by [`Self::new`]`(k)` holds once `k` distinct keys
+    /// shaped like `key` are resident: the struct, the heap array of
+    /// [`HHItem`]s, the parallel digest array, the digest index and each
+    /// resident's owned key bytes (a `String` or `Bytes` key's length; 0 for a
+    /// number).
+    ///
+    /// The arrays and the index reserve `min(k, 1024)` up front and double
+    /// from there, so for `k` above 1,024 the arrays hold the next doubling
+    /// at or above `k`. Evictions leave tombstones that can make the index
+    /// rehash into twice its buckets, which this count leaves out. The index
+    /// size follows std's `HashMap` table layout on the build target.
+    ///
+    /// A memory model calls this to price a top-k heap before building one.
+    pub fn estimated_memory_bytes(k: usize, key: &DataInput) -> usize {
+        let reserved = k.min(PREALLOCATED_SLOTS);
+        let array_slots = vec_slots_after_pushes(reserved, k, size_of::<HHItem>());
+        let digest_slots = vec_slots_after_pushes(reserved, k, size_of::<u64>());
+        size_of::<Self>()
+            + array_slots * size_of::<HHItem>()
+            + digest_slots * size_of::<u64>()
+            + hash_index_bytes::<(u64, Slot)>(reserved, k)
+            + k * owned_key_bytes(key)
+    }
+
     // -- index maintenance ---------------------------------------------------
 
     fn lookup(&self, slot: u64, matches: impl Fn(&HHItem) -> bool) -> Option<usize> {
@@ -301,6 +325,99 @@ fn drop_entry(positions: &mut Index, digest: u64, idx: usize) {
     }
 }
 
+/// Bytes a key shaped like `key` owns once [`input_to_owned`] copies it into
+/// a heap entry: a string's or a byte slice's length, 0 for a number.
+pub(crate) fn owned_key_bytes(key: &DataInput) -> usize {
+    match key {
+        DataInput::Str(s) => s.len(),
+        DataInput::String(s) => s.len(),
+        DataInput::Bytes(b) => b.len(),
+        _ => 0,
+    }
+}
+
+/// Element slots a `Vec` holds after growing from `reserved` slots to `len`
+/// elements one push at a time: std's amortized growth doubles the capacity
+/// (at least 4 slots, 8 for 1-byte elements, 1 above 1 KiB).
+pub(crate) fn vec_slots_after_pushes(reserved: usize, len: usize, element_bytes: usize) -> usize {
+    let minimum = match element_bytes {
+        1 => 8,
+        2..=1024 => 4,
+        _ => 1,
+    };
+    let mut slots = reserved;
+    while slots < len {
+        slots = (slots * 2).max(slots + 1).max(minimum);
+    }
+    slots
+}
+
+/// Width of std's `HashMap` control group on this target: 16 with SSE2 on
+/// x86, 8 elsewhere.
+const HASH_GROUP_WIDTH: usize = if cfg!(all(
+    any(target_arch = "x86", target_arch = "x86_64"),
+    target_feature = "sse2"
+)) {
+    16
+} else {
+    8
+};
+
+/// Buckets std's `HashMap` allocates for `capacity` entries of `entry_bytes`.
+fn hash_buckets_for(capacity: usize, entry_bytes: usize) -> usize {
+    if capacity == 0 {
+        return 0;
+    }
+    if capacity < 15 {
+        // Small tables round up to a minimum capacity so the control-group
+        // padding is not wasted.
+        let minimum = match (HASH_GROUP_WIDTH, entry_bytes) {
+            (16, 0..=1) => 14,
+            (16, 2..=3) => 7,
+            (8, 0..=1) => 7,
+            _ => 3,
+        };
+        return match minimum.max(capacity) {
+            0..=3 => 4,
+            4..=7 => 8,
+            _ => 16,
+        };
+    }
+    (capacity * 8 / 7).next_power_of_two()
+}
+
+/// Entries a table of `buckets` holds before it grows (7/8 of its buckets).
+fn hash_capacity_of(buckets: usize) -> usize {
+    if buckets <= 8 {
+        buckets - 1
+    } else {
+        buckets / 8 * 7
+    }
+}
+
+/// Bytes std's `HashMap` allocates for entries of type `T` when created with
+/// capacity `reserved` and then filled with `len` distinct keys: the entry
+/// slots padded to the control alignment, one control byte per bucket and one
+/// trailing control group. A full table grows to the buckets for one more
+/// entry than it holds.
+pub(crate) fn hash_index_bytes<T>(reserved: usize, len: usize) -> usize {
+    let entry_bytes = size_of::<T>();
+    let mut buckets = hash_buckets_for(reserved, entry_bytes);
+    while (buckets == 0 && len > 0) || (buckets > 0 && hash_capacity_of(buckets) < len) {
+        let held = if buckets == 0 {
+            0
+        } else {
+            hash_capacity_of(buckets)
+        };
+        buckets = hash_buckets_for(held + 1, entry_bytes);
+    }
+    if buckets == 0 {
+        return 0;
+    }
+    let control_align = align_of::<T>().max(HASH_GROUP_WIDTH);
+    (buckets * entry_bytes).next_multiple_of(control_align) + buckets + HASH_GROUP_WIDTH
+}
+
 /// Every element sits in heap order beneath the parent the array layout gives
 /// it, with the parent index computed here rather than taken from the heap.
 #[cfg(test)]

src/sketches/weighted_frequency.rs

@@ -55,17 +55,21 @@ impl Ord for Candidate {
 }
 /// Canonical identity tuple and its hash key: floating identities must be
 /// finite, and signed zeros share one identity.
+/// The identity is allocated at its exact length, so a retained candidate
+/// holds no spare slots.
 fn normalize(values: &[FrequencyIdentity]) -> Result<(Vec<FrequencyIdentity>, Vec<u8>), Error> {
-    let identity = values
-        .iter()
-        .map(|value| match value {
-            FrequencyIdentity::Float64(v) if !v.is_finite() => Err(Error::Invalid(
-                "non-finite weighted frequency identity".into(),
-            )),
-            FrequencyIdentity::Float64(v) if *v == 0.0 => Ok(FrequencyIdentity::Float64(0.0)),
-            value => Ok(value.clone()),
-        })
-        .collect::<Result<Vec<_>, _>>()?;
+    let mut identity = Vec::with_capacity(values.len());
+    for value in values {
+        identity.push(match value {
+            FrequencyIdentity::Float64(v) if !v.is_finite() => {
+                return Err(Error::Invalid(
+                    "non-finite weighted frequency identity".into(),
+                ));
+            }
+            FrequencyIdentity::Float64(v) if *v == 0.0 => FrequencyIdentity::Float64(0.0),
+            value => value.clone(),
+        });
+    }
     let key = bincode::serialize(&identity).map_err(|e| Error::Update(e.to_string()))?;
     Ok((identity, key))
 }
@@ -127,6 +131,44 @@ impl WeightedFrequency {
             .map(|_| ())
             .ok_or_else(|| Error::Invalid("invalid weighted frequency dimensions".into()))
     }
+    /// Bytes a state of shape `width × depth` holds once it retains
+    /// `capacity` candidates whose identity tuples are shaped like `identity`:
+    /// the struct, the `width × depth` 8-byte cells, the candidate heap's
+    /// slots, and per retained candidate its identity tuple (with a `Utf8`
+    /// value's bytes) and its encoded hash key. Only the identity's shape
+    /// matters (its variants and string lengths), not its values.
+    ///
+    /// The candidate heap is rebuilt on every update by pushing up to
+    /// `capacity + 1` candidates, so it holds `max(4, (capacity + 1)
+    /// rounded up to a power of two)` slots, and during an update the
+    /// previous heap's slots are held as well. A memory model calls this to
+    /// price a top-k frequency state before building one.
+    pub fn estimated_memory_bytes(
+        width: usize,
+        depth: usize,
+        capacity: usize,
+        identity: &[FrequencyIdentity],
+    ) -> usize {
+        let cells = width.saturating_mul(depth).saturating_mul(size_of::<f64>());
+        if capacity == 0 {
+            return size_of::<Self>() + cells;
+        }
+        let heap_slots =
+            crate::common::heap::vec_slots_after_pushes(0, capacity + 1, size_of::<Candidate>());
+        let identity_bytes = size_of_val(identity)
+            + identity
+                .iter()
+                .map(|value| match value {
+                    FrequencyIdentity::Utf8(text) => text.len(),
+                    _ => 0,
+                })
+                .sum::<usize>();
+        let key_bytes = bincode::serialized_size(identity).map_or(0, |bytes| bytes as usize);
+        size_of::<Self>()
+            + cells
+            + heap_slots * size_of::<Candidate>()
+            + capacity * (identity_bytes + key_bytes)
+    }
     pub fn new(
         algorithm: FrequencyAlgorithm,
         width: usize,

tests/heap_memory_bytes.rs

@@ -0,0 +1,128 @@
+//! `HHHeap::estimated_memory_bytes` and `WeightedFrequency::estimated_memory_bytes`
+//! checked against `size_of` and against the bytes the heap actually holds,
+//! counted by a global allocator. One `#[test]` so no other test allocates
+//! while a count is taken.
+
+use asap_sketchlib::{
+    DataInput, FrequencyAlgorithm, FrequencyIdentity, HHHeap, HHItem, WeightedFrequency,
+};
+use std::alloc::{GlobalAlloc, Layout, System};
+use std::sync::atomic::{AtomicIsize, Ordering};
+
+/// Counts the bytes currently allocated through it.
+struct CountingAllocator;
+
+static LIVE_BYTES: AtomicIsize = AtomicIsize::new(0);
+
+unsafe impl GlobalAlloc for CountingAllocator {
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        LIVE_BYTES.fetch_add(layout.size() as isize, Ordering::SeqCst);
+        unsafe { System.alloc(layout) }
+    }
+    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+        LIVE_BYTES.fetch_sub(layout.size() as isize, Ordering::SeqCst);
+        unsafe { System.dealloc(ptr, layout) }
+    }
+    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
+        LIVE_BYTES.fetch_add(new_size as isize - layout.size() as isize, Ordering::SeqCst);
+        unsafe { System.realloc(ptr, layout, new_size) }
+    }
+}
+
+#[global_allocator]
+static ALLOCATOR: CountingAllocator = CountingAllocator;
+
+fn live_bytes() -> isize {
+    LIVE_BYTES.load(Ordering::SeqCst)
+}
+
+/// Heap bytes `build` leaves allocated, plus the size of the value it returns.
+fn measured_bytes<T>(build: impl FnOnce() -> T) -> (T, usize) {
+    let before = live_bytes();
+    let value = build();
+    let heap = (live_bytes() - before) as usize;
+    (value, heap + std::mem::size_of::<T>())
+}
+
+fn int_identity(i: usize) -> Vec<FrequencyIdentity> {
+    vec![FrequencyIdentity::Int64(i as i64)]
+}
+
+fn text_identity(i: usize) -> Vec<FrequencyIdentity> {
+    vec![
+        FrequencyIdentity::Utf8(format!("{i:016}")),
+        FrequencyIdentity::Bool(i % 2 == 0),
+    ]
+}
+
+#[test]
+fn estimated_memory_bytes_match_size_of_and_the_allocation() {
+    // size_of: a resident is one `HHItem` in the heap array, one `u64` digest
+    // and one index entry; a heap of capacity 0 allocates nothing.
+    assert_eq!(
+        HHHeap::estimated_memory_bytes(0, &DataInput::U64(0)),
+        std::mem::size_of::<HHHeap>()
+    );
+    assert!(
+        HHHeap::estimated_memory_bytes(1024, &DataInput::U64(0))
+            > std::mem::size_of::<HHHeap>() + 1024 * (std::mem::size_of::<HHItem>() + 8)
+    );
+    assert_eq!(
+        WeightedFrequency::estimated_memory_bytes(64, 3, 0, &int_identity(0)),
+        std::mem::size_of::<WeightedFrequency>() + 64 * 3 * 8
+    );
+
+    // HHHeap filled with `k` distinct keys (no eviction), for capacities on
+    // both sides of the 1,024-slot up-front reservation.
+    for k in [1usize, 3, 7, 8, 100, 1024, 2000, 5000] {
+        let (_heap, bytes) = measured_bytes(|| {
+            let mut heap = HHHeap::new(k);
+            for i in 0..k as u64 {
+                heap.update(&DataInput::U64(i), i as i64 + 1);
+            }
+            heap
+        });
+        assert_eq!(
+            bytes,
+            HHHeap::estimated_memory_bytes(k, &DataInput::U64(0)),
+            "u64 keys, k = {k}"
+        );
+
+        let (_heap, bytes) = measured_bytes(|| {
+            let mut heap = HHHeap::new(k);
+            for i in 0..k as u64 {
+                let key = [&[3u8][..], &i.to_le_bytes()].concat();
+                heap.update(&DataInput::Bytes(&key), i as i64 + 1);
+            }
+            heap
+        });
+        assert_eq!(
+            bytes,
+            HHHeap::estimated_memory_bytes(k, &DataInput::Bytes(&[0; 9])),
+            "9-byte keys, k = {k}"
+        );
+    }
+
+    // WeightedFrequency after four times `capacity` distinct identities, so the
+    // candidate heap is full and has turned keys away.
+    for capacity in [1usize, 7, 64, 100, 1000] {
+        for (name, identity) in [
+            ("Int64", int_identity as fn(usize) -> Vec<FrequencyIdentity>),
+            ("Utf8 + Bool", text_identity),
+        ] {
+            let (_state, bytes) = measured_bytes(|| {
+                let mut state =
+                    WeightedFrequency::new(FrequencyAlgorithm::Cms, 64, 3, capacity).unwrap();
+                for i in 0..4 * capacity {
+                    state.update(&identity(i), 1.0).unwrap();
+                }
+                state
+            });
+            assert_eq!(
+                bytes,
+                WeightedFrequency::estimated_memory_bytes(64, 3, capacity, &identity(0)),
+                "{name} identities, capacity = {capacity}"
+            );
+        }
+    }
+}