use parity_scale_codec::{Decode, DecodeWithMemTracking, Encode};
use scale_info::TypeInfo;
use sp_runtime::traits::{Hash, TrailingZeroInput};
#[derive(Encode, Decode, DecodeWithMemTracking, TypeInfo)]
pub struct RandomNumberGenerator<Hashing: Hash> {
current: Hashing::Output,
offset: u32,
}
impl<Hashing: Hash> RandomNumberGenerator<Hashing> {
pub fn new(seed: Hashing::Output) -> Self {
Self { current: seed, offset: 0 }
}
fn offset(&self) -> usize {
self.offset as usize
}
pub fn pick_u32(&mut self, max: u32) -> u32 {
let needed = (4 - max.leading_zeros() / 8) as usize;
let top = ((1 << (needed as u64 * 8)) / (max as u64 + 1) * (max as u64 + 1) - 1) as u32;
loop {
if self.offset() + needed > self.current.as_ref().len() {
self.current = <Hashing as Hash>::hash(self.current.as_ref());
self.offset = 0;
}
let data = &self.current.as_ref()[self.offset()..self.offset() + needed];
self.offset += needed as u32;
let raw = u32::decode(&mut TrailingZeroInput::new(data)).unwrap_or(0);
if raw <= top {
break if max < u32::MAX { raw % (max + 1) } else { raw };
}
}
}
pub fn pick_non_zero_u32(&mut self, max: u32) -> u32 {
self.pick_u32(max - 1) + 1
}
pub fn pick_usize(&mut self, max: usize) -> usize {
self.pick_u32(max as u32) as usize
}
pub fn pick_item<'a, T>(&mut self, items: &'a [T]) -> Option<&'a T> {
if items.is_empty() {
None
} else {
Some(&items[self.pick_usize(items.len() - 1)])
}
}
}
#[cfg(test)]
mod tests {
use super::RandomNumberGenerator;
use sp_runtime::traits::{BlakeTwo256, Hash};
#[test]
fn does_not_panic_on_max() {
let seed = BlakeTwo256::hash(b"Fourty-two");
let _random = RandomNumberGenerator::<BlakeTwo256>::new(seed).pick_u32(u32::MAX);
}
}