add a seeded random number generator so that sequences of random numbers can be easily reproduced (for rust-lang#2379)

Author: Dretch

src/libcore/rand.rs

@@ -1,12 +1,14 @@
 #[doc = "Random number generation"];
 
-export rng, weighted, extensions;
+export rng, seed, seeded_rng, weighted, extensions;
 
 enum rctx {}
 
 #[abi = "cdecl"]
 native mod rustrt {
+    fn rand_seed() -> [u8];
     fn rand_new() -> *rctx;
+    fn rand_new_seeded(seed: [u8]) -> *rctx;
     fn rand_next(c: *rctx) -> u32;
     fn rand_free(c: *rctx);
 }
@@ -227,20 +229,50 @@ impl extensions for rng {
 
 }
 
-#[doc = "Create a random number generator"]
-fn rng() -> rng {
-    resource rand_res(c: *rctx) { rustrt::rand_free(c); }
+resource rand_res(c: *rctx) { rustrt::rand_free(c); }
 
-    impl of rng for @rand_res {
-        fn next() -> u32 { ret rustrt::rand_next(**self); }
-    }
+impl of rng for @rand_res {
+    fn next() -> u32 { ret rustrt::rand_next(**self); }
+}
 
+#[doc = "Create a new random seed for seeded_rng"]
+fn seed() -> [u8] {
+    rustrt::rand_seed()
+}
+
+#[doc = "Create a random number generator with a system specified seed"]
+fn rng() -> rng {
     @rand_res(rustrt::rand_new()) as rng
 }
 
+#[doc = "Create a random number generator using the specified seed. A \
+         generator constructed with a given seed will generate the same \
+         sequence of values as all other generators constructed with the \
+         same seed. The seed may be any length."]
+fn seeded_rng(seed: [u8]) -> rng {
+    @rand_res(rustrt::rand_new_seeded(seed)) as rng
+}
+
 #[cfg(test)]
 mod tests {
 
+    #[test]
+    fn rng_seeded() {
+        let seed = rand::seed();
+        let ra = rand::seeded_rng(seed);
+        let rb = rand::seeded_rng(seed);
+        assert ra.gen_str(100u) == rb.gen_str(100u);
+    }
+
+    #[test]
+    fn rng_seeded_custom_seed() {
+        // much shorter than generated seeds which are 1024 bytes
+        let seed = [2u8, 32u8, 4u8, 32u8, 51u8];
+        let ra = rand::seeded_rng(seed);
+        let rb = rand::seeded_rng(seed);
+        assert ra.gen_str(100u) == rb.gen_str(100u);
+    }
+
     #[test]
     fn gen_int_from() {
         let r = rand::rng();

src/rt/rust_builtin.cpp

@@ -178,16 +178,41 @@ rust_str_push(rust_vec** sp, uint8_t byte) {
     (*sp)->fill = fill + 1;
 }
 
+extern "C" CDECL rust_vec*
+rand_seed() {
+    size_t size = sizeof(ub4) * RANDSIZ;
+    rust_task *task = rust_get_current_task();
+    rust_vec *v = (rust_vec *) task->kernel->malloc(vec_size<uint8_t>(size),
+                                            "rand_seed");
+    v->fill = v->alloc = size;
+    isaac_seed((uint8_t*) &v->data);
+    return v;
+}
+
 extern "C" CDECL void *
 rand_new() {
     rust_task *task = rust_get_current_task();
     rust_sched_loop *thread = task->sched_loop;
-    randctx *rctx = (randctx *) task->malloc(sizeof(randctx), "randctx");
+    randctx *rctx = (randctx *) task->malloc(sizeof(randctx), "rand_new");
+    if (!rctx) {
+        task->fail();
+        return NULL;
+    }
+    isaac_init(thread->kernel, rctx, NULL);
+    return rctx;
+}
+
+extern "C" CDECL void *
+rand_new_seeded(rust_vec* seed) {
+    rust_task *task = rust_get_current_task();
+    rust_sched_loop *thread = task->sched_loop;
+    randctx *rctx = (randctx *) task->malloc(sizeof(randctx),
+                                             "rand_new_seeded");
     if (!rctx) {
         task->fail();
         return NULL;
     }
-    isaac_init(thread->kernel, rctx);
+    isaac_init(thread->kernel, rctx, seed);
     return rctx;
 }
 

src/rt/rust_sched_loop.cpp

@@ -29,7 +29,7 @@ rust_sched_loop::rust_sched_loop(rust_scheduler *sched,int id) :
     name("main")
 {
     LOGPTR(this, "new dom", (uintptr_t)this);
-    isaac_init(kernel, &rctx);
+    isaac_init(kernel, &rctx, NULL);
 
     if (!tls_initialized)
         init_tls();

src/rt/rust_util.h

@@ -32,46 +32,6 @@ align_to(T size, size_t alignment) {
     return x;
 }
 
-// Initialization helper for ISAAC RNG
-
-inline void
-isaac_init(rust_kernel *kernel, randctx *rctx)
-{
-        memset(rctx, 0, sizeof(randctx));
-
-        char *rust_seed = kernel->env->rust_seed;
-        if (rust_seed != NULL) {
-            ub4 seed = (ub4) atoi(rust_seed);
-            for (size_t i = 0; i < RANDSIZ; i ++) {
-                memcpy(&rctx->randrsl[i], &seed, sizeof(ub4));
-                seed = (seed + 0x7ed55d16) + (seed << 12);
-            }
-        } else {
-#ifdef __WIN32__
-            HCRYPTPROV hProv;
-            kernel->win32_require
-                (_T("CryptAcquireContext"),
-                 CryptAcquireContext(&hProv, NULL, NULL, PROV_RSA_FULL,
-                                     CRYPT_VERIFYCONTEXT|CRYPT_SILENT));
-            kernel->win32_require
-                (_T("CryptGenRandom"),
-                 CryptGenRandom(hProv, sizeof(rctx->randrsl),
-                                (BYTE*)(&rctx->randrsl)));
-            kernel->win32_require
-                (_T("CryptReleaseContext"),
-                 CryptReleaseContext(hProv, 0));
-#else
-            int fd = open("/dev/urandom", O_RDONLY);
-            assert(fd > 0);
-            assert(read(fd, (void*) &rctx->randrsl, sizeof(rctx->randrsl))
-                   == sizeof(rctx->randrsl));
-            assert(close(fd) == 0);
-#endif
-        }
-
-        randinit(rctx, 1);
-}
-
 // Interior vectors (rust-user-code level).
 
 struct
@@ -136,6 +96,53 @@ make_str_vec(rust_kernel* kernel, size_t nstrs, char **strs) {
     return v;
 }
 
+// Initialization helpers for ISAAC RNG
+
+inline void isaac_seed(uint8_t* dest)
+{
+    size_t size = sizeof(ub4) * RANDSIZ;
+#ifdef __WIN32__
+    HCRYPTPROV hProv;
+    kernel->win32_require
+        (_T("CryptAcquireContext"),
+         CryptAcquireContext(&hProv, NULL, NULL, PROV_RSA_FULL,
+                             CRYPT_VERIFYCONTEXT|CRYPT_SILENT));
+    kernel->win32_require
+        (_T("CryptGenRandom"), CryptGenRandom(hProv, size, (BYTE*) dest));
+    kernel->win32_require
+        (_T("CryptReleaseContext"), CryptReleaseContext(hProv, 0));
+#else
+    int fd = open("/dev/urandom", O_RDONLY);
+    assert(fd > 0);
+    assert(read(fd, dest, size) == (int) size);
+    assert(close(fd) == 0);
+#endif
+}
+
+inline void
+isaac_init(rust_kernel *kernel, randctx *rctx, rust_vec* user_seed)
+{
+    memset(rctx, 0, sizeof(randctx));
+
+    char *env_seed = kernel->env->rust_seed;
+    if (user_seed != NULL) {
+        // ignore bytes after the required length
+        size_t seed_len = user_seed->fill < sizeof(rctx->randrsl)
+            ? user_seed->fill : sizeof(rctx->randrsl);
+        memcpy(&rctx->randrsl, user_seed->data, seed_len);
+    } else if (env_seed != NULL) {
+        ub4 seed = (ub4) atoi(env_seed);
+        for (size_t i = 0; i < RANDSIZ; i ++) {
+            memcpy(&rctx->randrsl[i], &seed, sizeof(ub4));
+            seed = (seed + 0x7ed55d16) + (seed << 12);
+        }
+    } else {
+        isaac_seed((uint8_t*) &rctx->randrsl);
+    }
+
+    randinit(rctx, 1);
+}
+
 //
 // Local Variables:
 // mode: C++

src/rt/rustrt.def.in

@@ -26,7 +26,9 @@ rust_port_id_send
 rust_port_select
 rand_free
 rand_new
+rand_new_seeded
 rand_next
+rand_seed
 refcount
 rust_get_sched_id
 rust_new_sched