From 0174d4b157de4ada4c7b8a829ab963d7f30914fa Mon Sep 17 00:00:00 2001
From: Caio Raposo <caioraposo@disroot.org>
Date: Thu, 7 Sep 2023 23:28:44 -0300
Subject: [PATCH] Add std.randseed and use global RNG

---
 Cargo.lock              |  1 +
 Cargo.toml              |  1 +
 src/runtime/lib/rand.rs | 39 ++++++++++++++++++++++++++++++++++-----
 3 files changed, 36 insertions(+), 5 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index f8cf8d6..69bf332 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -173,6 +173,7 @@ dependencies = [
  "inventory",
  "os_pipe",
  "rand",
+ "rand_chacha",
  "regex",
  "serde",
  "serde_json",
diff --git a/Cargo.toml b/Cargo.toml
index 8ced449..b2293b1 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -29,6 +29,7 @@ serde_json = "1.0"
 base64 = "0.13"
 hex = "0.4"
 rand = "0.8.5"
+rand_chacha = "0.3.1"
 
 [dev-dependencies]
 assert_matches = "1.5"
diff --git a/src/runtime/lib/rand.rs b/src/runtime/lib/rand.rs
index 506e576..226b806 100644
--- a/src/runtime/lib/rand.rs
+++ b/src/runtime/lib/rand.rs
@@ -1,4 +1,13 @@
+//! This module uses the `ChaCha8Rng` pseudo-random generator from the rand
+//! crate. It is suitable for tasks such as simulation, but should not be used
+//! for applications like criptography or gambling games.
+//! see: <https://rust-random.github.io/rand/rand_chacha/struct.ChaCha8Rng.html>
+
+use std::cell::RefCell;
+
 use gc::{Finalize, Trace};
+use rand::{Rng, SeedableRng, thread_rng};
+use rand_chacha::ChaCha8Rng;
 
 use super::{
 	Float,
@@ -8,10 +17,12 @@ use super::{
 	Value,
 	CallContext,
 };
-use rand::{Rng, thread_rng};
 
 inventory::submit! { RustFun::from(Rand) }
 inventory::submit! { RustFun::from(RandInt) }
+inventory::submit! { RustFun::from(RandSeed) }
+
+thread_local!(static RNG: RefCell<ChaCha8Rng> = RefCell::new(ChaCha8Rng::from_rng(thread_rng()).unwrap()));
 
 #[derive(Trace, Finalize)]
 struct Rand;
@@ -19,14 +30,16 @@ struct Rand;
 #[derive(Trace, Finalize)]
 struct RandInt;
 
+#[derive(Trace, Finalize)]
+struct RandSeed;
+
 impl NativeFun for Rand {
 	fn name(&self) -> &'static str { "std.rand" }
 
 	fn call(&self, context: CallContext) -> Result<Value, Panic> {
-		let mut rng = thread_rng();
 		let args = context.args();
 		if args.is_empty() {
-			Ok(Value::Float(Float(rng.gen::<f64>())))
+			Ok(Value::Float(Float(RNG.with(|rng| rng.borrow_mut().gen::<f64>()))))
 		} else {
 			Err(Panic::invalid_args(args.len() as u32, 0, context.pos))
 		}
@@ -37,11 +50,27 @@ impl NativeFun for RandInt {
 	fn name(&self) -> &'static str { "std.randint" }
 
 	fn call(&self, context: CallContext) -> Result<Value, Panic> {
-		let mut rng = thread_rng();
 		match context.args() {
-			[ Value::Int(m), Value::Int(n) ] => Ok(Value::Int(rng.gen_range(*m..=*n))),
+			[ Value::Int(m), Value::Int(n) ] => Ok(Value::Int(
+				RNG.with(|rng| rng.borrow_mut().gen_range(*m..=*n))
+			)),
 			[ other, _ ] => Err(Panic::type_error(other.copy(), "int", context.pos)),
 			args => Err(Panic::invalid_args(args.len() as u32, 2, context.pos))
 		}
 	}
 }
+
+impl NativeFun for RandSeed {
+	fn name(&self) -> &'static str { "std.randseed" }
+
+	fn call(&self, context: CallContext) -> Result<Value, Panic> {
+		match context.args() {
+			[ Value::Int(n) ] => {
+				RNG.with(|rng| *rng.borrow_mut() = ChaCha8Rng::seed_from_u64(*n as u64));
+				Ok(Value::default())
+			},
+			[ other ] => Err(Panic::type_error(other.copy(), "int", context.pos)),
+			args => Err(Panic::invalid_args(args.len() as u32, 1, context.pos))
+		}
+	}
+}