nanorand/
buffer.rs

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 const 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	#[allow(clippy::missing_const_for_fn)]
36	pub fn into_inner(self) -> InternalGenerator {
37		self.rng
38	}
39
40	/// Returns how many unused bytes are currently buffered.
41	pub fn buffered(&self) -> usize {
42		self.buffer.len()
43	}
44}
45
46impl<InternalGenerator: Rng<OUTPUT>, const OUTPUT: usize> Rng<OUTPUT>
47	for BufferedRng<InternalGenerator, OUTPUT>
48{
49	fn rand(&mut self) -> [u8; OUTPUT] {
50		let mut out = [0_u8; OUTPUT];
51		self.fill_bytes(&mut out);
52		out
53	}
54
55	fn fill_bytes<Bytes>(&mut self, mut output: Bytes)
56	where
57		Bytes: AsMut<[u8]>,
58	{
59		let output = output.as_mut();
60		let mut remaining = output.len();
61		while remaining > 0 {
62			if self.buffer.is_empty() {
63				self.buffer.extend_from_slice(&self.rng.rand());
64			}
65			let to_copy = core::cmp::min(remaining, self.buffer.len());
66			let output_len = output.len();
67			let start_idx = output_len - remaining;
68			output[start_idx..start_idx + to_copy].copy_from_slice(&self.buffer[..to_copy]);
69			self.buffer.drain(..to_copy);
70			remaining = remaining.saturating_sub(to_copy);
71		}
72	}
73}
74
75#[cfg(feature = "std")]
76impl<InternalGenerator: Rng<OUTPUT>, const OUTPUT: usize> std::io::Read
77	for BufferedRng<InternalGenerator, OUTPUT>
78{
79	fn read(&mut self, output: &mut [u8]) -> std::io::Result<usize> {
80		self.fill_bytes(&mut *output);
81		Ok(output.len())
82	}
83
84	fn read_to_end(&mut self, buf: &mut Vec<u8>) -> std::io::Result<usize> {
85		buf.extend_from_slice(&self.buffer);
86		Ok(self.buffer.drain(..).count())
87	}
88
89	fn read_to_string(&mut self, _buf: &mut String) -> std::io::Result<usize> {
90		panic!("attempted to read an rng into a string")
91	}
92}
93
94impl<
95		InternalGenerator: SeedableRng<SEED_SIZE, OUTPUT>,
96		const OUTPUT: usize,
97		const SEED_SIZE: usize,
98	> SeedableRng<SEED_SIZE, OUTPUT> for BufferedRng<InternalGenerator, OUTPUT>
99{
100	fn reseed(&mut self, seed: [u8; SEED_SIZE]) {
101		self.rng.reseed(seed);
102	}
103}
104
105impl<InternalGenerator: Rng<OUTPUT> + Default, const OUTPUT: usize> Default
106	for BufferedRng<InternalGenerator, OUTPUT>
107{
108	fn default() -> Self {
109		Self::new(InternalGenerator::default())
110	}
111}
nanorand/buffer.rs

nanorand/
buffer.rs
