use core::ops::{Bound, RangeBounds};
use crate::{
env::internal,
unwrap::{UnwrapInfallible, UnwrapOptimized},
Bytes, BytesN, Env, IntoVal, Vec,
};
pub struct Prng {
env: Env,
}
impl Prng {
pub(crate) fn new(env: &Env) -> Prng {
Prng { env: env.clone() }
}
pub fn env(&self) -> &Env {
&self.env
}
pub fn seed(&self, seed: Bytes) {
let env = self.env();
assert_in_contract!(env);
internal::Env::prng_reseed(env, seed.into()).unwrap_infallible();
}
pub fn fill<T>(&self, v: &mut T)
where
T: Fill + ?Sized,
{
v.fill(self);
}
pub fn gen<T>(&self) -> T
where
T: Gen,
{
T::gen(self)
}
pub fn gen_len<T>(&self, len: T::Len) -> T
where
T: GenLen,
{
T::gen_len(self, len)
}
pub fn gen_range<T>(&self, r: impl RangeBounds<T::RangeBound>) -> T
where
T: GenRange,
{
T::gen_range(self, r)
}
#[deprecated(note = "use env.prng().gen_range(...)")]
pub fn u64_in_range(&self, r: impl RangeBounds<u64>) -> u64 {
self.gen_range(r)
}
pub fn shuffle<T>(&self, v: &mut T)
where
T: Shuffle,
{
v.shuffle(self);
}
}
impl<T> Shuffle for Vec<T> {
fn shuffle(&mut self, prng: &Prng) {
let env = prng.env();
assert_in_contract!(env);
let obj = internal::Env::prng_vec_shuffle(env, self.to_object()).unwrap_infallible();
*self = unsafe { Self::unchecked_new(env.clone(), obj) };
}
}
pub trait Fill {
fn fill(&mut self, prng: &Prng);
}
pub trait Gen {
fn gen(prng: &Prng) -> Self;
}
pub trait GenLen {
type Len;
fn gen_len(prng: &Prng, len: Self::Len) -> Self;
}
pub trait GenRange {
type RangeBound;
fn gen_range(prng: &Prng, r: impl RangeBounds<Self::RangeBound>) -> Self;
}
pub trait Shuffle {
fn shuffle(&mut self, prng: &Prng);
}
pub trait ToShuffled {
type Shuffled;
fn to_shuffled(&self, prng: &Prng) -> Self::Shuffled;
}
impl<T: Shuffle + Clone> ToShuffled for T {
type Shuffled = Self;
fn to_shuffled(&self, prng: &Prng) -> Self {
let mut copy = self.clone();
copy.shuffle(prng);
copy
}
}
impl Fill for u64 {
fn fill(&mut self, prng: &Prng) {
*self = Self::gen(prng);
}
}
impl Gen for u64 {
fn gen(prng: &Prng) -> Self {
let env = prng.env();
assert_in_contract!(env);
internal::Env::prng_u64_in_inclusive_range(env, u64::MIN, u64::MAX).unwrap_infallible()
}
}
impl GenRange for u64 {
type RangeBound = u64;
fn gen_range(prng: &Prng, r: impl RangeBounds<Self::RangeBound>) -> Self {
let env = prng.env();
assert_in_contract!(env);
let start_bound = match r.start_bound() {
Bound::Included(b) => *b,
Bound::Excluded(b) => *b + 1,
Bound::Unbounded => u64::MIN,
};
let end_bound = match r.end_bound() {
Bound::Included(b) => *b,
Bound::Excluded(b) => *b - 1,
Bound::Unbounded => u64::MAX,
};
internal::Env::prng_u64_in_inclusive_range(env, start_bound, end_bound).unwrap_infallible()
}
}
impl Fill for Bytes {
fn fill(&mut self, prng: &Prng) {
let env = prng.env();
assert_in_contract!(env);
let len: u32 = self.len();
let obj = internal::Env::prng_bytes_new(env, len.into()).unwrap_infallible();
*self = unsafe { Bytes::unchecked_new(env.clone(), obj) };
}
}
impl GenLen for Bytes {
type Len = u32;
fn gen_len(prng: &Prng, len: u32) -> Self {
let env = prng.env();
assert_in_contract!(env);
let obj = internal::Env::prng_bytes_new(env, len.into()).unwrap_infallible();
unsafe { Bytes::unchecked_new(env.clone(), obj) }
}
}
impl<const N: usize> Fill for BytesN<N> {
fn fill(&mut self, prng: &Prng) {
let bytesn = Self::gen(prng);
*self = bytesn;
}
}
impl<const N: usize> Gen for BytesN<N> {
fn gen(prng: &Prng) -> Self {
let env = prng.env();
assert_in_contract!(env);
let len: u32 = N.try_into().unwrap_optimized();
let obj = internal::Env::prng_bytes_new(env, len.into()).unwrap_infallible();
unsafe { BytesN::unchecked_new(env.clone(), obj) }
}
}
impl Fill for [u8] {
fn fill(&mut self, prng: &Prng) {
let env = prng.env();
assert_in_contract!(env);
let len: u32 = self.len().try_into().unwrap_optimized();
let bytes: Bytes = internal::Env::prng_bytes_new(env, len.into())
.unwrap_infallible()
.into_val(env);
bytes.copy_into_slice(self);
}
}
impl<const N: usize> Fill for [u8; N] {
fn fill(&mut self, prng: &Prng) {
let env = prng.env();
assert_in_contract!(env);
let len: u32 = N.try_into().unwrap_optimized();
let bytes: Bytes = internal::Env::prng_bytes_new(env, len.into())
.unwrap_infallible()
.into_val(env);
bytes.copy_into_slice(self);
}
}
impl<const N: usize> Gen for [u8; N] {
fn gen(prng: &Prng) -> Self {
let mut v = [0u8; N];
v.fill(prng);
v
}
}