histogram 1.6.4

use crate::Error;
use core::ops::RangeInclusive;

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

/// The configuration of a histogram which determines the bucketing strategy and
/// therefore the relative error and memory utilization of a histogram.
/// * `grouping_power` - controls the number of buckets that are used to span
///   consecutive powers of two. Lower values result in less memory usage since
///   fewer buckets will be created. However, this will result in larger
///   relative error as each bucket represents a wider range of values.
/// * `max_value_power` - controls the largest value which can be stored in the
///   histogram. `2^(max_value_power) - 1` is the inclusive upper bound for the
///   representable range of values.
///
/// # How to choose parameters for your data
/// Please see <https://observablehq.com/@iopsystems/h2histogram> for an
/// in-depth discussion about the bucketing strategy and an interactive
/// calculator that lets you explore how these parameters result in histograms
/// with varying error guarantees and memory utilization requirements.
///
/// # The short version
/// ## Grouping Power
/// `grouping_power` should be set such that `2^(-1 * grouping_power)` is an
/// acceptable relative error. Rephrased, we can plug-in the acceptable
/// relative error into `grouping_power = ceil(log2(1/e))`. For example, if we
/// want to limit the relative error to 0.1% (0.001), we should set
/// `grouping_power = 10`. This gives a relative error bound of
/// `100 / 2^10 = 0.09765625%`.
///
/// ## Max Value Power
/// `max_value_power` should be the closest power of 2 that is larger than the
/// largest value you expect in your data. If your only guarantee is that the
/// values are all `u64`, then setting this to `64` may be reasonable if you
/// can tolerate a bit of relative error.
///
/// ## Resulting size
///
/// If we want to allow any value in a range of unsigned types, the amount of
/// memory for the histogram is approximately:
///
/// | grouping_power | error |     u16 |     u32 |     u64 |
/// |----------------|-------|---------|---------|---------|
/// |     2 |   25% | 0.6 KiB |   1 KiB |   2 KiB |
/// |     3 | 12.5% |   1 KiB |   2 KiB |   4 KiB |
/// |     4 | 6.25% |   2 KiB |   4 KiB |   8 KiB |
/// |     5 | 3.13% |   3 KiB |   7 KiB |  15 KiB |
/// |     6 | 1.56% |   6 KiB |  14 KiB |  30 KiB |
/// |     7 | .781% |  10 KiB |  26 KiB |  58 KiB |
/// |     8 | .391% |  18 KiB |  50 KiB | 114 KiB |
/// |     9 | .195% |  32 KiB |  96 KiB | 224 KiB |
/// |    10 | .098% |  56 KiB | 184 KiB | 440 KiB |
/// |    11 | .049% |  96 KiB | 352 KiB | 864 KiB |
/// |    12 | .025% | 160 KiB | 672 KiB | 1.7 MiB |
///
/// These figures describe dense u64 counter storage. Halve counter storage for
/// `Histogram32` and `AtomicHistogram32`. Sparse/cumulative storage also includes
/// indices and vector capacity; changing counter width does not halve those costs.
///
/// # Constraints:
/// * `max_value_power` must be in the range `0..=64`
/// * `max_value_power` must be greater than `grouping_power`
/// * The total bucket count must fit in `u32`; larger geometries return [`Error::Overflow`].
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(Serialize))]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct Config {
    max: u64,
    grouping_power: u8,
    max_value_power: u8,
    cutoff_power: u8,
    cutoff_value: u64,
    lower_bin_count: u32,
    upper_bin_divisions: u32,
    upper_bin_count: u32,
}

#[cfg(feature = "serde")]
impl<'de> Deserialize<'de> for Config {
    fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
        // Preserve the original struct name, field names, and field order for
        // both map-based and sequence-based serialization formats.
        #[derive(Deserialize)]
        struct Config {
            max: u64,
            grouping_power: u8,
            max_value_power: u8,
            cutoff_power: u8,
            cutoff_value: u64,
            lower_bin_count: u32,
            upper_bin_divisions: u32,
            upper_bin_count: u32,
        }
        let wire = Config::deserialize(deserializer)?;
        let expected = Self::new(wire.grouping_power, wire.max_value_power)
            .map_err(serde::de::Error::custom)?;
        // histogram 0.8.0 through 0.11.3 stored `max` as `2^max_value_power`
        // when `max_value_power < 64`; 0.11.4 changed it to the inclusive
        // bound `2^max_value_power - 1`. Accept that encoding as the current
        // one so data serialized by 0.8.2 through 0.11.3 still decodes. The
        // `< 64` check must come first: it keeps the shift in range.
        let legacy_max = wire.max_value_power < 64 && wire.max == 1_u64 << wire.max_value_power;
        let decoded = Self {
            max: if legacy_max { expected.max } else { wire.max },
            grouping_power: wire.grouping_power,
            max_value_power: wire.max_value_power,
            cutoff_power: wire.cutoff_power,
            cutoff_value: wire.cutoff_value,
            lower_bin_count: wire.lower_bin_count,
            upper_bin_divisions: wire.upper_bin_divisions,
            upper_bin_count: wire.upper_bin_count,
        };
        if decoded != expected {
            return Err(serde::de::Error::custom(
                "serialized config has inconsistent derived fields",
            ));
        }
        Ok(expected)
    }
}

impl Config {
    /// Create a new histogram `Config` from the parameters. See the struct
    /// documentation [`crate::Config`] for the meaning of the parameters and
    /// their constraints.
    pub const fn new(grouping_power: u8, max_value_power: u8) -> Result<Self, Error> {
        // we only allow values up to 2^64
        if max_value_power > 64 {
            return Err(Error::MaxPowerTooHigh);
        }

        // check that the other parameters make sense together
        if grouping_power >= max_value_power {
            return Err(Error::MaxPowerTooLow);
        }

        // Bucket indices and the total bucket count use u32. Check the full
        // geometry before the shifts, products, or narrowing conversions below.
        let Some(upper_bin_divisions) = 1_u32.checked_shl(grouping_power as u32) else {
            return Err(Error::Overflow);
        };
        if upper_bin_divisions
            .checked_mul((max_value_power - grouping_power + 1) as u32)
            .is_none()
        {
            return Err(Error::Overflow);
        }

        // the cutoff is the point at which the linear range divisions and the
        // logarithmic range subdivisions diverge.
        //
        // for example:
        // when a = 0, the linear range has bins with width 1.
        // if b = 7 the logarithmic range has 128 subdivisions.
        // this means that for 0..128 we must be representing the values exactly
        // but we also represent 128..256 exactly since the subdivisions divide
        // that range into bins with the same width as the linear portion.
        //
        // therefore our cutoff power = a + b + 1

        // note: because a + b must be less than n which is a u8, a + b + 1 must
        // be less than or equal to u8::MAX. This means our cutoff power will
        // always fit in a u8
        let cutoff_power = grouping_power + 1;
        let cutoff_value = 2_u64.pow(cutoff_power as u32);

        let max = if max_value_power == 64 {
            u64::MAX
        } else {
            2_u64.pow(max_value_power as u32) - 1
        };

        let lower_bin_count = cutoff_value as u32;
        let upper_bin_count = (max_value_power - cutoff_power) as u32 * upper_bin_divisions;

        Ok(Self {
            max,
            grouping_power,
            max_value_power,
            cutoff_power,
            cutoff_value,
            lower_bin_count,
            upper_bin_divisions,
            upper_bin_count,
        })
    }

    /// Returns the grouping power that was used to create this configuration.
    #[inline]
    pub const fn grouping_power(&self) -> u8 {
        self.grouping_power
    }

    /// Returns the max value power that was used to create this configuration.
    #[inline]
    pub const fn max_value_power(&self) -> u8 {
        self.max_value_power
    }

    /// Returns the relative error (in percentage) of this configuration. This
    /// only applies to the logarithmic bins of the histogram (linear bins have
    /// a width of 1 and no error). For histograms with no logarithmic bins,
    /// error for the entire histogram is zero.
    pub fn error(&self) -> f64 {
        match self.grouping_power == self.max_value_power - 1 {
            true => 0.0,
            false => 100.0 / 2_u64.pow(self.grouping_power as u32) as f64,
        }
    }

    /// Return the total number of buckets needed for this config.
    #[inline]
    pub const fn total_buckets(&self) -> usize {
        (self.lower_bin_count + self.upper_bin_count) as usize
    }

    /// The bucket index of `value` with no range check.
    ///
    /// Below the cutoff the index is the value. Above it, with `s` the number
    /// of low bits a bucket discards (the value's bit length minus
    /// `grouping_power + 1`), the index is `(s << grouping_power) + (value >> s)`:
    /// one block of 2^`grouping_power` buckets per further power of two.
    ///
    /// The result is monotone in `value` and fits in `u64` (`s` is at most 63
    /// and `grouping_power` at most 31), even above `max`.
    /// Values above `max` map to indices at or beyond
    /// [`total_buckets`](Self::total_buckets).
    /// Callers using a slice lookup as the range check must use storage of
    /// exactly that length.
    #[inline]
    pub(crate) fn bucket_index(&self, value: u64) -> u64 {
        if value < self.cutoff_value {
            return value;
        }
        // Setting the low bit preserves the leading-zero count here and lets
        // the compiler treat the operand as nonzero.
        let s = (63 - (value | 1).leading_zeros()) - self.grouping_power as u32;
        ((s as u64) << self.grouping_power) + (value >> s)
    }

    /// Converts a value to a bucket index. Returns an error if the value is
    /// outside of the range for the config.
    #[inline]
    pub(crate) fn value_to_index(&self, value: u64) -> Result<usize, Error> {
        let index = self.bucket_index(value);
        if index >= self.total_buckets() as u64 {
            return Err(Error::OutOfRange);
        }
        Ok(index as usize)
    }

    /// Convert a bucket index to a lower bound.
    #[inline]
    pub(crate) fn index_to_lower_bound(&self, index: usize) -> u64 {
        let g = index as u64 >> self.grouping_power;
        let h = index as u64 - g * (1 << self.grouping_power);

        if g < 1 {
            h
        } else {
            (1 << (self.grouping_power as u64 + g - 1)) + (1 << (g - 1)) * h
        }
    }

    /// Convert a bucket index to a upper inclusive bound.
    #[inline]
    pub(crate) fn index_to_upper_bound(&self, index: usize) -> u64 {
        if index as u32 == self.lower_bin_count + self.upper_bin_count - 1 {
            return self.max;
        }
        let g = index as u64 >> self.grouping_power;
        let h = index as u64 - g * (1 << self.grouping_power) + 1;

        if g < 1 {
            h - 1
        } else {
            (1 << (self.grouping_power as u64 + g - 1)) + (1 << (g - 1)) * h - 1
        }
    }

    /// Convert a bucket index to a range.
    #[inline]
    pub(crate) fn index_to_range(&self, index: usize) -> RangeInclusive<u64> {
        self.index_to_lower_bound(index)..=self.index_to_upper_bound(index)
    }
}

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

    #[cfg(target_pointer_width = "64")]
    #[test]
    fn sizes() {
        assert_eq!(std::mem::size_of::<Config>(), 32);
    }

    #[test]
    // Test that the number of buckets matches the expected count
    fn total_buckets() {
        let config = Config::new(2, 64).unwrap();
        assert_eq!(config.total_buckets(), 252);

        let config = Config::new(7, 64).unwrap();
        assert_eq!(config.total_buckets(), 7424);

        let config = Config::new(14, 64).unwrap();
        assert_eq!(config.total_buckets(), 835_584);

        let config = Config::new(2, 4).unwrap();
        assert_eq!(config.total_buckets(), 12);
    }

    /// Indexing oracle with explicit log-bin offsets and a separate maximum check.
    fn reference_index(config: &Config, value: u64) -> Option<usize> {
        if value < config.cutoff_value {
            return Some(value as usize);
        }
        if value > config.max {
            return None;
        }
        let power = 63 - value.leading_zeros();
        let log_bin = power - config.cutoff_power as u32;
        let offset = (value - (1 << power)) >> (power - config.grouping_power as u32);
        Some(
            (config.lower_bin_count + log_bin * config.upper_bin_divisions + offset as u32)
                as usize,
        )
    }

    #[test]
    // Sampled values check index agreement, monotonicity, and rejection above
    // the maximum for valid geometries with grouping powers 0..=15.
    fn bucket_index_matches_reference_on_all_boundaries() {
        for max_value_power in 1..=64u8 {
            for grouping_power in 0..max_value_power.min(16) {
                let config = Config::new(grouping_power, max_value_power).unwrap();
                let total = config.total_buckets() as u64;
                let mut values = vec![0u64, 1, 2, 3, config.max, u64::MAX];
                for k in 0..64u32 {
                    let p = 1u64 << k;
                    values.extend([p - 1, p, p + 1, p + (p >> 1)]);
                }
                if config.max < u64::MAX {
                    values.push(config.max + 1);
                }
                values.sort_unstable();
                let mut previous = 0u64;
                for value in values {
                    let index = config.bucket_index(value);
                    assert!(
                        index >= previous,
                        "gp{grouping_power}/max{max_value_power} value {value} not monotone"
                    );
                    previous = index;
                    match reference_index(&config, value) {
                        Some(expected) => {
                            assert_eq!(
                                index, expected as u64,
                                "gp{grouping_power}/max{max_value_power} value {value}"
                            );
                            assert_eq!(config.value_to_index(value), Ok(expected));
                        }
                        None => {
                            assert!(
                                index >= total,
                                "gp{grouping_power}/max{max_value_power} value {value} indexed {index} inside {total} buckets"
                            );
                            assert_eq!(config.value_to_index(value), Err(Error::OutOfRange));
                        }
                    }
                }
            }
        }
    }

    #[test]
    // Test value to index conversions
    fn value_to_idx() {
        let config = Config::new(7, 64).unwrap();
        assert_eq!(config.value_to_index(0), Ok(0));
        assert_eq!(config.value_to_index(1), Ok(1));
        assert_eq!(config.value_to_index(256), Ok(256));
        assert_eq!(config.value_to_index(257), Ok(256));
        assert_eq!(config.value_to_index(258), Ok(257));
        assert_eq!(config.value_to_index(512), Ok(384));
        assert_eq!(config.value_to_index(515), Ok(384));
        assert_eq!(config.value_to_index(516), Ok(385));
        assert_eq!(config.value_to_index(1024), Ok(512));
        assert_eq!(config.value_to_index(1031), Ok(512));
        assert_eq!(config.value_to_index(1032), Ok(513));
        assert_eq!(config.value_to_index(u64::MAX - 1), Ok(7423));
        assert_eq!(config.value_to_index(u64::MAX), Ok(7423));
    }

    #[test]
    // Test index to lower bound conversion
    fn idx_to_lower_bound() {
        let config = Config::new(7, 64).unwrap();
        assert_eq!(config.index_to_lower_bound(0), 0);
        assert_eq!(config.index_to_lower_bound(1), 1);
        assert_eq!(config.index_to_lower_bound(256), 256);
        assert_eq!(config.index_to_lower_bound(384), 512);
        assert_eq!(config.index_to_lower_bound(512), 1024);
        assert_eq!(
            config.index_to_lower_bound(7423),
            18_374_686_479_671_623_680
        );
    }

    #[test]
    // Test index to upper bound conversion
    fn idx_to_upper_bound() {
        let config = Config::new(7, 64).unwrap();
        assert_eq!(config.index_to_upper_bound(0), 0);
        assert_eq!(config.index_to_upper_bound(1), 1);
        assert_eq!(config.index_to_upper_bound(256), 257);
        assert_eq!(config.index_to_upper_bound(384), 515);
        assert_eq!(config.index_to_upper_bound(512), 1031);
        assert_eq!(config.index_to_upper_bound(7423), u64::MAX);
    }

    #[test]
    // Test index to range conversion
    fn idx_to_range() {
        let config = Config::new(7, 64).unwrap();
        assert_eq!(config.index_to_range(0), 0..=0);
        assert_eq!(config.index_to_range(1), 1..=1);
        assert_eq!(config.index_to_range(256), 256..=257);
        assert_eq!(config.index_to_range(384), 512..=515);
        assert_eq!(config.index_to_range(512), 1024..=1031);
        assert_eq!(
            config.index_to_range(7423),
            18_374_686_479_671_623_680..=u64::MAX
        );
    }
}