quantile-sketch 0.1.0

use proptest::{collection::vec, prelude::*, test_runner::Config as ProptestConfig};

use super::{ConcurrentDDSketch, Ordering, BLOCK_SIZE};

const ERROR: f64 = 0.01;
const PROPERTY_CASES: u32 = if cfg!(miri) { 8 } else { 8192 };
const MAX_VALUES: usize = if cfg!(miri) { 9 } else { 197 };

fn property_config() -> ProptestConfig {
    let config = ProptestConfig {
        cases: PROPERTY_CASES,
        ..ProptestConfig::default()
    };
    #[cfg(miri)]
    // Miri isolation blocks Proptest's filesystem-backed regression loader.
    let config = ProptestConfig {
        failure_persistence: None,
        ..config
    };
    config
}

fn datasets() -> impl Strategy<Value = Vec<u16>> {
    vec(any::<u16>(), 0..MAX_VALUES)
}

fn populated(values: &[u16]) -> ConcurrentDDSketch {
    let sketch = ConcurrentDDSketch::with_err_and_range(ERROR, 0.0, u16::MAX as f64);
    for &value in values {
        sketch.insert(value as f64);
    }
    sketch
}

fn quantile_probes(random: u16) -> [f64; 3] {
    [0.0, random as f64 / u16::MAX as f64, 1.0]
}

fn bulk_quantile_probes(random: u16) -> [f64; 4] {
    let q = random as f64 / u16::MAX as f64 / 2.0;
    if random & 1 == 0 {
        [0.0, q, q, 0.5]
    } else {
        [0.5, 0.5 + q, 0.5 + q, 1.0]
    }
}

fn rank_probes(random: u16) -> [f64; 3] {
    [0.0, random as f64, u16::MAX as f64]
}

fn nonzero_bins(sketch: &ConcurrentDDSketch) -> Vec<(usize, u64)> {
    sketch
        .allocated_blocks()
        .flat_map(|(block_index, block)| {
            block.counts.iter().enumerate().filter_map(move |(count_index, count)| {
                let count = count.load(Ordering::Relaxed);
                (count != 0).then_some((block_index * BLOCK_SIZE + count_index, count))
            })
        })
        .collect()
}

proptest! {
    #![proptest_config(property_config())]

    #[test]
    fn quantiles_match_an_exact_sorted_reference(values in datasets(), random_q in any::<u16>()) {
        let sketch = populated(&values);
        let mut sorted = values.clone();
        sorted.sort_unstable();

        for q in quantile_probes(random_q) {
            if sorted.is_empty() {
                prop_assert_eq!(sketch.quantile(q), None);
                continue;
            }

            let rank = (q * (sorted.len() - 1) as f64) as usize;
            let exact = sorted[rank] as f64;
            let estimate = sketch.quantile(q).expect("non-empty sketches have quantiles");
            if exact == 0.0 {
                prop_assert_eq!(estimate, 0.0);
            } else {
                let relative_error = (estimate - exact).abs() / exact;
                prop_assert!(relative_error <= ERROR * 1.001, "q={q}, exact={exact}, estimate={estimate}");
            }
        }

        let quantiles = bulk_quantile_probes(random_q);
        let mut estimates = [None; 4];
        sketch.quantiles(&quantiles, &mut estimates);
        for (&q, estimate) in quantiles.iter().zip(estimates) {
            super::assert_estimate_eq(estimate, sketch.quantile(q));
        }
    }

    #[test]
    fn percentile_ranks_match_an_exact_bucket_model(values in datasets(), random_rank in any::<u16>()) {
        let sketch = populated(&values);
        let mut previous = 0.0;

        for probe in rank_probes(random_rank) {
            if values.is_empty() {
                prop_assert_eq!(sketch.percentile_rank(probe), None);
                continue;
            }

            let target = sketch.coord(probe);
            let included = values.iter().filter(|&&value| sketch.coord(value as f64) <= target).count();
            let expected = included as f64 / values.len() as f64;
            let actual = sketch.percentile_rank(probe).expect("non-empty sketches have ranks");
            prop_assert_eq!(actual, expected);
            prop_assert!(actual >= previous, "ranks must be monotonic");
            previous = actual;
        }
    }

    #[test]
    fn merging_matches_direct_insertion(
        values in datasets(),
        split_code in any::<u8>(),
        random_q in any::<u16>(),
        random_rank in any::<u16>(),
    ) {
        let split = split_code as usize % (values.len() + 1);
        let merged = populated(&values[..split]);
        let source = populated(&values[split..]);
        merged.merge(&source).expect("identical configurations can be merged");

        let direct = populated(&values);
        prop_assert_eq!(nonzero_bins(&merged), nonzero_bins(&direct));
        for q in quantile_probes(random_q) {
            super::assert_estimate_eq(merged.quantile(q), direct.quantile(q));
        }
        for probe in rank_probes(random_rank) {
            prop_assert_eq!(merged.percentile_rank(probe), direct.percentile_rank(probe));
        }
    }
}