ep_core/
random_number_generator.rs

1// This file is part of Substrate.
2
3// Copyright (C) 2017-2021 Parity Technologies (UK) Ltd.
4// SPDX-License-Identifier: Apache-2.0
5
6// Licensed under the Apache License, Version 2.0 (the "License");
7// you may not use this file except in compliance with the License.
8// You may obtain a copy of the License at
9//
10// 	http://www.apache.org/licenses/LICENSE-2.0
11//
12// Unless required by applicable law or agreed to in writing, software
13// distributed under the License is distributed on an "AS IS" BASIS,
14// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15// See the License for the specific language governing permissions and
16// limitations under the License.
17
18//! A simple pseudo random number generator that allows a stream of random numbers to be efficiently
19//! created from a single initial seed hash.
20
21use parity_scale_codec::{Decode, DecodeWithMemTracking, Encode};
22use scale_info::TypeInfo;
23use sp_runtime::traits::{Hash, TrailingZeroInput};
24
25/// Pseudo-random number streamer. This retains the state of the random number stream. It's as
26/// secure as the combination of the seed with which it is constructed and the hash function it uses
27/// to cycle elements.
28///
29/// It can be saved and later reloaded using the Codec traits.
30///
31/// (It is recommended to use the `rand_chacha` crate as an alternative to this where possible.)
32///
33/// Example:
34/// ```
35/// use sp_runtime::traits::{Hash, BlakeTwo256};
36/// use ep_core::RandomNumberGenerator;
37/// let random_seed = BlakeTwo256::hash(b"Sixty-nine");
38/// let mut rng = <RandomNumberGenerator<BlakeTwo256>>::new(random_seed);
39/// assert_eq!(rng.pick_u32(100), 59);
40/// assert_eq!(rng.pick_item(&[1, 2, 3]), Some(&1));
41/// ```
42///
43/// This can use any cryptographic `Hash` function as the means of entropy-extension, and avoids
44/// needless extensions of entropy.
45///
46/// If you're persisting it over blocks, be aware that the sequence will start to repeat. This won't
47/// be a practical issue unless you're using tiny hash types (e.g. 64-bit) and pulling hundred of
48/// megabytes of data from it.
49#[derive(Encode, Decode, DecodeWithMemTracking, TypeInfo)]
50pub struct RandomNumberGenerator<Hashing: Hash> {
51	current: Hashing::Output,
52	offset: u32,
53}
54
55impl<Hashing: Hash> RandomNumberGenerator<Hashing> {
56	/// A new source of random data.
57	pub fn new(seed: Hashing::Output) -> Self {
58		Self { current: seed, offset: 0 }
59	}
60
61	fn offset(&self) -> usize {
62		self.offset as usize
63	}
64
65	/// Returns a number at least zero, at most `max`.
66	pub fn pick_u32(&mut self, max: u32) -> u32 {
67		let needed = (4 - max.leading_zeros() / 8) as usize;
68		let top = ((1 << (needed as u64 * 8)) / (max as u64 + 1) * (max as u64 + 1) - 1) as u32;
69		loop {
70			if self.offset() + needed > self.current.as_ref().len() {
71				// rehash
72				self.current = <Hashing as Hash>::hash(self.current.as_ref());
73				self.offset = 0;
74			}
75			let data = &self.current.as_ref()[self.offset()..self.offset() + needed];
76			self.offset += needed as u32;
77			let raw = u32::decode(&mut TrailingZeroInput::new(data)).unwrap_or(0);
78			if raw <= top {
79				break if max < u32::MAX { raw % (max + 1) } else { raw };
80			}
81		}
82	}
83
84	/// Returns a number at least 1, at most `max`.
85	pub fn pick_non_zero_u32(&mut self, max: u32) -> u32 {
86		self.pick_u32(max - 1) + 1
87	}
88
89	/// Returns a number at least zero, at most `max`.
90	///
91	/// This returns a `usize`, but internally it only uses `u32` so avoid consensus problems.
92	pub fn pick_usize(&mut self, max: usize) -> usize {
93		self.pick_u32(max as u32) as usize
94	}
95
96	/// Pick a random element from an array of `items`.
97	///
98	/// This is guaranteed to return `Some` except in the case that the given array `items` is
99	/// empty.
100	pub fn pick_item<'a, T>(&mut self, items: &'a [T]) -> Option<&'a T> {
101		if items.is_empty() {
102			None
103		} else {
104			Some(&items[self.pick_usize(items.len() - 1)])
105		}
106	}
107}
108
109#[cfg(test)]
110mod tests {
111	use super::RandomNumberGenerator;
112	use sp_runtime::traits::{BlakeTwo256, Hash};
113
114	#[test]
115	fn does_not_panic_on_max() {
116		let seed = BlakeTwo256::hash(b"Fourty-two");
117		let _random = RandomNumberGenerator::<BlakeTwo256>::new(seed).pick_u32(u32::MAX);
118	}
119}
ep_core/random_number_generator.rs

ep_core/
random_number_generator.rs
