Skip to content

Commit dcc55b6

Browse files
newpavlovnewpavlov
committed
use index level
1 parent f314152 commit dcc55b6

File tree

4 files changed

+220
-140
lines changed

4 files changed

+220
-140
lines changed

‎rand_chacha/src/chacha.rs

Lines changed: 32 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -8,10 +8,11 @@
88

99
//! The ChaCha random number generator.
1010
11+
// TODO: remove
1112
#[cfg(not(feature = "std"))] use core;
1213
#[cfg(feature = "std")] use std as core;
1314

14-
use self::core::fmt;
15+
use self::core::{fmt, slice};
1516
use crate::guts::ChaCha;
1617
use rand_core::block::{BlockRng, BlockRngCore};
1718
use rand_core::{CryptoRng, Error, RngCore, SeedableRng};
@@ -25,38 +26,45 @@ const BUF_BLOCKS: u8 = 4;
2526
const BLOCK_WORDS: u8 = 16;
2627

2728
/// Type representing result of the ChaCha core iteration
28-
#[derive(Eq, PartialEq)]
29-
#[repr(align(4))]
30-
pub struct Results([u8; 256]);
29+
#[derive(Eq, PartialEq, Default, Clone, Copy, Debug)]
30+
pub struct Results([u64; 32]);
3131

32-
impl Default for Results {
33-
fn default() -> Self {
34-
Self([0; 256])
35-
}
36-
}
3732
impl AsRef<[u8]> for Results {
33+
#[inline(always)]
3834
fn as_ref(&self) -> &[u8] {
39-
&self.0
35+
unsafe {
36+
slice::from_raw_parts(
37+
self.0.as_ptr() as *const u8,
38+
8 * self.0.len(),
39+
)
40+
}
4041
}
4142
}
4243

43-
impl AsMut<[u8]> for Results {
44-
fn as_mut(&mut self) -> &mut [u8] {
45-
&mut self.0
44+
impl AsRef<[u32]> for Results {
45+
#[inline(always)]
46+
fn as_ref(&self) -> &[u32] {
47+
unsafe {
48+
slice::from_raw_parts(
49+
self.0.as_ptr() as *const u32,
50+
2 * self.0.len(),
51+
)
52+
}
4653
}
4754
}
48-
49-
impl Clone for Results {
50-
fn clone(&self) -> Self {
51-
let mut new = Self::default();
52-
new.0.copy_from_slice(&self.0);
53-
new
55+
impl AsRef<[u64]> for Results {
56+
#[inline(always)]
57+
fn as_ref(&self) -> &[u64] {
58+
&self.0
5459
}
5560
}
5661

57-
impl fmt::Debug for Results {
58-
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
59-
write!(f, "Results {{}}")
62+
impl AsMut<[u8; 256]> for Results {
63+
#[inline(always)]
64+
fn as_mut(&mut self) -> &mut [u8; 256] {
65+
unsafe {
66+
&mut *(self.0.as_mut_ptr() as *mut [u8; 256])
67+
}
6068
}
6169
}
6270

@@ -79,11 +87,8 @@ macro_rules! chacha_impl {
7987
type Results = Results;
8088
#[inline]
8189
fn generate(&mut self, r: &mut Self::Results) {
82-
// Fill slice of words by writing to equivalent slice of bytes, then fixing endianness.
83-
self.state.refill4($rounds, &mut r.0);
84-
for x in r.as_mut() {
85-
*x = x.to_le();
86-
}
90+
let r: &mut [u8; 256] = r.as_mut();
91+
self.state.refill4($rounds, r);
8792
}
8893
}
8994

‎rand_chacha/src/lib.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
#![deny(missing_docs)]
1717
#![deny(missing_debug_implementations)]
1818
#![doc(test(attr(allow(unused_variables), deny(warnings))))]
19-
#![cfg_attr(not(feature = "std"), no_std)]
19+
#![cfg_attr(not(feature = "std"), no_std)] // TODO: use `#![no_std]`
2020

2121
pub use rand_core;
2222

‎rand_core/src/block.rs

Lines changed: 118 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -27,8 +27,13 @@
2727
//!
2828
//! struct MyRngCore;
2929
//!
30+
//! #[derive(Eq, PartialEq, Default)]
31+
//! pub struct Results([u64; 8]);
32+
//!
33+
//! // implement `AsRef<[u8]>`, `AsRef<[u32]>`, and `AsRef<[u64]>` for `Results`
34+
//!
3035
//! impl BlockRngCore for MyRngCore {
31-
//! type Results = [u8; 32];
36+
//! type Results = Results;
3237
//!
3338
//! fn generate(&mut self, results: &mut Self::Results) {
3439
//! unimplemented!()
@@ -54,7 +59,7 @@
5459
//! [`fill_bytes`]: RngCore::fill_bytes
5560
5661
use crate::{CryptoRng, Error, RngCore, SeedableRng};
57-
use core::convert::{AsRef, TryInto};
62+
use core::convert::AsRef;
5863
use core::fmt;
5964
#[cfg(feature = "serde1")]
6065
use serde::{Deserialize, Serialize};
@@ -67,10 +72,49 @@ use serde::{Deserialize, Serialize};
6772
pub trait BlockRngCore {
6873
/// Results type. This is the 'block' an RNG implementing `BlockRngCore`
6974
/// generates, which will usually be an array like `[u8; 64]`.
70-
type Results: AsRef<[u8]> + AsMut<[u8]> + Default;
75+
type Results: AsRef<[u8]> + AsRef<[u32]> + AsRef<[u64]> + Default + Sized;
7176

7277
/// Generate a new block of results.
7378
fn generate(&mut self, results: &mut Self::Results);
79+
80+
/// Try to generate a new block of results.
81+
#[inline]
82+
fn try_generate(&mut self, results: &mut Self::Results) -> Result<(), Error> {
83+
self.generate(results);
84+
Ok(())
85+
}
86+
}
87+
88+
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
89+
enum IndexLevel {
90+
Bit,
91+
Byte,
92+
U32,
93+
U64,
94+
}
95+
96+
impl IndexLevel {
97+
#[inline(always)]
98+
fn convert(&mut self, index: usize, level: Self) -> usize {
99+
use IndexLevel::*;
100+
let res = match (*self, level) {
101+
(Bit, Bit) | (Byte, Byte) | (U32, U32) | (U64, U64) => return index,
102+
(Bit, Byte) => (index / 8) + ((index & 0b111) != 0) as usize,
103+
(Bit, U32) => (index / 32) + ((index & 0b1_1111) != 0) as usize,
104+
(Bit, U64) => (index / 64) + ((index & 0b11_1111) != 0) as usize,
105+
(Byte, Bit) => 8 * index,
106+
(Byte, U32) => (index / 4) + ((index & 0b11) != 0) as usize,
107+
(Byte, U64) => (index / 8) + ((index & 0b111) != 0) as usize,
108+
(U32, Bit) => 32 * index,
109+
(U32, Byte) => 4 * index,
110+
(U32, U64) => (index / 2) + ((index & 0b1) != 0) as usize,
111+
(U64, Bit) => 64 * index,
112+
(U64, Byte) => 8 * index,
113+
(U64, U32) => 2 * index,
114+
};
115+
*self = level;
116+
res
117+
}
74118
}
75119

76120
/// A wrapper type implementing [`RngCore`] for some type implementing
@@ -110,11 +154,12 @@ pub trait BlockRngCore {
110154
/// [`try_fill_bytes`]: RngCore::try_fill_bytes
111155
#[derive(Clone, Eq, PartialEq)]
112156
#[cfg_attr(feature = "serde1", derive(Serialize, Deserialize))]
113-
pub struct BlockRng<R: BlockRngCore + ?Sized> {
114-
results: R::Results,
115-
index: usize,
157+
pub struct BlockRng<R: BlockRngCore + Sized> {
116158
/// The *core* part of the RNG, implementing the `generate` function.
117159
pub core: R,
160+
results: R::Results,
161+
index: usize,
162+
level: IndexLevel,
118163
}
119164

120165
// Custom Debug implementation that does not expose the contents of `results`.
@@ -133,76 +178,96 @@ impl<R: BlockRngCore> BlockRng<R> {
133178
/// `BlockRngCore`. Results will be generated on first use.
134179
#[inline]
135180
pub fn new(core: R) -> BlockRng<R> {
136-
let results_empty = R::Results::default();
137-
BlockRng {
138-
core,
139-
index: 8 * results_empty.as_ref().len(),
140-
results: results_empty,
141-
}
181+
let results = R::Results::default();
182+
let index = AsRef::<[u8]>::as_ref(&results).len();
183+
BlockRng { core, results, index, level: IndexLevel::Byte }
184+
}
185+
186+
#[inline(always)]
187+
fn get_results_u8(&self) -> &[u8] {
188+
AsRef::<[u8]>::as_ref(&self.results)
189+
}
190+
191+
#[inline(always)]
192+
fn get_results_u32(&self) -> &[u32] {
193+
AsRef::<[u32]>::as_ref(&self.results)
194+
}
195+
196+
#[inline(always)]
197+
fn get_results_u64(&self) -> &[u64] {
198+
AsRef::<[u64]>::as_ref(&self.results)
142199
}
143200
}
144201

145202
impl<R: BlockRngCore> RngCore for BlockRng<R>
146203
where
147-
<R as BlockRngCore>::Results: AsRef<[u8]> + AsMut<[u8]>,
204+
<R as BlockRngCore>::Results: AsRef<[u8]> + AsRef<[u32]> + AsRef<[u64]> + Default + Sized,
148205
{
149206
#[inline]
150207
fn next_bool(&mut self) -> bool {
151-
let mut index = self.index;
152-
153-
if index / 8 >= self.results.as_ref().len() {
154-
self.core.generate(&mut self.results);
155-
index = 0;
208+
let index = self.level.convert(self.index, IndexLevel::Bit);
209+
match self.get_results_u8().get(index / 8) {
210+
Some(&v) => {
211+
self.index = index + 1;
212+
(v >> (index % 8) & 1) != 0
213+
}
214+
None => {
215+
self.core.generate(&mut self.results);
216+
self.index = 1;
217+
(self.get_results_u8()[0] & 1) != 0
218+
}
156219
}
157-
158-
let res = (self.results.as_ref()[index / 8] >> (index % 8)) & 0b1 != 0;
159-
self.index = index + 1;
160-
res
161220
}
162221

163222
#[inline]
164223
fn next_u32(&mut self) -> u32 {
165-
let mut index = self.index;
166-
index = 4 * ((index / 32) + ((index & 0b1_1111) != 0) as usize);
167-
168-
if index + 4 > self.results.as_ref().len() {
169-
self.core.generate(&mut self.results);
170-
index = 0;
224+
let index = self.level.convert(self.index, IndexLevel::U32);
225+
match self.get_results_u32().get(index) {
226+
Some(&v) => {
227+
self.index = index + 1;
228+
v
229+
}
230+
None => {
231+
self.core.generate(&mut self.results);
232+
self.index = 1;
233+
self.get_results_u32()[0]
234+
}
171235
}
172-
173-
let buf = self.results.as_ref()[index..index + 4].try_into().unwrap();
174-
self.index = 8 * (index + 4);
175-
u32::from_le_bytes(buf)
176236
}
177237

178238
#[inline]
179239
fn next_u64(&mut self) -> u64 {
180-
let mut index = self.index;
181-
index = 8 * ((index / 64) + ((index & 0b11_1111) != 0) as usize);
182-
183-
if index + 8 > self.results.as_ref().len() {
184-
self.core.generate(&mut self.results);
185-
index = 0;
240+
let index = self.level.convert(self.index, IndexLevel::U64);
241+
match self.get_results_u64().get(index) {
242+
Some(&v) => {
243+
self.index = index + 1;
244+
v
245+
}
246+
None => {
247+
self.core.generate(&mut self.results);
248+
self.index = 1;
249+
self.get_results_u64()[0]
250+
}
186251
}
252+
}
187253

188-
let buf = self.results.as_ref()[index..index + 8].try_into().unwrap();
189-
self.index = 8 * (index + 8);
190-
u64::from_le_bytes(buf)
254+
#[inline]
255+
fn fill_bytes(&mut self, dest: &mut [u8]) {
256+
self.try_fill_bytes(dest).unwrap();
191257
}
192258

193259
#[inline]
194-
fn fill_bytes(&mut self, mut dest: &mut [u8]) {
195-
let mut index = self.index;
196-
index = (index / 8) + ((index & 0b111) != 0) as usize;
260+
fn try_fill_bytes(&mut self, mut dest: &mut [u8]) -> Result<(), Error> {
261+
let index = self.level.convert(self.index, IndexLevel::Byte);
197262

198-
let rlen = self.results.as_ref().len();
263+
let rlen = self.get_results_u8().len();
199264
if index < rlen {
200265
let dlen = dest.len();
201-
let res = self.results.as_ref();
266+
let res = self.get_results_u8();
202267
if dlen <= rlen - index {
203268
dest.copy_from_slice(&res[index..index + dlen]);
204-
self.index = 8*(index + dlen);
205-
return;
269+
self.index = index + dlen;
270+
return Ok(());
206271
} else {
207272
let (l, r) = dest.split_at_mut(rlen - index);
208273
l.copy_from_slice(&res[index..]);
@@ -211,26 +276,20 @@ where
211276
}
212277

213278
let mut chunks = dest.chunks_exact_mut(rlen);
214-
215279
for chunk in &mut chunks {
216280
let mut buf = R::Results::default();
217-
self.core.generate(&mut buf);
281+
self.core.try_generate(&mut buf)?;
218282
chunk.copy_from_slice(buf.as_ref());
219283
}
220284

221285
let rem = chunks.into_remainder();
222286
if !rem.is_empty() {
223-
self.core.generate(&mut self.results);
224-
rem.copy_from_slice(&self.results.as_ref()[..rem.len()]);
225-
self.index = 8 * rem.len();
287+
self.core.try_generate(&mut self.results)?;
288+
rem.copy_from_slice(&self.get_results_u8()[..rem.len()]);
289+
self.index = rem.len();
226290
} else {
227-
self.index = 8 * rlen;
291+
self.index = rlen;
228292
}
229-
}
230-
231-
#[inline(always)]
232-
fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), Error> {
233-
self.fill_bytes(dest);
234293
Ok(())
235294
}
236295
}

0 commit comments

Comments
 (0)