From 20c42392171dad71fd030381c25d760855b8af7c Mon Sep 17 00:00:00 2001
From: Will Crichton <wcrichto@cs.stanford.edu>
Date: Wed, 10 Jul 2024 22:11:37 -0700
Subject: [PATCH] Make RNG seedable

---
 Cargo.lock                | 50 +++++++++++++++++++++++++++++++++++++++
 Cargo.toml                |  3 ++-
 crates/chatbot/Cargo.toml |  3 ++-
 crates/chatbot/src/lib.rs | 14 +++++++++--
 crates/server/Cargo.toml  |  4 ++++
 5 files changed, 70 insertions(+), 4 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 5bd5a0f..d44b29f 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -138,6 +138,7 @@ checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd"
 name = "chatbot"
 version = "0.1.0"
 dependencies = [
+ "rand",
  "tokio",
 ]
 
@@ -189,6 +190,17 @@ dependencies = [
  "pin-utils",
 ]
 
+[[package]]
+name = "getrandom"
+version = "0.2.15"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c4567c8db10ae91089c99af84c68c38da3ec2f087c3f82960bcdbf3656b6f4d7"
+dependencies = [
+ "cfg-if",
+ "libc",
+ "wasi",
+]
+
 [[package]]
 name = "gimli"
 version = "0.29.0"
@@ -433,6 +445,12 @@ version = "0.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184"
 
+[[package]]
+name = "ppv-lite86"
+version = "0.2.17"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5b40af805b3121feab8a3c29f04d8ad262fa8e0561883e7653e024ae4479e6de"
+
 [[package]]
 name = "proc-macro2"
 version = "1.0.86"
@@ -451,6 +469,36 @@ dependencies = [
  "proc-macro2",
 ]
 
+[[package]]
+name = "rand"
+version = "0.8.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404"
+dependencies = [
+ "libc",
+ "rand_chacha",
+ "rand_core",
+]
+
+[[package]]
+name = "rand_chacha"
+version = "0.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88"
+dependencies = [
+ "ppv-lite86",
+ "rand_core",
+]
+
+[[package]]
+name = "rand_core"
+version = "0.6.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c"
+dependencies = [
+ "getrandom",
+]
+
 [[package]]
 name = "redox_syscall"
 version = "0.5.2"
@@ -543,6 +591,8 @@ version = "0.1.0"
 dependencies = [
  "anyhow",
  "axum",
+ "chatbot",
+ "rand",
  "serde",
  "tokio",
 ]
diff --git a/Cargo.toml b/Cargo.toml
index 55ab10e..5dcb0ed 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -5,4 +5,5 @@ resolver = "2"
 [workspace.dependencies]
 tokio = { version = "1.38.0", features = ["full"] }
 axum = "0.7.5"
-serde = { version = "1.0.204", features = ["derive"] }
\ No newline at end of file
+serde = { version = "1.0.204", features = ["derive"] }
+rand = { version = "0.8.5", features = ["small_rng"] }
\ No newline at end of file
diff --git a/crates/chatbot/Cargo.toml b/crates/chatbot/Cargo.toml
index d3e30ed..3a6f742 100644
--- a/crates/chatbot/Cargo.toml
+++ b/crates/chatbot/Cargo.toml
@@ -4,4 +4,5 @@ version = "0.1.0"
 edition = "2021"
 
 [dependencies]
-tokio = { workspace = true }
\ No newline at end of file
+tokio = { workspace = true }
+rand = { workspace = true }
\ No newline at end of file
diff --git a/crates/chatbot/src/lib.rs b/crates/chatbot/src/lib.rs
index bc96bf0..7fb9d57 100644
--- a/crates/chatbot/src/lib.rs
+++ b/crates/chatbot/src/lib.rs
@@ -1,12 +1,22 @@
-use std::time::Duration;
+use rand::{rngs::SmallRng, Rng, SeedableRng};
+use std::{cell::RefCell, time::Duration};
 use tokio::time::sleep;
 
+thread_local! {
+    static RNG: RefCell<SmallRng> = RefCell::new(SmallRng::from_entropy());
+}
+
+/// Seeds the thread-local RNG used by [`gen_random_number`].
+pub fn seed_rng(seed: u64) {
+    RNG.with(|rng| *rng.borrow_mut() = SmallRng::seed_from_u64(seed));
+}
+
 /// Generates a random `usize`.
 ///
 /// Warning: may take a few seconds!
 pub async fn gen_random_number() -> usize {
     sleep(Duration::from_secs(2)).await;
-    4
+    RNG.with(|rng| rng.borrow_mut().gen())
 }
 
 /// Generates a list of possible responses given the current chat.
diff --git a/crates/server/Cargo.toml b/crates/server/Cargo.toml
index 9b30df7..aa5a7d0 100644
--- a/crates/server/Cargo.toml
+++ b/crates/server/Cargo.toml
@@ -8,3 +8,7 @@ anyhow = "1.0.86"
 axum = { workspace = true }
 serde = { workspace = true }
 tokio = { workspace = true }
+chatbot = { path = "../chatbot" }
+
+[dev-dependencies]
+rand = { workspace = true }
\ No newline at end of file