Skip to content

Commit 14a6a62

Browse files
committed
std::rand: simplify/safe-ify the default Rng.fill_bytes.
The `&[u8]` -> `&[u64]` and `&[u32]` casts were not nice: they ignored alignment requirements and are generally very unsafe.
1 parent ae0905a commit 14a6a62

File tree

1 file changed

+44
-45
lines changed

1 file changed

+44
-45
lines changed

src/libstd/rand/mod.rs

Lines changed: 44 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -52,8 +52,6 @@ fn main () {
5252
```
5353
*/
5454

55-
use mem::size_of;
56-
use unstable::raw::Slice;
5755
use cast;
5856
use container::Container;
5957
use iter::{Iterator, range};
@@ -136,46 +134,26 @@ pub trait Rng {
136134
/// }
137135
/// ```
138136
fn fill_bytes(&mut self, dest: &mut [u8]) {
139-
let mut slice: Slice<u64> = unsafe { cast::transmute_copy(&dest) };
140-
slice.len /= size_of::<u64>();
141-
let as_u64: &mut [u64] = unsafe { cast::transmute(slice) };
142-
for dest in as_u64.mut_iter() {
143-
*dest = self.next_u64();
144-
}
145-
146-
// the above will have filled up the vector as much as
147-
// possible in multiples of 8 bytes.
148-
let mut remaining = dest.len() % 8;
149-
150-
// space for a u32
151-
if remaining >= 4 {
152-
let mut slice: Slice<u32> = unsafe { cast::transmute_copy(&dest) };
153-
slice.len /= size_of::<u32>();
154-
let as_u32: &mut [u32] = unsafe { cast::transmute(slice) };
155-
as_u32[as_u32.len() - 1] = self.next_u32();
156-
remaining -= 4;
157-
}
158-
// exactly filled
159-
if remaining == 0 { return }
160-
161-
// now we know we've either got 1, 2 or 3 spots to go,
162-
// i.e. exactly one u32 is enough.
163-
let rand = self.next_u32();
164-
let remaining_index = dest.len() - remaining;
165-
match dest.mut_slice_from(remaining_index) {
166-
[ref mut a] => {
167-
*a = rand as u8;
137+
// this could, in theory, be done by transmuting dest to a
138+
// [u64], but this is (1) likely to be undefined behaviour for
139+
// LLVM, (2) has to be very careful about alignment concerns,
140+
// (3) adds more `unsafe` that needs to be checked, (4)
141+
// probably doesn't give much performance gain if
142+
// optimisations are on.
143+
let mut count = 0;
144+
let mut num = 0;
145+
for byte in dest.mut_iter() {
146+
if count == 0 {
147+
// we could micro-optimise here by generating a u32 if
148+
// we only need a few more bytes to fill the vector
149+
// (i.e. at most 4).
150+
num = self.next_u64();
151+
count = 8;
168152
}
169-
[ref mut a, ref mut b] => {
170-
*a = rand as u8;
171-
*b = (rand >> 8) as u8;
172-
}
173-
[ref mut a, ref mut b, ref mut c] => {
174-
*a = rand as u8;
175-
*b = (rand >> 8) as u8;
176-
*c = (rand >> 16) as u8;
177-
}
178-
_ => fail!("Rng.fill_bytes: the impossible occurred: remaining != 1, 2 or 3")
153+
154+
*byte = (num & 0xff) as u8;
155+
num >>= 8;
156+
count -= 1;
179157
}
180158
}
181159

@@ -749,14 +727,35 @@ pub fn random<T: Rand>() -> T {
749727
mod test {
750728
use iter::{Iterator, range};
751729
use option::{Option, Some};
730+
use vec;
752731
use super::*;
753732

733+
struct ConstRng { i: u64 }
734+
impl Rng for ConstRng {
735+
fn next_u32(&mut self) -> u32 { self.i as u32 }
736+
fn next_u64(&mut self) -> u64 { self.i }
737+
738+
// no fill_bytes on purpose
739+
}
740+
754741
#[test]
755742
fn test_fill_bytes_default() {
756-
let mut r = weak_rng();
757-
758-
let mut v = [0u8, .. 100];
759-
r.fill_bytes(v);
743+
let mut r = ConstRng { i: 0x11_22_33_44_55_66_77_88 };
744+
745+
// check every remainder mod 8, both in small and big vectors.
746+
let lengths = [0, 1, 2, 3, 4, 5, 6, 7,
747+
80, 81, 82, 83, 84, 85, 86, 87];
748+
for &n in lengths.iter() {
749+
let mut v = vec::from_elem(n, 0u8);
750+
r.fill_bytes(v);
751+
752+
// use this to get nicer error messages.
753+
for (i, &byte) in v.iter().enumerate() {
754+
if byte == 0 {
755+
fail!("byte {} of {} is zero", i, n)
756+
}
757+
}
758+
}
760759
}
761760

762761
#[test]

0 commit comments

Comments
 (0)