ddsketch/agent/

sketch.rs

//! Agent-specific DDSketch implementation.

use std::cmp::Ordering;

use datadog_protos::metrics::Dogsketch;
use ordered_float::OrderedFloat;
use smallvec::SmallVec;

use super::bin::Bin;
use super::bucket::Bucket;
use super::config::{
    Config, DDSKETCH_CONF_BIN_LIMIT, DDSKETCH_CONF_GAMMA_LN, DDSKETCH_CONF_GAMMA_V, DDSKETCH_CONF_NORM_BIAS,
    DDSKETCH_CONF_NORM_MIN,
};
use crate::canonical::mapping::LogarithmicMapping;
use crate::common::float_eq;

static SKETCH_CONFIG: Config = Config::new(
    DDSKETCH_CONF_BIN_LIMIT,
    DDSKETCH_CONF_GAMMA_V,
    DDSKETCH_CONF_GAMMA_LN,
    DDSKETCH_CONF_NORM_MIN,
    DDSKETCH_CONF_NORM_BIAS,
);
const MAX_BIN_WIDTH: u32 = u32::MAX;

#[cfg(feature = "serde")]
mod serde_helpers {
    use serde::{Deserialize, Deserializer};

    /// Deserializes an `f64` field, restoring JSON `null` to `NaN`.
    ///
    /// `serde_json` serializes non-finite floating-point values as `null`. This preserves the value's non-finite
    /// semantics when a sketch is round-tripped through JSON without fabricating a zero value.
    pub(super) fn null_as_nan<'de, D>(deserializer: D) -> Result<f64, D::Error>
    where
        D: Deserializer<'de>,
    {
        Ok(Option::<f64>::deserialize(deserializer)?.unwrap_or(f64::NAN))
    }
}

/// [DDSketch][ddsketch] implementation based on the [Datadog Agent][ddagent].
///
/// This implementation is subtly different from the open-source implementations of `DDSketch`, as Datadog made some
/// slight tweaks to configuration values and in-memory layout to optimize it for insertion performance within the
/// agent.
///
/// We've mimicked the agent version of `DDSketch` here in order to support a future where we can take sketches shipped
/// by the agent, handle them internally, merge them, and so on, without any loss of accuracy, eventually forwarding
/// them to Datadog ourselves.
///
/// As such, this implementation is constrained in the same ways: the configuration parameters can't be changed, the
/// collapsing strategy is fixed, and we support a limited number of methods for inserting into the sketch.
///
/// Importantly, we've a special function, again taken from the agent version, to allow us to interpolate histograms,
/// specifically our own aggregated histograms, into a sketch so that we can emit useful default quantiles, rather than
/// having to ship the buckets -- upper bound and count -- to a downstream system that might have no native way to do
/// the same thing, basically providing no value as they have no way to render useful data from them.
///
/// # Features
///
/// This crate exposes a single feature, `serde`, which enables serialization and deserialization of `DDSketch` with
/// `serde`. This feature isn't enabled by default, as it can be slightly risky to use. This is primarily due to the
/// fact that the format of `DDSketch` isn't promised to be stable over time. If you enable this feature, you should
/// take care to avoid storing serialized `DDSketch` data for long periods of time, as deserializing it in the future
/// may work but could lead to incorrect/unexpected behavior or issues with correctness.
///
/// [ddsketch]: https://www.vldb.org/pvldb/vol12/p2195-masson.pdf
/// [ddagent]: https://github.com/DataDog/datadog-agent
#[derive(Clone, Debug)]
#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
pub struct DDSketch {
    /// The bins within the sketch.
    bins: SmallVec<[Bin; 4]>,

    /// The number of observations within the sketch.
    count: u64,

    /// The minimum value of all observations within the sketch.
    #[cfg_attr(feature = "serde", serde(deserialize_with = "serde_helpers::null_as_nan"))]
    min: f64,

    /// The maximum value of all observations within the sketch.
    #[cfg_attr(feature = "serde", serde(deserialize_with = "serde_helpers::null_as_nan"))]
    max: f64,

    /// The sum of all observations within the sketch.
    #[cfg_attr(feature = "serde", serde(deserialize_with = "serde_helpers::null_as_nan"))]
    sum: f64,

    /// The average value of all observations within the sketch.
    #[cfg_attr(feature = "serde", serde(deserialize_with = "serde_helpers::null_as_nan"))]
    avg: f64,
}

impl DDSketch {
    /// Returns the canonical DDSketch mapping that aligns with the agent sketch's key space.
    pub fn remap_mapping() -> LogarithmicMapping {
        LogarithmicMapping::new_with_gamma_and_offset(DDSKETCH_CONF_GAMMA_V, f64::from(DDSKETCH_CONF_NORM_BIAS) + 0.5)
            .expect("agent sketch gamma and offset are always valid")
    }

    /// Returns the representative value for the given agent sketch key.
    pub fn value_for_key(key: i16) -> f64 {
        SKETCH_CONFIG.bin_lower_bound(key)
    }

    /// Returns the number of bins in the sketch.
    pub fn bin_count(&self) -> usize {
        self.bins.len()
    }

    /// Whether or not this sketch is empty.
    pub fn is_empty(&self) -> bool {
        self.count == 0
    }

    /// Number of samples currently represented by this sketch.
    pub fn count(&self) -> u64 {
        self.count
    }

    /// Overrides the sample count tracked by this sketch.
    pub fn set_count(&mut self, count: u64) {
        self.count = count;
    }

    /// Overrides the sum of all values tracked by this sketch.
    pub fn set_sum(&mut self, sum: f64) {
        self.sum = sum;
    }

    /// Overrides the average value tracked by this sketch.
    pub fn set_avg(&mut self, avg: f64) {
        self.avg = avg;
    }

    /// Overrides the minimum value tracked by this sketch.
    pub fn set_min(&mut self, min: f64) {
        self.min = min;
    }

    /// Overrides the maximum value tracked by this sketch.
    pub fn set_max(&mut self, max: f64) {
        self.max = max;
    }

    /// Minimum value seen by this sketch.
    ///
    /// Returns `None` if the sketch is empty.
    pub fn min(&self) -> Option<f64> {
        if self.is_empty() {
            None
        } else {
            Some(self.min)
        }
    }

    /// Maximum value seen by this sketch.
    ///
    /// Returns `None` if the sketch is empty.
    pub fn max(&self) -> Option<f64> {
        if self.is_empty() {
            None
        } else {
            Some(self.max)
        }
    }

    /// Sum of all values seen by this sketch.
    ///
    /// Returns `None` if the sketch is empty.
    pub fn sum(&self) -> Option<f64> {
        if self.is_empty() {
            None
        } else {
            Some(self.sum)
        }
    }

    /// Returns the stored sum summary field without considering the sample count.
    pub fn stored_sum(&self) -> f64 {
        self.sum
    }

    /// Average value seen by this sketch.
    ///
    /// Returns `None` if the sketch is empty.
    pub fn avg(&self) -> Option<f64> {
        if self.is_empty() {
            None
        } else {
            Some(self.avg)
        }
    }

    /// Returns the stored average summary field without considering the sample count.
    pub fn stored_avg(&self) -> f64 {
        self.avg
    }

    /// Returns the stored minimum summary field without considering the sample count.
    pub fn stored_min(&self) -> f64 {
        self.min
    }

    /// Returns the stored maximum summary field without considering the sample count.
    pub fn stored_max(&self) -> f64 {
        self.max
    }

    /// Returns the current bins of this sketch.
    pub fn bins(&self) -> &[Bin] {
        &self.bins
    }

    /// Clears the sketch, removing all bins and resetting all statistics.
    pub fn clear(&mut self) {
        self.count = 0;
        self.min = f64::MAX;
        self.max = f64::MIN;
        self.avg = 0.0;
        self.sum = 0.0;
        self.bins.clear();
    }

    fn adjust_basic_stats(&mut self, v: f64, n: u64) {
        // Every insert path funnels through here, so this is where we guard that the incoming sample is finite. If
        // this is not true something has gone wrong with the DogStatsD codec.
        saluki_antithesis::always!(v.is_finite(), "DDSketch sample is finite at insert");

        if v < self.min {
            self.min = v;
        }

        if v > self.max {
            self.max = v;
        }

        saluki_antithesis::always_le!(self.min, self.max, "DDSketch min does not exceed max after insert");

        self.count += n;
        // It's possible that self.sum will be INF after this multiplication, even though we've demonstrated that `v`
        // is finite. The Datadog Agent sketch sum behaves the same way, so we do not assert that self.sum is itself
        // finite.
        self.sum += v * n as f64;

        if n == 1 {
            self.avg += (v - self.avg) / self.count as f64;
        } else {
            // TODO: From the Agent source code, this method apparently loses precision when the
            // two averages -- v and self.avg -- are close.  Is there a better approach?
            self.avg = self.avg + (v - self.avg) * n as f64 / self.count as f64;
        }
    }

    fn insert_key_counts(&mut self, counts: &[(i16, u64)]) {
        let mut temp = SmallVec::<[Bin; 4]>::new();

        let mut bins_idx = 0;
        let mut key_idx = 0;
        let bins_len = self.bins.len();
        let counts_len = counts.len();

        // PERF TODO: there's probably a fast path to be had where could check if all if the counts have existing bins
        // that aren't yet full, and we just update them directly, although we'd still be doing a linear scan to find
        // them since keys aren't 1:1 with their position in `self.bins` but using this method just to update one or two
        // bins is clearly suboptimal and we wouldn't really want to scan them all just to have to back out and actually
        // do the non-fast path.. maybe a first pass could be checking if the first/last key falls within our known
        // min/max key, and if it doesn't, then we know we have to go through the non-fast path, and if it passes, we do
        // the scan to see if we can just update bins directly?
        while bins_idx < bins_len && key_idx < counts_len {
            let bin = self.bins[bins_idx];
            let vk = counts[key_idx].0;
            let kn = counts[key_idx].1;

            match bin.k.cmp(&vk) {
                Ordering::Greater => {
                    generate_bins(&mut temp, vk, kn);
                    key_idx += 1;
                }
                Ordering::Less => {
                    temp.push(bin);
                    bins_idx += 1;
                }
                Ordering::Equal => {
                    generate_bins(&mut temp, bin.k, u64::from(bin.n) + kn);
                    bins_idx += 1;
                    key_idx += 1;
                }
            }
        }

        temp.extend_from_slice(&self.bins[bins_idx..]);

        while key_idx < counts_len {
            let vk = counts[key_idx].0;
            let kn = counts[key_idx].1;
            generate_bins(&mut temp, vk, kn);
            key_idx += 1;
        }

        trim_left(&mut temp, SKETCH_CONFIG.bin_limit);

        // PERF TODO: This is where we might do a mem::swap instead so that we could shove the bin vector into an object
        // pool but I'm not sure this actually matters at the moment.
        self.bins = temp;
    }

    fn insert_keys(&mut self, mut keys: Vec<i16>) {
        // Updating more than 4 billion keys would be very very weird and likely indicative of something horribly
        // broken.
        //
        // TODO: I don't actually understand why I wrote this assertion in this way. Either the code can handle
        // collapsing values in order to maintain the relative error bounds, or we have to cap it to the maximum allowed
        // number of bins. Gotta think about this some more.
        assert!(keys.len() <= u32::MAX.try_into().expect("we don't support 16-bit systems"));

        keys.sort_unstable();

        let mut temp = SmallVec::<[Bin; 4]>::new();

        let mut bins_idx = 0;
        let mut key_idx = 0;
        let bins_len = self.bins.len();
        let keys_len = keys.len();

        // PERF TODO: there's probably a fast path to be had where could check if all if the counts have existing bins
        // that aren't yet full, and we just update them directly, although we'd still be doing a linear scan to find
        // them since keys aren't 1:1 with their position in `self.bins` but using this method just to update one or two
        // bins is clearly suboptimal and we wouldn't really want to scan them all just to have to back out and actually
        // do the non-fast path.. maybe a first pass could be checking if the first/last key falls within our known
        // min/max key, and if it doesn't, then we know we have to go through the non-fast path, and if it passes, we do
        // the scan to see if we can just update bins directly?
        while bins_idx < bins_len && key_idx < keys_len {
            let bin = self.bins[bins_idx];
            let vk = keys[key_idx];

            match bin.k.cmp(&vk) {
                Ordering::Greater => {
                    let kn = buf_count_leading_equal(&keys, key_idx);
                    generate_bins(&mut temp, vk, kn);
                    key_idx += kn as usize;
                }
                Ordering::Less => {
                    temp.push(bin);
                    bins_idx += 1;
                }
                Ordering::Equal => {
                    let kn = buf_count_leading_equal(&keys, key_idx);
                    generate_bins(&mut temp, bin.k, u64::from(bin.n) + kn);
                    bins_idx += 1;
                    key_idx += kn as usize;
                }
            }
        }

        temp.extend_from_slice(&self.bins[bins_idx..]);

        while key_idx < keys_len {
            let vk = keys[key_idx];
            let kn = buf_count_leading_equal(&keys, key_idx);
            generate_bins(&mut temp, vk, kn);
            key_idx += kn as usize;
        }

        trim_left(&mut temp, SKETCH_CONFIG.bin_limit);

        // PERF TODO: This is where we might do a mem::swap instead so that we could shove the bin vector into an object
        // pool but I'm not sure this actually matters at the moment.
        self.bins = temp;
    }

    /// Inserts a single value into the sketch.
    pub fn insert(&mut self, v: f64) {
        // TODO: This should return a result that makes sure we have enough room to actually add 1 more sample without
        // hitting `self.config.max_count()`
        self.adjust_basic_stats(v, 1);

        let key = SKETCH_CONFIG.key(v);

        let mut insert_at = None;

        for (bin_idx, b) in self.bins.iter_mut().enumerate() {
            if b.k == key {
                if b.n < MAX_BIN_WIDTH {
                    // Fast path for adding to an existing bin without overflow.
                    b.n += 1;
                    return;
                } else {
                    insert_at = Some(bin_idx);
                    break;
                }
            }
            if b.k > key {
                insert_at = Some(bin_idx);
                break;
            }
        }

        if let Some(bin_idx) = insert_at {
            self.bins.insert(bin_idx, Bin { k: key, n: 1 });
        } else {
            self.bins.push(Bin { k: key, n: 1 });
        }
        trim_left(&mut self.bins, SKETCH_CONFIG.bin_limit);
    }

    /// Inserts many values into the sketch.
    pub fn insert_many(&mut self, vs: &[f64]) {
        // TODO: This should return a result that makes sure we have enough room to actually add N more samples without
        // hitting `self.config.bin_limit`.
        let mut keys = Vec::with_capacity(vs.len());
        for v in vs {
            self.adjust_basic_stats(*v, 1);
            keys.push(SKETCH_CONFIG.key(*v));
        }
        self.insert_keys(keys);
    }

    /// Inserts a single value into the sketch `n` times.
    pub fn insert_n(&mut self, v: f64, n: u64) {
        // TODO: This should return a result that makes sure we have enough room to actually add N more samples without
        // hitting `self.config.max_count()`.
        if n == 1 {
            self.insert(v);
        } else {
            self.adjust_basic_stats(v, n);

            let key = SKETCH_CONFIG.key(v);
            self.insert_key_counts(&[(key, n)]);
        }
    }

    fn insert_interpolate_bucket(&mut self, lower: f64, upper: f64, count: u64) {
        // Find the keys for the bins where the lower bound and upper bound would end up, and collect all of the keys in
        // between, inclusive.
        let lower_key = SKETCH_CONFIG.key(lower);
        let upper_key = SKETCH_CONFIG.key(upper);
        let keys = (lower_key..=upper_key).collect::<Vec<_>>();

        let mut key_counts = Vec::new();
        let mut remaining_count = count;
        let distance = upper - lower;
        let mut start_idx = 0;
        let mut end_idx = 1;
        let mut lower_bound = SKETCH_CONFIG.bin_lower_bound(keys[start_idx]);
        let mut remainder = 0.0;

        while end_idx < keys.len() && remaining_count > 0 {
            // For each key, map the total distance between the input lower/upper bound against the sketch lower/upper
            // bound for the current sketch bin, which tells us how much of the input count to apply to the current
            // sketch bin.
            let upper_bound = SKETCH_CONFIG.bin_lower_bound(keys[end_idx]);
            let fkn = ((upper_bound - lower_bound) / distance) * count as f64;
            if fkn > 1.0 {
                remainder += fkn - fkn.trunc();
            }

            // SAFETY: This integer cast is intentional: we want to get the non-fractional part, as we've captured the
            // fractional part in the above conditional.
            #[allow(clippy::cast_possible_truncation)]
            let mut kn = fkn as u64;
            if remainder > 1.0 {
                kn += 1;
                remainder -= 1.0;
            }

            if kn > 0 {
                if kn > remaining_count {
                    kn = remaining_count;
                }

                self.adjust_basic_stats(lower_bound, kn);
                key_counts.push((keys[start_idx], kn));

                remaining_count -= kn;
                start_idx = end_idx;
                lower_bound = upper_bound;
            }

            end_idx += 1;
        }

        if remaining_count > 0 {
            let last_key = keys[start_idx];
            lower_bound = SKETCH_CONFIG.bin_lower_bound(last_key);
            self.adjust_basic_stats(lower_bound, remaining_count);
            key_counts.push((last_key, remaining_count));
        }

        // Sort the key counts first, as that's required by `insert_key_counts`.
        key_counts.sort_unstable_by_key(|(k, _)| *k);

        self.insert_key_counts(&key_counts);
    }

    /// Inserts histogram buckets into the sketch via linear interpolation.
    ///
    /// ## Errors
    ///
    /// Returns an error if a bucket size is greater that `u32::MAX`.
    pub fn insert_interpolate_buckets(&mut self, mut buckets: Vec<Bucket>) -> Result<(), &'static str> {
        // Buckets need to be sorted from lowest to highest so that we can properly calculate the rolling lower/upper
        // bounds.
        buckets.sort_by(|a, b| {
            let oa = OrderedFloat(a.upper_limit);
            let ob = OrderedFloat(b.upper_limit);

            oa.cmp(&ob)
        });

        let mut lower = f64::NEG_INFINITY;

        if buckets.iter().any(|bucket| bucket.count > u64::from(u32::MAX)) {
            return Err("bucket size greater than u32::MAX");
        }

        for bucket in buckets {
            let original_upper = bucket.upper_limit;
            let mut upper = bucket.upper_limit;
            if upper.is_sign_positive() && upper.is_infinite() {
                upper = lower;
            } else if lower.is_sign_negative() && lower.is_infinite() {
                lower = upper;
            } else if lower == 0.0 && upper > 0.0 {
                // OpenTelemetry explicit buckets use (lower, upper]. After a zero boundary, the
                // next bucket is (0, upper] and excludes zero. Collapse interpolation to `upper`
                // so it does not add samples at zero.
                lower = upper;
            }

            self.insert_interpolate_bucket(lower, upper, bucket.count);
            lower = original_upper;
        }

        Ok(())
    }

    /// Adds a bin directly into the sketch.
    ///
    /// Used only for unit testing so that we can create a sketch with an exact layout, which allows testing around the
    /// resulting bins when feeding in specific values, as well as generating explicitly bad layouts for testing.
    #[allow(dead_code)]
    pub(crate) fn insert_raw_bin(&mut self, k: i16, n: u32) {
        let v = SKETCH_CONFIG.bin_lower_bound(k);
        self.adjust_basic_stats(v, u64::from(n));
        self.bins.push(Bin { k, n });
    }

    /// Gets the value at a given quantile.
    pub fn quantile(&self, q: f64) -> Option<f64> {
        if self.count == 0 {
            return None;
        }

        if q <= 0.0 {
            return Some(self.min);
        }

        if q >= 1.0 {
            return Some(self.max);
        }

        let mut n = 0.0;
        let mut estimated = None;
        let wanted_rank = rank(self.count, q);

        for (i, bin) in self.bins.iter().enumerate() {
            n += f64::from(bin.n);
            if n <= wanted_rank {
                continue;
            }

            let weight = (n - wanted_rank) / f64::from(bin.n);
            let mut v_low = SKETCH_CONFIG.bin_lower_bound(bin.k);
            let mut v_high = v_low * SKETCH_CONFIG.gamma_v;

            if i == self.bins.len() {
                v_high = self.max;
            } else if i == 0 {
                v_low = self.min;
            }

            estimated = Some(v_low * weight + v_high * (1.0 - weight));
            break;
        }

        estimated.map(|v| v.clamp(self.min, self.max)).or(Some(f64::NAN))
    }

    /// Merges another sketch into this sketch, without a loss of accuracy.
    ///
    /// All samples present in the other sketch will be correctly represented in this sketch, and summary statistics
    /// such as the sum, average, count, min, and max, will represent the sum of samples from both sketches.
    ///
    /// A zero-count sketch can retain bins after an upstream cumulative-to-delta conversion. Its summary fields do not
    /// describe samples in the current interval, so they do not contribute to the merged summary. Its bins are still
    /// merged.
    pub fn merge(&mut self, other: &DDSketch) {
        // Zero-count destination matches other summary but does not affect an existing summary
        if self.count == 0 {
            self.count = other.count;
            self.min = other.min;
            self.max = other.max;
            self.sum = other.sum;
            self.avg = other.avg;
        } else if other.count > 0 {
            self.count += other.count;
            if other.max > self.max {
                self.max = other.max;
            }
            if other.min < self.min {
                self.min = other.min;
            }
            self.sum += other.sum;
            self.avg = self.avg + (other.avg - self.avg) * other.count as f64 / self.count as f64;
        }

        // Merge the bins regardless of the summary count.
        let mut temp = SmallVec::<[Bin; 4]>::new();

        let mut bins_idx = 0;
        for other_bin in &other.bins {
            let start = bins_idx;
            while bins_idx < self.bins.len() && self.bins[bins_idx].k < other_bin.k {
                bins_idx += 1;
            }

            temp.extend_from_slice(&self.bins[start..bins_idx]);

            if bins_idx >= self.bins.len() || self.bins[bins_idx].k > other_bin.k {
                temp.push(*other_bin);
            } else if self.bins[bins_idx].k == other_bin.k {
                generate_bins(
                    &mut temp,
                    other_bin.k,
                    u64::from(other_bin.n) + u64::from(self.bins[bins_idx].n),
                );
                bins_idx += 1;
            }
        }

        temp.extend_from_slice(&self.bins[bins_idx..]);
        trim_left(&mut temp, SKETCH_CONFIG.bin_limit);

        self.bins = temp;
    }

    /// Merges this sketch into the `Dogsketch` Protocol Buffers representation.
    pub fn merge_to_dogsketch(&self, dogsketch: &mut Dogsketch) {
        dogsketch.set_cnt(i64::try_from(self.count).unwrap_or(i64::MAX));
        dogsketch.set_min(self.min);
        dogsketch.set_max(self.max);
        dogsketch.set_avg(self.avg);
        dogsketch.set_sum(self.sum);

        let mut k = Vec::new();
        let mut n = Vec::new();

        for bin in &self.bins {
            k.push(i32::from(bin.k));
            n.push(bin.n);
        }

        dogsketch.set_k(k);
        dogsketch.set_n(n);
    }
}

impl PartialEq for DDSketch {
    fn eq(&self, other: &Self) -> bool {
        // We skip checking the configuration because we don't allow creating configurations by hand, and it's always
        // locked to the constants used by the Datadog Agent.  We only check the configuration equality manually in
        // `DDSketch::merge`, to protect ourselves in the future if different configurations become allowed.
        //
        // Additionally, we also use floating-point-specific relative comparisons for sum/avg because they can be
        // minimally different between sketches purely due to floating-point behavior, despite being fed the same exact
        // data in terms of recorded samples.
        self.count == other.count
            && float_eq(self.min, other.min)
            && float_eq(self.max, other.max)
            && float_eq(self.sum, other.sum)
            && float_eq(self.avg, other.avg)
            && self.bins == other.bins
    }
}

impl Default for DDSketch {
    fn default() -> Self {
        Self {
            bins: SmallVec::new(),
            count: 0,
            min: f64::MAX,
            max: f64::MIN,
            sum: 0.0,
            avg: 0.0,
        }
    }
}

impl Eq for DDSketch {}

impl TryFrom<Dogsketch> for DDSketch {
    type Error = &'static str;

    fn try_from(value: Dogsketch) -> Result<Self, Self::Error> {
        let mut sketch = DDSketch {
            count: u64::try_from(value.cnt).map_err(|_| "sketch count overflows u64 or is negative")?,
            min: value.min,
            max: value.max,
            avg: value.avg,
            sum: value.sum,
            ..Default::default()
        };

        let k = value.k;
        let n = value.n;

        if k.len() != n.len() {
            return Err("k and n bin vectors have differing lengths");
        }

        for (k, n) in k.into_iter().zip(n) {
            let k = i16::try_from(k).map_err(|_| "bin key overflows i16")?;

            sketch.bins.push(Bin { k, n });
        }

        Ok(sketch)
    }
}

fn rank(count: u64, q: f64) -> f64 {
    let rank = q * (count - 1) as f64;
    rank.round_ties_even()
}

#[allow(clippy::cast_possible_truncation)]
fn buf_count_leading_equal(keys: &[i16], start_idx: usize) -> u64 {
    if start_idx == keys.len() - 1 {
        return 1;
    }

    let mut idx = start_idx;
    while idx < keys.len() && keys[idx] == keys[start_idx] {
        idx += 1;
    }

    // SAFETY: We limit the size of the vector (used to provide the slice given to us here) to be no larger than 2^32,
    // so we can't exceed u64 here.
    (idx - start_idx) as u64
}

fn trim_left(bins: &mut SmallVec<[Bin; 4]>, bin_limit: u16) {
    // We won't ever support Vector running on anything other than a 32-bit platform and above, I imagine, so this
    // should always be safe.
    let bin_limit = bin_limit as usize;
    if bin_limit == 0 || bins.len() <= bin_limit {
        return;
    }

    let num_to_remove = bins.len() - bin_limit;
    let mut missing: u64 = 0;

    // Sum all mass from the bins being removed. Per CollapsingLowestDenseStore in sketches-go,
    // all removed mass collapses into the first kept bin (the new minimum index).
    for bin in bins.iter().take(num_to_remove) {
        missing += u64::from(bin.n);
    }

    // Fold the accumulated mass into the first kept bin, matching Go's `bins[newMinIndex] += n`.
    // Any remainder that overflows u32::MAX is discarded—this requires >4B observations in a
    // single collapsed bin and is an intentional divergence from the Datadog Agent (which uses
    // float64 counts and never loses mass).
    bins[num_to_remove].increment(missing);

    // Drop the removed prefix, leaving exactly bin_limit bins.
    bins.drain(0..num_to_remove);

    // This is the one place every mutating method routes through, so asserting here guards the bin-count bound for
    // all of them.
    saluki_antithesis::reachable!("DDSketch bin collapse reached");
    saluki_antithesis::always_le!(bins.len(), bin_limit, "DDSketch bin count within bin_limit");
}

#[allow(clippy::cast_possible_truncation)]
fn generate_bins(bins: &mut SmallVec<[Bin; 4]>, k: i16, n: u64) {
    if n < u64::from(MAX_BIN_WIDTH) {
        // SAFETY: `n < MAX_BIN_WIDTH = u32::MAX`, so it fits in u32.
        bins.push(Bin { k, n: n as u32 });
    } else {
        let overflow = n % u64::from(MAX_BIN_WIDTH);
        if overflow != 0 {
            bins.push(Bin {
                k,
                // SAFETY: `overflow = n % u32::MAX`, so overflow <= u32::MAX - 1, which fits in u32.
                n: overflow as u32,
            });
        }

        for _ in 0..(n / u64::from(MAX_BIN_WIDTH)) {
            bins.push(Bin { k, n: MAX_BIN_WIDTH });
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    // Helper: build a SmallVec<[Bin; 4]> from (k, n) pairs. Assumes input is already sorted by k.
    fn make_bins(pairs: &[(i16, u32)]) -> SmallVec<[Bin; 4]> {
        pairs.iter().map(|&(k, n)| Bin { k, n }).collect()
    }

    // Helper: extract (k, n) pairs from a SmallVec for easy assertion.
    fn to_pairs(bins: &SmallVec<[Bin; 4]>) -> Vec<(i16, u32)> {
        bins.iter().map(|b| (b.k, b.n)).collect()
    }

    /// Basic collapse: when bins exceed the limit, the mass from removed bins is merged
    /// into the first kept bin (lowest surviving key). This mirrors the CollapsingLowestDenseStore
    /// semantics from sketches-go: all bins with index < (`maxIndex` - limit + 1) collapse into
    /// the bin at (`maxIndex` - limit + 1).
    ///
    /// Input:  [(0,2), (1,3), (2,4), (3,5)]  limit=2  →  remove 2 bins
    /// missing = n[0] + n[1] = 2 + 3 = 5
    /// first kept bin: k=2, n=4 → increment by 5 → n=9
    /// Result: [(2,9), (3,5)]
    #[test]
    fn trim_left_collapses_removed_mass_into_first_kept_bin() {
        let mut bins = make_bins(&[(0, 2), (1, 3), (2, 4), (3, 5)]);
        trim_left(&mut bins, 2);
        assert_eq!(to_pairs(&bins), vec![(2, 9), (3, 5)]);
    }

    /// Total count is preserved exactly when the collapse fits within a single u32 bin.
    ///
    /// Input:  [(10,5), (20,3), (30,7)]  limit=2  →  remove 1 bin
    /// missing = 5; first kept bin k=20, n=3 → n=8
    /// Result: [(20,8), (30,7)], total=15 == 5+3+7
    #[test]
    fn trim_left_preserves_total_count_when_no_overflow() {
        let mut bins = make_bins(&[(10, 5), (20, 3), (30, 7)]);
        let total_before: u64 = bins.iter().map(|b| u64::from(b.n)).sum();
        trim_left(&mut bins, 2);
        let total_after: u64 = bins.iter().map(|b| u64::from(b.n)).sum();
        assert_eq!(to_pairs(&bins), vec![(20, 8), (30, 7)]);
        assert_eq!(total_before, total_after);
    }

    /// With u32 bin counts, collapsed mass from multiple removed bins fits in a single bin
    /// without saturation for typical weights, fully preserving the total count.
    ///
    /// Input:  [(0,50000), (1,50000), (2,1)]  limit=1  →  remove 2 bins
    /// missing = 100000; bins[2].increment(100000): 100001 < u32::MAX → n=100001
    /// bins.drain(0..2). Final: [(2,100001)], all mass preserved.
    ///
    /// With the old u16 layout, this same input would have saturated at 65535 and discarded
    /// 34466 observations. u32 eliminates that loss for any per-bin count below ~4.3 billion.
    #[test]
    fn trim_left_preserves_exact_count_with_u32_bins() {
        let mut bins = make_bins(&[(0, 50000), (1, 50000), (2, 1)]);
        let total_before: u64 = bins.iter().map(|b| u64::from(b.n)).sum();
        trim_left(&mut bins, 1);
        let total_after: u64 = bins.iter().map(|b| u64::from(b.n)).sum();
        assert_eq!(bins.len(), 1);
        assert_eq!(bins[0].k, 2);
        assert_eq!(bins[0].n, 100001);
        assert_eq!(total_before, total_after, "all mass must be preserved with u32 bins");
    }

    /// When already at or under the limit, `trim_left` is a no-op.
    #[test]
    fn trim_left_no_op_when_within_limit() {
        let original = make_bins(&[(5, 10), (6, 20)]);
        let mut bins = original.clone();
        trim_left(&mut bins, 2);
        assert_eq!(to_pairs(&bins), to_pairs(&original));
        trim_left(&mut bins, 3);
        assert_eq!(to_pairs(&bins), to_pairs(&original));
    }

    /// Regression test for `trim_left` bin count with large per-sample weights.
    ///
    /// With the old `u16` layout, a sample weight of ~260M would generate `ceil(260M / 65535)`, or 3969, bins per key,
    /// causing bin count explosion and an encoder panic. With `u32`, the same weight fits in a single bin (up to
    /// `u32::MAX`, ~4.3 billion, samples per bin), so one `insert_n` call produces exactly one bin per key and the bin
    /// limit is trivially respected.
    ///
    /// This test inserts several values with a weight representative of what ADP receives when clamping an incoming
    /// sample rate of `3e-9` to its minimum of `3.845e-9` (~260M per sample), then asserts the bin count never exceeds
    /// `DDSKETCH_CONF_BIN_LIMIT`.
    #[test]
    fn trim_left_respects_bin_limit_with_large_weights() {
        // Weight corresponding to ADP's minimum safe sample rate (1 / 3.845e-9 ≈ 260_078_024).
        // With u32 bins, this fits in a single bin per key (260_078_024 < u32::MAX).
        let weight: u64 = 260_078_024;
        let bin_limit = usize::from(DDSKETCH_CONF_BIN_LIMIT);

        let mut sketch = DDSketch::default();

        // Insert enough distinct values to repeatedly trigger trim_left.  Ten values is more
        // than sufficient; two already exceed the bin limit without the fix.
        for i in 1..=10_i32 {
            sketch.insert_n(f64::from(i), weight);
            assert!(
                sketch.bins().len() <= bin_limit,
                "bin count {} exceeded limit {} after inserting {} value(s) at weight {}",
                sketch.bins().len(),
                bin_limit,
                i,
                weight,
            );
        }
    }

    #[test]
    fn interpolate_buckets_does_not_seed_zero_bin_for_positive_bucket_after_zero_bound() {
        let mut sketch = DDSketch::default();

        sketch
            .insert_interpolate_buckets(vec![
                Bucket {
                    upper_limit: 0.0,
                    count: 0,
                },
                Bucket {
                    upper_limit: 10.0,
                    count: 10,
                },
            ])
            .expect("buckets should interpolate");

        assert_eq!(sketch.count(), 10);
        assert!(
            sketch.min().expect("sketch should not be empty") > 0.0,
            "positive bucket after a zero lower bound must not seed a zero minimum"
        );
    }

    /// When the accumulated missing mass plus the first kept bin's existing count exceeds
    /// u32::MAX, increment saturates at u32::MAX and the remainder is discarded.
    ///
    /// Input:  [(0, u32::MAX), (1, 1)]  limit=1  →  remove 1 bin
    /// missing = u32::MAX; bins[1].increment(u32::MAX): next = u32::MAX+1 > u32::MAX
    /// → n = u32::MAX, remainder = 1 (discarded)
    /// Final: [(1, u32::MAX)]—1 observation lost
    #[test]
    fn trim_left_saturates_first_kept_bin_and_discards_remainder() {
        let mut bins = make_bins(&[(0, u32::MAX), (1, 1)]);
        trim_left(&mut bins, 1);
        assert_eq!(bins.len(), 1);
        assert_eq!(bins[0].k, 1);
        assert_eq!(bins[0].n, u32::MAX);
    }

    /// Collapsing works correctly with negative bin keys. The highest key always survives;
    /// lower (more negative) keys collapse into the first kept key.
    /// Adapted from TestAddIntDatasets in sketches-go which covers negative indices.
    ///
    /// Input:  [(-3,5), (-2,3), (-1,2), (0,1)]  limit=2  →  remove 2 bins
    /// missing = 5+3=8; bins[-1].increment(8): 2+8=10 < u32::MAX → n=10
    /// Final: [(-1,10), (0,1)]
    #[test]
    fn trim_left_collapses_negative_keys_correctly() {
        let mut bins = make_bins(&[(-3, 5), (-2, 3), (-1, 2), (0, 1)]);
        trim_left(&mut bins, 2);
        assert_eq!(to_pairs(&bins), vec![(-1, 10), (0, 1)]);
    }

    /// Monotonic ascending sequence: inserting keys 0..N where N > limit collapses the
    /// lowest keys into the first kept key, with their total mass summed there.
    /// Adapted from TestAddMonotonous in sketches-go.
    ///
    /// Input:  [(0,1),(1,1),(2,1),(3,1),(4,1),(5,1),(6,1),(7,1),(8,1),(9,1)]  limit=4
    /// num_to_remove=6; missing=6; bins[6].increment(6): 1+6=7 → n=7
    /// Final: [(6,7),(7,1),(8,1),(9,1)] —only top 4 keys kept, collapsed mass in first
    #[test]
    fn trim_left_monotonic_ascending_keeps_top_keys() {
        let pairs: Vec<(i16, u32)> = (0..10).map(|i| (i, 1)).collect();
        let mut bins = make_bins(&pairs);
        trim_left(&mut bins, 4);
        assert_eq!(to_pairs(&bins), vec![(6, 7), (7, 1), (8, 1), (9, 1)]);
    }

    #[test]
    fn merge_handles_zero_count_summaries_without_dropping_bins() {
        for (name, destination_is_zero_count, source_is_zero_count, expected_summary) in [
            ("zero-count source", false, true, 10.0),
            ("zero-count destination", true, false, 20.0),
        ] {
            let mut destination = DDSketch::default();
            destination.insert(10.0);

            let mut source = DDSketch::default();
            source.insert(20.0);

            if destination_is_zero_count {
                destination.set_count(0);
                destination.set_min(-1.0);
                destination.set_max(100.0);
                destination.set_sum(999.0);
                destination.set_avg(f64::NAN);
            }

            if source_is_zero_count {
                source.set_count(0);
                source.set_min(-1.0);
                source.set_max(100.0);
                source.set_sum(999.0);
                source.set_avg(f64::NAN);
            }

            destination.merge(&source);

            assert_eq!(destination.count(), 1, "{name}");
            assert_eq!(destination.stored_min(), expected_summary, "{name}");
            assert_eq!(destination.stored_max(), expected_summary, "{name}");
            assert_eq!(destination.stored_sum(), expected_summary, "{name}");
            assert_eq!(destination.stored_avg(), expected_summary, "{name}");
            assert_eq!(destination.bin_count(), 2, "{name}");
        }
    }

    #[cfg(feature = "serde")]
    #[test]
    fn json_null_summary_round_trips_as_nan() {
        let mut sketch = DDSketch::default();
        sketch.insert(42.0);
        sketch.set_avg(f64::NAN);

        let encoded = serde_json::to_string(&sketch).expect("sketch should serialize");
        assert!(encoded.contains("\"avg\":null"));

        let decoded: DDSketch = serde_json::from_str(&encoded).expect("sketch should deserialize");
        assert!(decoded.avg().expect("sketch should not be empty").is_nan());
    }
}

/// Property-based tests for `trim_left`, adapted from TestAddFuzzy / TestAddIntFuzzy /
/// TestMergeFuzzy in sketches-go/ddsketch/store/store_test.go.
///
/// The sketches-go suite generates random (index, count) inputs, passes them through
/// CollapsingLowestDenseStore, and asserts structural invariants using the `collapsingLowest`
/// oracle transform. We do the same here: generate random sorted distinct-key bins, run
/// `trim_left`, and check the same invariants.
#[cfg(test)]
mod property_tests {
    use proptest::prelude::*;

    use super::*;

    /// Strategy: a non-empty sorted vec of distinct (k: i16, n: u32) pairs.
    /// Keys are drawn from i16 without repetition; counts are 1..=u32::MAX.
    fn arb_bins(max_len: usize) -> impl Strategy<Value = SmallVec<[Bin; 4]>> {
        proptest::collection::btree_map(any::<i16>(), 1u32..=u32::MAX, 1..=max_len)
            .prop_map(|map| map.into_iter().map(|(k, n)| Bin { k, n }).collect())
    }

    /// Strategy: a `bin_limit` in 1..=32 (small enough to exercise collapsing frequently).
    fn arb_limit() -> impl Strategy<Value = u16> {
        1u16..=32
    }

    proptest! {
        /// After `trim_left`, the bin count must never exceed `bin_limit`.
        ///
        /// Mirrors the core bin-count invariant checked throughout the sketches-go suite:
        /// every store operation must leave the store within its configured capacity.
        #[test]
        fn prop_bin_count_never_exceeds_limit(
            mut bins in arb_bins(64),
            limit in arb_limit(),
        ) {
            trim_left(&mut bins, limit);
            prop_assert!(
                bins.len() <= limit as usize,
                "bin count {} exceeded limit {}",
                bins.len(),
                limit,
            );
        }

        /// After `trim_left`, bins must remain sorted by key with no duplicate keys.
        ///
        /// `trim_left` only drains a prefix and modifies bins[num_to_remove].n in place;
        /// it must not disturb the ordering of the surviving bins.
        #[test]
        fn prop_output_bins_are_sorted_and_distinct(
            mut bins in arb_bins(64),
            limit in arb_limit(),
        ) {
            trim_left(&mut bins, limit);
            for window in bins.windows(2) {
                prop_assert!(
                    window[0].k < window[1].k,
                    "bins not strictly sorted after trim_left: {:?} >= {:?}",
                    window[0].k,
                    window[1].k,
                );
            }
        }

        /// When bins.len() <= limit, `trim_left` is a no-op: bins is unchanged.
        #[test]
        fn prop_no_op_when_within_limit(
            bins in arb_bins(16),
            extra in 0u16..=16,
        ) {
            let limit = bins.len() as u16 + extra;
            let original = bins.clone();
            let mut bins = bins;
            trim_left(&mut bins, limit);
            prop_assert_eq!(bins, original);
        }

        /// Total count is preserved exactly when no u32 overflow occurs.
        ///
        /// This is the key invariant from sketches-go's assertEncodeBins:
        /// `store.TotalCount()` must equal the sum of all inserted counts.
        /// We restrict counts to keep the collapsed sum safely below u32::MAX
        /// (sketches-go never loses mass because it uses float64; we match that
        /// guarantee for all inputs where the collapsed bin doesn't overflow u32).
        ///
        /// Strategy: counts capped at 1000 so that even if all 64 bins collapse
        /// into one the sum (≤ 64_000) is far below u32::MAX.
        #[test]
        fn prop_total_count_preserved_when_no_overflow(
            mut bins in proptest::collection::btree_map(any::<i16>(), 1u32..=1000, 1..=64usize)
                .prop_map(|map| -> SmallVec<[Bin; 4]> {
                    map.into_iter().map(|(k, n)| Bin { k, n }).collect()
                }),
            limit in arb_limit(),
        ) {
            let total_before: u64 = bins.iter().map(|b| u64::from(b.n)).sum();
            trim_left(&mut bins, limit);
            let total_after: u64 = bins.iter().map(|b| u64::from(b.n)).sum();
            prop_assert_eq!(
                total_before, total_after,
                "total count changed: before={}, after={}",
                total_before, total_after,
            );
        }

        /// The surviving bins are always the `bin_limit` highest-key bins from the input.
        ///
        /// This is the direct encoding of the `collapsingLowest` oracle from sketches-go:
        /// minCollapsedIndex = `maxIndex` - limit + 1; all bins with key < minCollapsedIndex
        /// are removed (their mass folds into minCollapsedIndex). The output keys must
        /// exactly match the top min(len, limit) keys of the input.
        #[test]
        fn prop_output_keys_are_highest_from_input(
            mut bins in arb_bins(64),
            limit in arb_limit(),
        ) {
            let all_keys: Vec<i16> = bins.iter().map(|b| b.k).collect();
            let expected_len = all_keys.len().min(limit as usize);
            let expected_keys: Vec<i16> = all_keys[all_keys.len() - expected_len..].to_vec();

            trim_left(&mut bins, limit);

            let actual_keys: Vec<i16> = bins.iter().map(|b| b.k).collect();
            prop_assert_eq!(
                actual_keys, expected_keys,
                "output keys don't match top {} keys of input",
                limit,
            );
        }
    }
}

/// Direct unit tests for the agent `DDSketch`'s public API.
///
/// The `tests`/`property_tests` modules above only cover the internal `trim_left` collapse helper; this module
/// exercises the public surface (insert/insert_n/insert_many, quantile, merge, the basic-statistics accessors,
/// clear, histogram-bucket interpolation, and `Dogsketch` conversion) that all production callers actually use.
#[cfg(test)]
mod public_api_tests {
    use datadog_protos::metrics::Dogsketch;

    use super::*;
    // `remap_mapping` returns a `LogarithmicMapping`; bring the trait into scope so its `gamma()` accessor is callable.
    use crate::canonical::mapping::IndexMapping as _;

    // The agent sketch is configured with eps = 1/128, giving gamma_v = 1.015625 and a per-bin relative accuracy of
    // ~1.56%. Interior quantiles are therefore approximate; min/max/sum/avg are tracked exactly.
    const QUANTILE_ABS_TOLERANCE: f64 = 2.0;

    #[test]
    fn default_sketch_is_empty() {
        let sketch = DDSketch::default();

        assert!(sketch.is_empty());
        assert_eq!(sketch.count(), 0);
        assert_eq!(sketch.bin_count(), 0);
        assert_eq!(sketch.min(), None);
        assert_eq!(sketch.max(), None);
        assert_eq!(sketch.sum(), None);
        assert_eq!(sketch.avg(), None);
        assert_eq!(sketch.quantile(0.5), None);
    }

    #[test]
    fn insert_tracks_exact_basic_statistics() {
        let mut sketch = DDSketch::default();
        sketch.insert(1.0);
        sketch.insert(2.0);
        sketch.insert(3.0);

        assert!(!sketch.is_empty());
        assert_eq!(sketch.count(), 3);
        // min/max/sum/avg are recorded from the raw samples, so they're exact (not bucketed).
        assert_eq!(sketch.min(), Some(1.0));
        assert_eq!(sketch.max(), Some(3.0));
        assert_eq!(sketch.sum(), Some(6.0));
        assert_eq!(sketch.avg(), Some(2.0));
    }

    #[test]
    fn insert_n_applies_the_weight_to_every_statistic() {
        let mut sketch = DDSketch::default();
        sketch.insert_n(10.0, 5);

        assert_eq!(sketch.count(), 5);
        assert_eq!(sketch.min(), Some(10.0));
        assert_eq!(sketch.max(), Some(10.0));
        assert_eq!(sketch.sum(), Some(50.0));
        assert_eq!(sketch.avg(), Some(10.0));
    }

    #[test]
    fn insert_many_matches_repeated_single_inserts() {
        let mut many = DDSketch::default();
        many.insert_many(&[1.0, 2.0, 3.0, 4.0]);

        let mut single = DDSketch::default();
        for v in [1.0, 2.0, 3.0, 4.0] {
            single.insert(v);
        }

        assert_eq!(many.count(), 4);
        // Feeding identical samples through either entry point must produce an identical sketch.
        assert_eq!(many, single);
    }

    #[test]
    fn quantile_at_extremes_returns_exact_min_and_max() {
        let mut sketch = DDSketch::default();
        for i in 1..=100 {
            sketch.insert(f64::from(i));
        }

        // q <= 0 and q >= 1 short-circuit to the exactly-tracked min/max.
        assert_eq!(sketch.quantile(0.0), Some(1.0));
        assert_eq!(sketch.quantile(-1.0), Some(1.0));
        assert_eq!(sketch.quantile(1.0), Some(100.0));
        assert_eq!(sketch.quantile(2.0), Some(100.0));
    }

    #[test]
    fn quantile_estimates_interior_percentiles_within_relative_accuracy() {
        let mut sketch = DDSketch::default();
        for i in 1..=100 {
            sketch.insert(f64::from(i));
        }

        let median = sketch.quantile(0.5).expect("non-empty sketch has a median");
        assert!(
            (median - 50.0).abs() <= QUANTILE_ABS_TOLERANCE,
            "median {} should be within {} of 50",
            median,
            QUANTILE_ABS_TOLERANCE
        );

        let p90 = sketch.quantile(0.9).expect("non-empty sketch has a p90");
        assert!(
            (p90 - 90.0).abs() <= QUANTILE_ABS_TOLERANCE,
            "p90 {} should be near 90",
            p90
        );
    }

    #[test]
    fn merge_combines_counts_bounds_and_sums() {
        let mut sketch1 = DDSketch::default();
        sketch1.insert(1.0);
        sketch1.insert(2.0);

        let mut sketch2 = DDSketch::default();
        sketch2.insert(3.0);
        sketch2.insert(4.0);

        sketch1.merge(&sketch2);

        assert_eq!(sketch1.count(), 4);
        assert_eq!(sketch1.min(), Some(1.0));
        assert_eq!(sketch1.max(), Some(4.0));
        assert_eq!(sketch1.sum(), Some(10.0));
        assert_eq!(sketch1.avg(), Some(2.5));
    }

    #[test]
    fn merge_produces_same_sketch_as_inserting_all_samples() {
        let mut merged = DDSketch::default();
        let mut left = DDSketch::default();
        let mut right = DDSketch::default();
        for i in 1..=50 {
            left.insert(f64::from(i));
        }
        for i in 51..=100 {
            right.insert(f64::from(i));
        }
        merged.merge(&left);
        merged.merge(&right);

        let mut all_at_once = DDSketch::default();
        for i in 1..=100 {
            all_at_once.insert(f64::from(i));
        }

        assert_eq!(merged.count(), all_at_once.count());
        assert_eq!(merged.bins(), all_at_once.bins());
    }

    #[test]
    fn clear_resets_the_sketch_to_empty() {
        let mut sketch = DDSketch::default();
        sketch.insert(1.0);
        sketch.insert(2.0);

        sketch.clear();

        assert!(sketch.is_empty());
        assert_eq!(sketch.count(), 0);
        assert_eq!(sketch.bin_count(), 0);
        assert_eq!(sketch.min(), None);
    }

    #[test]
    fn setters_override_the_tracked_statistics() {
        let mut sketch = DDSketch::default();
        sketch.insert(1.0);

        sketch.set_count(100);
        sketch.set_sum(42.0);
        sketch.set_avg(0.42);
        sketch.set_min(-5.0);
        sketch.set_max(500.0);

        assert_eq!(sketch.count(), 100);
        assert_eq!(sketch.sum(), Some(42.0));
        assert_eq!(sketch.avg(), Some(0.42));
        assert_eq!(sketch.min(), Some(-5.0));
        assert_eq!(sketch.max(), Some(500.0));
    }

    #[test]
    fn insert_interpolate_buckets_preserves_total_count() {
        let mut sketch = DDSketch::default();
        sketch
            .insert_interpolate_buckets(vec![
                Bucket {
                    upper_limit: 10.0,
                    count: 4,
                },
                Bucket {
                    upper_limit: 20.0,
                    count: 6,
                },
            ])
            .expect("well-formed buckets should interpolate successfully");

        // Every observation from the histogram buckets must be represented in the sketch.
        assert_eq!(sketch.count(), 10);
        let median = sketch.quantile(0.5).expect("non-empty sketch has a median");
        assert!(
            (5.0..=25.0).contains(&median),
            "median {} should land within the interpolated 10..20 range",
            median
        );
    }

    #[test]
    fn insert_interpolate_buckets_rejects_oversized_buckets() {
        let mut sketch = DDSketch::default();
        let result = sketch.insert_interpolate_buckets(vec![Bucket {
            upper_limit: 10.0,
            count: u64::from(u32::MAX) + 1,
        }]);

        assert_eq!(result, Err("bucket size greater than u32::MAX"));
    }

    #[test]
    fn dogsketch_round_trip_preserves_the_sketch() {
        let mut sketch = DDSketch::default();
        for i in 1..=25 {
            sketch.insert(f64::from(i));
        }

        let mut dogsketch = Dogsketch::new();
        sketch.merge_to_dogsketch(&mut dogsketch);
        let recovered = DDSketch::try_from(dogsketch).expect("dogsketch produced by merge_to_dogsketch is valid");

        assert_eq!(sketch, recovered);
    }

    #[test]
    fn try_from_dogsketch_rejects_mismatched_bin_vectors() {
        let mut dogsketch = Dogsketch::new();
        dogsketch.set_cnt(3);
        dogsketch.set_k(vec![1, 2]);
        dogsketch.set_n(vec![1]); // fewer counts than keys

        let result = DDSketch::try_from(dogsketch);
        assert_eq!(result, Err("k and n bin vectors have differing lengths"));
    }

    #[test]
    fn value_for_key_is_zero_at_the_origin_and_monotonic() {
        assert_eq!(DDSketch::value_for_key(0), 0.0);

        // Higher keys map to strictly larger representative values within the positive range.
        let low = DDSketch::value_for_key(100);
        let high = DDSketch::value_for_key(200);
        assert!(low > 0.0, "positive key should map to a positive value, got {}", low);
        assert!(high > low, "value_for_key should be monotonic: {} !> {}", high, low);
    }

    #[test]
    fn remap_mapping_uses_the_agent_gamma() {
        let mapping = DDSketch::remap_mapping();
        assert!(crate::common::float_eq(mapping.gamma(), DDSKETCH_CONF_GAMMA_V));
    }
}