randomize 2.2.0

//! Module for bounded randomization.
//!
//! This module has two types for ranges. They each have some quirks.
//!
//! * `DebiasedIntMulRange` can't be made to cover the full range of the integer
//!   type it's a range for. It also takes a division to construct, so you maybe
//!   don't want to be making them "on the fly" as much since division is slow.
//! * `BitmaskRange` is fairly cheap to construct and also can cover the full
//!   range of the integer type used. However, if the range "width" between low
//!   and high isn't very close to the next higher power of two (eg, `1..=20` is
//!   only 19 out of 32) then there will be many rejected values.
//!
//! My advice is to use `DebiasedIntMulRange` as much as you can, and fall back
//! to `BitmaskRange` only when necessary. The vital operation, `convert` is
//! provided via a trait so that functions can be generic over the exact range
//! type used.
//!
//! **Note:** The two different range types might or might not generate
//! different results with their `convert` function using a particular size of
//! range. All that is assured is that a uniform input will give uniform output.
//! The exact translation from input to output is not the same between the two
//! range types.

/// Fiddly to use, but totally gets you the minimum without branching.
macro_rules! branchless_min {
  ($x:expr, $y:expr, $u:ty, $zero:expr) => {
    $y ^ (($x ^ $y) & ($zero.wrapping_sub(($x < $y) as $u)))
  };
}

/// Fiddly to use, but totally gets you the maximum without branching.
macro_rules! branchless_max {
  ($x:expr, $y:expr, $u:ty, $zero:expr) => {
    $x ^ (($x ^ $y) & ($zero.wrapping_sub(($x < $y) as $u)))
  };
}

// Note(Lokathor): We want to test the macros, but their existence is an
// implementation detail only, so we provide the test for it right here instead
// of in the `tests/` folder
#[test]
fn test_branchless_min_and_max() {
  for x in core::u8::MIN..=core::u8::MAX {
    for y in core::u8::MIN..=core::u8::MAX {
      assert_eq!(branchless_min!(x, y, u8, 0u8), x.min(y));
      assert_eq!(branchless_max!(x, y, u8, 0u8), x.max(y));
    }
  }
}

/// Things that produce an unbiased value within some range.
///
/// The details of the setup and technique are left up to the implementing type.
pub trait RandomRange<T> {
  /// Given a uniformly distributed input integer, produce
  /// `Some(value_in_range)` or `None`.
  fn convert(&self, x: T) -> Option<T>;
}

/// A random range that uses multiplication and rejection to avoid bias.
///
/// This range doesn't support having the range be as large as the entire type
/// being ranged over. It must be at least 1 less than the type's range. Also,
/// this range requires a division to construct, to so it's much slower to build
/// on the fly than the `BitmaskRange`. Because the intermediate step needs to
/// do a multiply at 1 higher type size than the actual type, this isn't
/// implemented for `u128`.
///
/// All that said, because the output has to be rejected less often than with a
/// `BitmaskRange`, this has a faster overall speed.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct DebiasedIntMulRange<T> {
  // obviously might be zero
  base: T,
  // never zero
  range: T,
  // is zero if low==high
  threshold: T,
}
// TODO(Lokathor): In a 3.0 version we might want to make this type
// pre-specialized to each bit size with a proc-macro or something and then make
// the `range` field be a `NonZeroUFoo` field. It's unlikely that we need the
// extra niche though, so we can ignore that for now.

macro_rules! impl_DebiasedIntMulRange {
  ($u:ty, $zero:expr, $u_bits:expr,  $big_u:ty) => {
    impl DebiasedIntMulRange<$u> {
      /// Constructs a new range.
      ///
      /// ## Panics
      ///
      /// If you try to construct a range across the integer's full range this
      /// will cause a divide by zero panic.
      #[inline]
      pub const fn new(a: $u, b: $u) -> Self {
        let (low, high) = (branchless_min!(a, b, $u, $zero), branchless_max!(a, b, $u, $zero));
        let base = low;
        let range = high.wrapping_sub(low).wrapping_add(1);
        let threshold = $zero.wrapping_sub(range) % range;
        Self { base, range, threshold }
      }

      /// The inclusive low end of this range.
      #[inline]
      pub const fn low(&self) -> $u {
        self.base
      }

      /// The inclusive high end of this range.
      #[inline]
      pub const fn high(&self) -> $u {
        self.base.wrapping_add(self.range).wrapping_sub(1)
      }
    }

    impl RandomRange<$u> for DebiasedIntMulRange<$u> {
      #[inline]
      fn convert(&self, x: $u) -> Option<$u> {
        let m: $big_u = (x as $big_u).wrapping_mul(self.range as $big_u);
        let l: $u = m as $u;
        if l < self.threshold {
          None
        } else {
          Some((m >> $u_bits) as $u + self.base)
        }
      }
    }
  };
}

impl_DebiasedIntMulRange! {u8, 0u8, 8, u16}
impl_DebiasedIntMulRange! {u16, 0u16, 16, u32}
impl_DebiasedIntMulRange! {u32, 0u32, 32, u64}
impl_DebiasedIntMulRange! {u64, 0u64, 64, u128}

/// A random range that uses bit masking and rejection to avoid bias.
///
/// This range can be used for _any_ size range, even the full size of the type.
/// However, its speed depends on how close your range size is to the next
/// higher power of two. If there is a sizable gap you can get a lot of
/// rejections, which can make the total generation time very poor.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct BitmaskRange<T> {
  // obviously might be zero
  base: T,
  // zero if low==high
  range_minus_one: T,
  // never zero
  mask: T,
}
// TODO(Lokathor): As with the `DebiasedIntMulRange` type, we might want to
// generate this via proc-macro and use a `NonZeroUFoo` type for the `mask`
// field.

macro_rules! impl_BitmaskRange {
  ($u:ty, $zero:expr, $max_val:expr) => {
    impl BitmaskRange<$u> {
      /// Constructs a new range.
      #[inline]
      pub const fn new(a: $u, b: $u) -> Self {
        let (low, high) = (branchless_min!(a, b, $u, $zero), branchless_max!(a, b, $u, $zero));
        let base = low;
        let range_minus_one = high.wrapping_sub(low);
        let mask = $max_val >> (range_minus_one | 1).leading_zeros();
        Self { base, range_minus_one, mask }
      }

      /// The inclusive low end of this range.
      #[inline]
      pub const fn low(&self) -> $u {
        self.base
      }

      /// The inclusive high end of this range.
      #[inline]
      pub const fn high(&self) -> $u {
        self.base.wrapping_add(self.range_minus_one)
      }
    }

    impl RandomRange<$u> for BitmaskRange<$u> {
      #[inline]
      fn convert(&self, x: $u) -> Option<$u> {
        let out = x & self.mask;
        if out > self.range_minus_one {
          None
        } else {
          Some(self.base + out)
        }
      }
    }
  };
}

impl_BitmaskRange! {u8, 0u8, core::u8::MAX}
impl_BitmaskRange! {u16, 0u16, core::u16::MAX}
impl_BitmaskRange! {u32, 0u32, core::u32::MAX}
impl_BitmaskRange! {u64, 0u64, core::u64::MAX}
impl_BitmaskRange! {u128, 0u128, core::u128::MAX}