hdrhistogram 7.6.0

//! This is used in tests (both unit tests and integration tests) to provide useful distributions
//! of random numbers.

use rand::Rng;
use rand::distr::Distribution;
use rand::distr::uniform::Uniform;

/// Smallest number in our varint encoding that takes the given number of bytes
pub fn smallest_number_in_n_byte_varint(byte_length: usize) -> u64 {
    assert!((1..=9).contains(&byte_length));

    match byte_length {
        1 => 0,
        // one greater than the largest of the previous length
        _ => largest_number_in_n_byte_varint(byte_length - 1) + 1,
    }
}

/// Largest number in our varint encoding that takes the given number of bytes
pub fn largest_number_in_n_byte_varint(byte_length: usize) -> u64 {
    assert!((1..=9).contains(&byte_length));

    match byte_length {
        9 => u64::MAX,
        _ => largest_number_in_7_bit_chunk(byte_length - 1),
    }
}

/// The largest in the set of numbers that have at least 1 bit set in the n'th chunk of 7 bits.
fn largest_number_in_7_bit_chunk(chunk_index: usize) -> u64 {
    // Our 9-byte varints do different encoding in the last byte, so we don't handle them here
    assert!(chunk_index <= 7);

    // 1 in every bit below the lowest bit in this chunk
    let lower_bits = match chunk_index {
        0 => 0,
        _ => largest_number_in_7_bit_chunk(chunk_index - 1),
    };

    // 1 in every bit in this chunk
    let this_chunk = 0x7F_u64 << (chunk_index * 7);

    lower_bits | this_chunk
}

// Evenly distributed random numbers end up biased heavily towards longer encoded byte lengths:
// there are a lot more large numbers than there are small (duh), but for exercising serialization
// code paths, we'd like many at all byte lengths. This is also arguably more representative of
// real data. This should emit values whose varint lengths are uniformly distributed across the
// whole length range (1 to 9).
pub struct RandomVarintEncodedLengthIter<R: Rng> {
    ranges: [Uniform<u64>; 9],
    range_for_picking_range: Uniform<usize>,
    rng: R,
}

impl<R: Rng> RandomVarintEncodedLengthIter<R> {
    pub fn new(rng: R) -> RandomVarintEncodedLengthIter<R> {
        RandomVarintEncodedLengthIter {
            // Each varint length has a nonempty value range, so constructing
            // the `Uniform` (fallible as of rand 0.9) cannot fail here.
            ranges: std::array::from_fn(|i| {
                let byte_length = i + 1;
                let smallest = smallest_number_in_n_byte_varint(byte_length);
                let largest = largest_number_in_n_byte_varint(byte_length);
                Uniform::new_inclusive(smallest, largest).expect("range is non-empty")
            }),
            range_for_picking_range: Uniform::new(0, 9).expect("range is non-empty"),
            rng,
        }
    }
}

impl<R: Rng> Iterator for RandomVarintEncodedLengthIter<R> {
    type Item = u64;

    fn next(&mut self) -> Option<Self::Item> {
        // pick the range we'll use
        let value_range = self.ranges[self.range_for_picking_range.sample(&mut self.rng)];

        Some(value_range.sample(&mut self.rng))
    }
}

#[test]
fn largest_number_in_7_bit_chunk_correct() {
    // 8 chunks (indices 0-7) of 7 bits gets you to 56 bits. Last byte in varint is handled
    // differently, so we don't test that here.
    for i in 0..8 {
        let largest = largest_number_in_7_bit_chunk(i);
        assert_eq!((i as u32 + 1) * 7, largest.count_ones());

        assert_eq!(64 - ((i as u32) + 1) * 7, largest.leading_zeros());
        // any larger and it will be in the next chunk
        assert_eq!(largest.leading_zeros() - 1, (largest + 1).leading_zeros());
    }
}