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}