library/compiler-builtins/libm-test/src/generate/random.rs

use std::env;
use std::ops::RangeInclusive;
use std::sync::LazyLock;

use libm::support::{Float, Int};
use rand::distr::uniform::SampleUniform;
use rand::distr::{Alphanumeric, StandardUniform};
use rand::prelude::Distribution;
use rand::{RngExt, SeedableRng};
use rand_chacha::ChaCha8Rng;

use crate::CheckCtx;
use crate::generate::{KnownSize, product2, product3};
use crate::num::full_range;
use crate::run_cfg::{int_range, iteration_count};

pub(crate) const SEED_ENV: &str = "LIBM_SEED";

pub static SEED: LazyLock<[u8; 32]> = LazyLock::new(|| {
    let s = env::var(SEED_ENV).unwrap_or_else(|_| {
        let mut rng = rand::rng();
        (0..32).map(|_| rng.sample(Alphanumeric) as char).collect()
    });

    s.as_bytes().try_into().unwrap_or_else(|_| {
        panic!("Seed must be 32 characters, got `{s}`");
    })
});

/// Generate a sequence of random values of this type.
pub trait RandomInput: Sized {
    fn get_cases(ctx: &CheckCtx) -> (impl Iterator<Item = Self> + Send, u64);
}

/// Generate a sequence of deterministically random floats.
fn random_floats<F: Float>(count: u64) -> impl Iterator<Item = F> + Clone
where
    StandardUniform: Distribution<F::Int>,
{
    let mut rng = ChaCha8Rng::from_seed(*SEED);

    // Generate integers to get a full range of bitpatterns (including NaNs), then convert back
    // to the float type.
    (0..count).map(move |_| F::from_bits(rng.random::<F::Int>()))
}

/// Generate a sequence of deterministically random `i32`s within a specified range.
fn random_ints<I>(count: u64, range: RangeInclusive<I>) -> impl Iterator<Item = I> + Clone
where
    I: Int + SampleUniform,
{
    let mut rng = ChaCha8Rng::from_seed(*SEED);
    (0..count).map(move |_| rng.random_range::<I, _>(range.clone()))
}

macro_rules! impl_random_input {
    ($fty:ty) => {
        impl RandomInput for ($fty,) {
            fn get_cases(ctx: &CheckCtx) -> (impl Iterator<Item = Self>, u64) {
                let count = iteration_count(ctx, 0);
                let iter = random_floats(count).map(|f: $fty| (f,));
                (iter, count)
            }
        }

        impl RandomInput for ($fty, $fty) {
            fn get_cases(ctx: &CheckCtx) -> (impl Iterator<Item = Self>, u64) {
                let count0 = iteration_count(ctx, 0);
                let count1 = iteration_count(ctx, 1);
                let iter0 = random_floats(count0);
                let iter1 = random_floats(count1);
                let iter = product2(iter0, iter1);
                let count = count0.strict_mul(count1);
                (iter, count)
            }
        }

        impl RandomInput for ($fty, $fty, $fty) {
            fn get_cases(ctx: &CheckCtx) -> (impl Iterator<Item = Self>, u64) {
                let count0 = iteration_count(ctx, 0);
                let count1 = iteration_count(ctx, 1);
                let count2 = iteration_count(ctx, 2);
                let iter0 = random_floats(count0);
                let iter1 = random_floats(count1);
                let iter2 = random_floats(count2);
                let iter = product3(iter0, iter1, iter2);
                let count = count0.strict_mul(count1).strict_mul(count2);
                (iter, count)
            }
        }

        impl RandomInput for (i32, $fty) {
            fn get_cases(ctx: &CheckCtx) -> (impl Iterator<Item = Self>, u64) {
                let count0 = iteration_count(ctx, 0);
                let count1 = iteration_count(ctx, 1);
                let range0 = int_range::<i32>(ctx, 0).unwrap_or(full_range());
                let iter0 = random_ints(count0, range0);
                let iter1 = random_floats(count1);
                let iter = product2(iter0, iter1);
                let count = count0.strict_mul(count1);
                (iter, count)
            }
        }

        impl RandomInput for ($fty, i32) {
            fn get_cases(ctx: &CheckCtx) -> (impl Iterator<Item = Self>, u64) {
                let count0 = iteration_count(ctx, 0);
                let count1 = iteration_count(ctx, 1);
                let range1 = int_range::<i32>(ctx, 1).unwrap_or(full_range());
                let iter0 = random_floats(count0);
                let iter1 = random_ints(count1, range1.clone());
                let iter = product2(iter0, iter1);
                let count = count0.strict_mul(count1);
                (iter, count)
            }
        }
    };
}

#[cfg(f16_enabled)]
impl_random_input!(f16);
impl_random_input!(f32);
impl_random_input!(f64);
#[cfg(f128_enabled)]
impl_random_input!(f128);

macro_rules! impl_random_input_int {
    (@skip_u32 $ity:ty) => {
        impl RandomInput for ($ity,) {
            fn get_cases(ctx: &CheckCtx) -> (impl Iterator<Item = Self>, u64) {
                let count = iteration_count(ctx, 0);
                let range = int_range::<$ity>(ctx, 0).unwrap_or(full_range());
                let iter = random_ints(count, range).map(|f: $ity| (f,));
                (iter, count)
            }
        }

        impl RandomInput for ($ity, $ity) {
            fn get_cases(ctx: &CheckCtx) -> (impl Iterator<Item = Self>, u64) {
                let count0 = iteration_count(ctx, 0);
                let count1 = iteration_count(ctx, 1);
                let range0 = int_range::<$ity>(ctx, 0).unwrap_or(full_range());
                let range1 = int_range::<$ity>(ctx, 1).unwrap_or(full_range());
                let iter0 = random_ints(count0, range0);
                let iter1 = random_ints(count1, range1.clone());
                let iter = product2(iter0, iter1);
                (iter, count0 * count1)
            }
        }
    };
    ($ity:ty) => {
        impl_random_input_int!(@skip_u32 $ity);

        impl RandomInput for ($ity, u32) {
            fn get_cases(ctx: &CheckCtx) -> (impl Iterator<Item = Self>, u64) {
                let count0 = iteration_count(ctx, 0);
                let count1 = iteration_count(ctx, 1);
                let range0 = int_range::<$ity>(ctx, 0).unwrap_or(full_range());
                let range1 = int_range::<u32>(ctx, 1).unwrap_or(full_range());
                let iter0 = random_ints(count0, range0);
                let iter1 = random_ints(count1, range1.clone());
                let iter = product2(iter0, iter1);
                (iter, count0 * count1)
            }
        }
    };
}

impl_random_input_int!(i32);
impl_random_input_int!(i64);
impl_random_input_int!(i128);
impl_random_input_int!(@skip_u32 u32);
impl_random_input_int!(u64);
impl_random_input_int!(u128);

/// Create a test case iterator.
pub fn get_test_cases<RustArgs: RandomInput>(
    ctx: &CheckCtx,
) -> (
    impl Iterator<Item = RustArgs> + Send + use<'_, RustArgs>,
    u64,
) {
    let (iter, count) = RustArgs::get_cases(ctx);

    // Wrap in `KnownSize` so we get an assertion if the cuunt is wrong.
    (KnownSize::new(iter, count), count)
}