use crate::TurboRand;
pub(crate) struct IncreasingUniformIter<R: TurboRand> {
rng: R,
n: u64,
chunk: u64,
chunk_remaining: u8,
len: usize,
}
impl<R: TurboRand> IncreasingUniformIter<R> {
#[inline]
pub(crate) fn new(rng: R, n: u64, len: usize) -> Self {
let chunk_remaining = u8::from(n == 0);
Self {
rng,
n,
chunk: 0,
chunk_remaining,
len,
}
}
#[inline]
fn next_swap_index(&mut self) -> usize {
let next_n = self.n + 1;
let next_chunk_remaining = self.chunk_remaining.checked_sub(1).unwrap_or_else(|| {
let (bound, remaining) = calculate_bound_u64(next_n);
self.chunk = self.rng.u64(..bound);
remaining - 1
});
let result = if next_chunk_remaining == 0 {
self.chunk as usize
} else {
let random = self.chunk % next_n;
self.chunk /= next_n;
random as usize
};
self.chunk_remaining = next_chunk_remaining;
self.n = next_n;
result
}
}
impl<R: TurboRand> Iterator for IncreasingUniformIter<R> {
type Item = (usize, usize);
#[inline]
fn next(&mut self) -> Option<Self::Item> {
if self.len == (self.n as usize) {
None
} else {
Some((self.n as usize, self.next_swap_index()))
}
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
let len = self.len - self.n as usize;
(len, Some(len))
}
}
impl<R: TurboRand> ExactSizeIterator for IncreasingUniformIter<R> {
#[inline]
fn len(&self) -> usize {
self.size_hint().0
}
}
#[inline]
fn calculate_bound_u64(min: u64) -> (u64, u8) {
debug_assert!(min > 1);
#[inline]
const fn inner(min: u64) -> (u64, u8) {
let mut product = min;
let mut current = min + 1;
while let Some(p) = product.checked_mul(current) {
product = p;
current += 1;
}
let count = (current - min) as u8;
(product, count)
}
const RESULT2: (u64, u8) = inner(2);
match min {
2 => RESULT2,
min => inner(min),
}
}