nanorand/
buffer.rs

Help
1use crate::rand::{Rng, SeedableRng};
2use alloc::vec::Vec;
3use core::default::Default;
4
5/// A buffered wrapper for any [Rng] implementation.
6/// It will keep unused bytes from the last call to [`Rng::rand`], and use them
7/// for subsequent randomness if needed, rather than throwing them away.
8///
9/// ```rust
10/// use nanorand::{Rng, BufferedRng, WyRand};
11///
12/// let mut thingy = [0u8; 5];
13/// let mut rng = BufferedRng::new(WyRand::new());
14/// rng.fill(&mut thingy);
15/// // As WyRand generates 8 bytes of output, and our target is only 5 bytes,
16/// // 3 bytes will remain in the buffer.
17/// assert_eq!(rng.buffered(), 3);
18/// ```
19#[derive(Clone)]
20pub struct BufferedRng<InternalGenerator: Rng<OUTPUT>, const OUTPUT: usize> {
21	rng: InternalGenerator,
22	buffer: Vec<u8>,
23}
24
25impl<InternalGenerator: Rng<OUTPUT>, const OUTPUT: usize> BufferedRng<InternalGenerator, OUTPUT> {
26	/// Wraps a [`Rng`] InternalGenerator in a [`BufferedRng`] instance.
27	pub fn new(rng: InternalGenerator) -> Self {
28		Self {
29			rng,
30			buffer: Vec::new(),
31		}
32	}
33
34	/// Returns the internal RNG, dropping the buffer.
35	pub fn into_inner(self) -> InternalGenerator {
36		self.rng
37	}
38
39	/// Returns how many unused bytes are currently buffered.
40	pub fn buffered(&self) -> usize {
41		self.buffer.len()
42	}
43}
44
45impl<InternalGenerator: Rng<OUTPUT>, const OUTPUT: usize> Rng<OUTPUT>
46	for BufferedRng<InternalGenerator, OUTPUT>
47{
48	fn rand(&mut self) -> [u8; OUTPUT] {
49		let mut out = [0_u8; OUTPUT];
50		self.fill_bytes(&mut out);
51		out
52	}
53
54	fn fill_bytes<Bytes>(&mut self, mut output: Bytes)
55	where
56		Bytes: AsMut<[u8]>,
57	{
58		let output = output.as_mut();
59		let mut remaining = output.len();
60		while remaining > 0 {
61			if self.buffer.is_empty() {
62				self.buffer.extend_from_slice(&self.rng.rand());
63			}
64			let to_copy = core::cmp::min(remaining, self.buffer.len());
65			let output_len = output.len();
66			let start_idx = output_len - remaining;
67			output[start_idx..start_idx + to_copy].copy_from_slice(&self.buffer[..to_copy]);
68			self.buffer.drain(..to_copy);
69			remaining = remaining.saturating_sub(to_copy);
70		}
71	}
72}
73
74#[cfg(feature = "std")]
75impl<InternalGenerator: Rng<OUTPUT>, const OUTPUT: usize> std::io::Read
76	for BufferedRng<InternalGenerator, OUTPUT>
77{
78	fn read(&mut self, output: &mut [u8]) -> std::io::Result<usize> {
79		self.fill_bytes(&mut *output);
80		Ok(output.len())
81	}
82
83	fn read_to_end(&mut self, buf: &mut Vec<u8>) -> std::io::Result<usize> {
84		buf.extend_from_slice(&self.buffer);
85		Ok(self.buffer.drain(..).count())
86	}
87
88	fn read_to_string(&mut self, _buf: &mut String) -> std::io::Result<usize> {
89		panic!("attempted to read an rng into a string")
90	}
91}
92
93impl<
94		InternalGenerator: SeedableRng<SEED_SIZE, OUTPUT>,
95		const OUTPUT: usize,
96		const SEED_SIZE: usize,
97	> SeedableRng<SEED_SIZE, OUTPUT> for BufferedRng<InternalGenerator, OUTPUT>
98{
99	fn reseed(&mut self, seed: [u8; SEED_SIZE]) {
100		self.rng.reseed(seed);
101	}
102}
103
104impl<InternalGenerator: Rng<OUTPUT> + Default, const OUTPUT: usize> Default
105	for BufferedRng<InternalGenerator, OUTPUT>
106{
107	fn default() -> Self {
108		Self::new(InternalGenerator::default())
109	}
110}
nanorand/buffer.rs

nanorand/
buffer.rs
