From 79aaf117ea20de16e3e2fd64ac861155883f7f64 Mon Sep 17 00:00:00 2001
From: Ilya Grigoriev <ilyagr@users.noreply.github.com>
Date: Thu, 29 Dec 2022 19:59:01 -0800
Subject: [PATCH] Put an RNG in user settings with configurable seed

Some hijinx are utilized to make it possible to generate
random values using an immutable `&UserSettings` reference.
---
 Cargo.lock          |  1 +
 lib/Cargo.toml      |  3 ++-
 lib/src/settings.rs | 65 ++++++++++++++++++++++++++++++++++++++++++---
 src/config.rs       |  5 ++++
 4 files changed, 70 insertions(+), 4 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index fa8b8a026..328a1b70d 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -846,6 +846,7 @@ dependencies = [
  "pest_derive",
  "prost",
  "rand",
+ "rand_chacha",
  "regex",
  "serde_json",
  "tempfile",
diff --git a/lib/Cargo.toml b/lib/Cargo.toml
index 7b48b0d94..3ff4fd801 100644
--- a/lib/Cargo.toml
+++ b/lib/Cargo.toml
@@ -31,6 +31,8 @@ maplit = "1.0.2"
 once_cell = "1.17.0"
 pest = "2.5.2"
 pest_derive = "2.5.2"
+rand = "0.8.5"
+rand_chacha = "0.3.1"
 regex = "1.7.0"
 serde_json = "1.0.91"
 tempfile = "3.3.0"
@@ -46,7 +48,6 @@ prost = "0.11.5"
 assert_matches = "1.5.0"
 insta = "1.23.0"
 num_cpus = "1.15.0"
-rand = "0.8.5"
 test-case = "2.2.2"
 testutils = { path = "testutils" }
 
diff --git a/lib/src/settings.rs b/lib/src/settings.rs
index e2f3862a7..dccf90100 100644
--- a/lib/src/settings.rs
+++ b/lib/src/settings.rs
@@ -13,8 +13,11 @@
 // limitations under the License.
 
 use std::path::Path;
+use std::sync::{Arc, Mutex};
 
 use chrono::DateTime;
+use rand::prelude::*;
+use rand_chacha::ChaCha20Rng;
 
 use crate::backend::{Signature, Timestamp};
 
@@ -22,6 +25,7 @@ use crate::backend::{Signature, Timestamp};
 pub struct UserSettings {
     config: config::Config,
     timestamp: Option<Timestamp>,
+    rng: Arc<JJRng>,
 }
 
 #[derive(Debug, Clone)]
@@ -39,10 +43,22 @@ fn get_timestamp_config(config: &config::Config, key: &str) -> Option<Timestamp>
     }
 }
 
+fn get_rng_seed_config(config: &config::Config) -> Option<u64> {
+    config
+        .get_string("debug.randomness-seed")
+        .ok()
+        .and_then(|str| str.parse().ok())
+}
+
 impl UserSettings {
     pub fn from_config(config: config::Config) -> Self {
         let timestamp = get_timestamp_config(&config, "user.timestamp");
-        UserSettings { config, timestamp }
+        let rng_seed = get_rng_seed_config(&config);
+        UserSettings {
+            config,
+            timestamp,
+            rng: Arc::new(JJRng::new(rng_seed)),
+        }
     }
 
     pub fn incorporate_toml_strings(
@@ -54,8 +70,13 @@ impl UserSettings {
             config_builder =
                 config_builder.add_source(config::File::from_str(s, config::FileFormat::Toml));
         }
-        self.config = config_builder.build()?;
-        self.timestamp = get_timestamp_config(&self.config, "user.timestamp");
+        let new_config = config_builder.build()?;
+        let new_rng_seed = get_rng_seed_config(&new_config);
+        if new_rng_seed != get_rng_seed_config(&self.config) {
+            self.rng.reset(new_rng_seed);
+        }
+        self.timestamp = get_timestamp_config(&new_config, "user.timestamp");
+        self.config = new_config;
         Ok(())
     }
 
@@ -71,6 +92,10 @@ impl UserSettings {
         Ok(RepoSettings { _config: config })
     }
 
+    pub fn get_rng(&self) -> Arc<JJRng> {
+        self.rng.clone()
+    }
+
     pub fn user_name(&self) -> String {
         self.config
             .get_string("user.name")
@@ -146,3 +171,37 @@ impl UserSettings {
         &self.config
     }
 }
+
+/// This Rng uses interior mutability to allow generating random values using an
+/// immutable reference. It also fixes a specific seedable RNG for
+/// reproducibility.
+#[derive(Debug)]
+pub struct JJRng(Mutex<ChaCha20Rng>);
+impl JJRng {
+    /// Wraps Rng::gen but only requires an immutable reference.
+    pub fn gen<T>(&self) -> T
+    where
+        rand::distributions::Standard: rand::distributions::Distribution<T>,
+    {
+        let mut rng = self.0.lock().unwrap();
+        rng.gen()
+    }
+
+    /// Creates a new RNGs. Could be made public, but we'd like to encourage all
+    /// RNGs references to point to the same RNG.
+    fn new(seed: Option<u64>) -> Self {
+        Self(Mutex::new(JJRng::internal_rng_from_seed(seed)))
+    }
+
+    fn reset(&self, seed: Option<u64>) {
+        let mut rng = self.0.lock().unwrap();
+        *rng = JJRng::internal_rng_from_seed(seed)
+    }
+
+    fn internal_rng_from_seed(seed: Option<u64>) -> ChaCha20Rng {
+        match seed {
+            Some(seed) => ChaCha20Rng::seed_from_u64(seed),
+            None => ChaCha20Rng::from_entropy(),
+        }
+    }
+}
diff --git a/src/config.rs b/src/config.rs
index 7ddf20b7e..bfbd1a2e5 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -100,6 +100,11 @@ fn env_overrides() -> config::Config {
     if let Ok(value) = env::var("JJ_TIMESTAMP") {
         builder = builder.set_override("user.timestamp", value).unwrap();
     }
+    if let Ok(value) = env::var("JJ_RANDOMNESS_SEED") {
+        builder = builder
+            .set_override("debug.randomness-seed", value)
+            .unwrap();
+    }
     if let Ok(value) = env::var("JJ_OP_TIMESTAMP") {
         builder = builder.set_override("operation.timestamp", value).unwrap();
     }