Skip to content

Commit cdf8391

Browse files
mamcxgefjon
andauthored
Add rng support to ProcedureContext on Rust (#3865)
# Description of Changes As the title says, required for #3538. # Expected complexity level and risk 1 # Testing Will be done on the `uuid` pr. --------- Signed-off-by: Mario Montoya <mamcx@elmalabarista.com> Co-authored-by: Phoebe Goldman <phoebe@clockworklabs.io>
1 parent e5f1558 commit cdf8391

3 files changed

Lines changed: 61 additions & 7 deletions

File tree

crates/bindings/src/lib.rs

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1084,13 +1084,25 @@ pub struct ProcedureContext {
10841084

10851085
/// Methods for performing HTTP requests.
10861086
pub http: crate::http::HttpClient,
1087-
// TODO: Add rng?
1087+
// TODO: Change rng?
10881088
// Complex and requires design because we may want procedure RNG to behave differently from reducer RNG,
10891089
// as it could actually be seeded by OS randomness rather than a deterministic source.
1090+
#[cfg(feature = "rand08")]
1091+
rng: std::cell::OnceCell<StdbRng>,
10901092
}
10911093

10921094
#[cfg(feature = "unstable")]
10931095
impl ProcedureContext {
1096+
fn new(sender: Identity, connection_id: Option<ConnectionId>, timestamp: Timestamp) -> Self {
1097+
Self {
1098+
sender,
1099+
timestamp,
1100+
connection_id,
1101+
http: http::HttpClient {},
1102+
#[cfg(feature = "rand08")]
1103+
rng: std::cell::OnceCell::new(),
1104+
}
1105+
}
10941106
/// Read the current module's [`Identity`].
10951107
pub fn identity(&self) -> Identity {
10961108
// Hypothetically, we *could* read the module identity out of the system tables.

crates/bindings/src/rng.rs

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
#[cfg(feature = "unstable")]
2+
use crate::ProcedureContext;
13
use crate::{rand, ReducerContext};
24
use core::cell::UnsafeCell;
35
use core::marker::PhantomData;
@@ -49,6 +51,51 @@ impl ReducerContext {
4951
}
5052
}
5153

54+
#[cfg(feature = "unstable")]
55+
impl ProcedureContext {
56+
/// Generates a random value.
57+
///
58+
/// Similar to [`rand::random()`], but using [`StdbRng`] instead.
59+
///
60+
/// See also [`ProcedureContext::rng()`].
61+
#[cfg(feature = "unstable")]
62+
pub fn random<T>(&self) -> T
63+
where
64+
Standard: Distribution<T>,
65+
{
66+
Standard.sample(&mut self.rng())
67+
}
68+
69+
/// Retrieve the random number generator for this procedure transaction,
70+
/// seeded by the timestamp of the procedure call.
71+
///
72+
/// If you only need a single random value, you can use [`ProcedureContext::random()`].
73+
///
74+
/// # Examples
75+
/// ```no_run
76+
/// # #[cfg(target_arch = "wasm32")] mod demo {
77+
/// use spacetimedb::{procedure, ProcedureContext};
78+
/// use rand::Rng;
79+
///
80+
/// #[spacetimedb::procedure]
81+
/// fn rng_demo(ctx: &spacetimedb::ProcedureContext) {
82+
/// // Can be used in method chaining style:
83+
/// let digit = ctx.rng().gen_range(0..=9);
84+
///
85+
/// // Or, cache locally for reuse:
86+
/// let mut rng = ctx.rng();
87+
/// let floats: Vec<f32> = rng.sample_iter(rand::distributions::Standard).collect();
88+
/// }
89+
/// # }
90+
/// ```
91+
///
92+
/// For more information, see [`StdbRng`] and [`rand::Rng`].
93+
#[cfg(feature = "unstable")]
94+
pub fn rng(&self) -> &StdbRng {
95+
self.rng.get_or_init(|| StdbRng::seed_from_ts(self.timestamp))
96+
}
97+
}
98+
5299
/// A reference to the random number generator for this reducer call.
53100
///
54101
/// An instance can be obtained via [`ReducerContext::rng()`]. Import

crates/bindings/src/rt.rs

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1034,12 +1034,7 @@ extern "C" fn __call_procedure__(
10341034
let timestamp = Timestamp::from_micros_since_unix_epoch(timestamp as i64);
10351035

10361036
// Assemble the `ProcedureContext`.
1037-
let ctx = ProcedureContext {
1038-
connection_id: conn_id,
1039-
sender,
1040-
timestamp,
1041-
http: crate::http::HttpClient {},
1042-
};
1037+
let ctx = ProcedureContext::new(sender, conn_id, timestamp);
10431038

10441039
// Grab the list of procedures, which is populated by the preinit functions.
10451040
let procedures = PROCEDURES.get().unwrap();

0 commit comments

Comments
 (0)