tp_runtime/
random_number_generator.rs

1// This file is part of Tetcore.
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 codec::{Encode, Decode};
22use crate::traits::{Hash, TrailingZeroInput};
23
24/// Pseudo-random number streamer. This retains the state of the random number stream. It's as
25/// secure as the combination of the seed with which it is constructed and the hash function it uses
26/// to cycle elements.
27///
28/// It can be saved and later reloaded using the Codec traits.
29///
30/// Example:
31/// ```
32/// use tp_runtime::traits::{Hash, BlakeTwo256};
33/// use tp_runtime::RandomNumberGenerator;
34/// let random_seed = BlakeTwo256::hash(b"Sixty-nine");
35/// let mut rng = <RandomNumberGenerator<BlakeTwo256>>::new(random_seed);
36/// assert_eq!(rng.pick_u32(100), 59);
37/// assert_eq!(rng.pick_item(&[1, 2, 3]), Some(&1));
38/// ```
39///
40/// This can use any cryptographic `Hash` function as the means of entropy-extension, and avoids
41/// needless extensions of entropy.
42///
43/// If you're persisting it over blocks, be aware that the sequence will start to repeat. This won't
44/// be a practical issue unless you're using tiny hash types (e.g. 64-bit) and pulling hundred of
45/// megabytes of data from it.
46#[derive(Encode, Decode)]
47pub struct RandomNumberGenerator<Hashing: Hash> {
48	current: Hashing::Output,
49	offset: u32,
50}
51
52impl<Hashing: Hash> RandomNumberGenerator<Hashing> {
53	/// A new source of random data.
54	pub fn new(seed: Hashing::Output) -> Self {
55		Self {
56			current: seed,
57			offset: 0,
58		}
59	}
60
61	fn offset(&self) -> usize { self.offset as usize }
62
63	/// Returns a number at least zero, at most `max`.
64	pub fn pick_u32(&mut self, max: u32) -> u32 {
65		let needed = (4 - max.leading_zeros() / 8) as usize;
66		let top = ((1 << (needed as u64 * 8)) / ((max + 1) as u64) * ((max + 1) as u64) - 1) as u32;
67		loop {
68			if self.offset() + needed > self.current.as_ref().len() {
69				// rehash
70				self.current = <Hashing as Hash>::hash(self.current.as_ref());
71				self.offset = 0;
72			}
73			let data = &self.current.as_ref()[self.offset()..self.offset() + needed];
74			self.offset += needed as u32;
75			let raw = u32::decode(&mut TrailingZeroInput::new(data)).unwrap_or(0);
76			if raw <= top {
77				break if max < u32::max_value() {
78					raw % (max + 1)
79				} else {
80					raw
81				}
82			}
83		}
84	}
85
86	/// Returns a number at least zero, at most `max`.
87	///
88	/// This returns a `usize`, but internally it only uses `u32` so avoid consensus problems.
89	pub fn pick_usize(&mut self, max: usize) -> usize {
90		self.pick_u32(max as u32) as usize
91	}
92
93	/// Pick a random element from an array of `items`.
94	///
95	/// This is guaranteed to return `Some` except in the case that the given array `items` is
96	/// empty.
97	pub fn pick_item<'a, T>(&mut self, items: &'a [T]) -> Option<&'a T> {
98		if items.is_empty() {
99			None
100		} else {
101			Some(&items[self.pick_usize(items.len() - 1)])
102		}
103	}
104}