use alloc::vec::Vec;
use core::num::NonZeroU32;
use super::{Sampler, Xoshiro256};
#[derive(Debug)]
pub(crate) struct FragmentChooser {
fragment_count: NonZeroU32,
checksum: u32,
degrees: Sampler,
}
impl FragmentChooser {
pub(crate) fn new(fragment_count: NonZeroU32, checksum: u32) -> Self {
let weights = (1..=fragment_count.get())
.map(|x| 1.0 / f64::from(x))
.collect();
Self {
fragment_count,
checksum,
degrees: Sampler::new(weights),
}
}
#[allow(
clippy::cast_possible_truncation,
reason = "degree <= fragment_count, which is u32-bounded"
)]
pub(crate) fn choose(&self, sequence: NonZeroU32) -> Vec<u32> {
let sequence = sequence.get();
let k = self.fragment_count.get();
if sequence <= k {
return alloc::vec![sequence - 1];
}
let mut seed = [0_u8; 8];
seed[0..4].copy_from_slice(&sequence.to_be_bytes());
seed[4..8].copy_from_slice(&self.checksum.to_be_bytes());
let mut xoshiro = Xoshiro256::from(seed.as_slice());
let degree = self.degrees.next(&mut xoshiro) + 1;
let remaining: Vec<u32> = (0..k).collect();
let mut indexes = xoshiro.shuffled(remaining, degree as usize);
indexes.sort_unstable();
indexes
}
}
#[must_use]
#[allow(
clippy::cast_possible_truncation,
reason = "fragment indexes are u32-bounded and fit usize on all supported targets"
)]
#[cfg(test)]
pub(crate) fn choose_fragments(
sequence: NonZeroU32,
fragment_count: NonZeroU32,
checksum: u32,
) -> Vec<usize> {
FragmentChooser::new(fragment_count, checksum)
.choose(sequence)
.into_iter()
.map(|i| i as usize)
.collect()
}