1use core::sync::atomic::{AtomicBool, Ordering};
9use std::sync::Arc;
10
11#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
13pub struct Parallelism {
14 pub max_threads: NonZeroThreadCount,
16}
17
18#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
20pub struct NonZeroThreadCount(u32);
21
22impl NonZeroThreadCount {
23 #[must_use]
25 pub const fn new(n: u32) -> Option<Self> {
26 if n == 0 { None } else { Some(Self(n)) }
27 }
28
29 #[must_use]
31 pub const fn one() -> Self {
32 Self(1)
33 }
34
35 #[must_use]
37 pub const fn get(self) -> u32 {
38 self.0
39 }
40}
41
42impl Parallelism {
43 #[must_use]
45 pub const fn serial() -> Self {
46 Self { max_threads: NonZeroThreadCount::one() }
47 }
48
49 #[must_use]
51 pub const fn bounded(max_threads: NonZeroThreadCount) -> Self {
52 Self { max_threads }
53 }
54}
55
56#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
58pub enum Determinism {
59 PreferFast,
61 Strict,
63}
64
65#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
67pub struct MemoryBudget {
68 pub soft_limit_bytes: Option<u64>,
70 pub hard_limit_bytes: Option<u64>,
72}
73
74impl MemoryBudget {
75 #[must_use]
77 pub const fn unlimited() -> Self {
78 Self { soft_limit_bytes: None, hard_limit_bytes: None }
79 }
80}
81
82#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
84pub struct KernelPolicy {
85 pub allow_portable_optimized: bool,
87 pub allow_arch_simd: bool,
89 pub force_scalar: bool,
91}
92
93impl KernelPolicy {
94 #[must_use]
96 pub const fn default_policy() -> Self {
97 Self { allow_portable_optimized: true, allow_arch_simd: true, force_scalar: false }
98 }
99
100 #[must_use]
102 pub const fn scalar_only() -> Self {
103 Self { allow_portable_optimized: false, allow_arch_simd: false, force_scalar: true }
104 }
105}
106
107#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
109pub struct CachePolicy {
110 pub enabled: bool,
112 pub max_bytes: Option<u64>,
114}
115
116impl CachePolicy {
117 #[must_use]
119 pub const fn disabled() -> Self {
120 Self { enabled: false, max_bytes: None }
121 }
122
123 #[must_use]
125 pub const fn enabled(max_bytes: Option<u64>) -> Self {
126 Self { enabled: true, max_bytes }
127 }
128}
129
130#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
132pub struct CacheBudget {
133 pub max_bytes: u64,
135 pub used_bytes: u64,
137}
138
139impl CacheBudget {
140 #[must_use]
142 pub const fn new(max_bytes: u64) -> Self {
143 Self { max_bytes, used_bytes: 0 }
144 }
145
146 #[must_use]
148 pub const fn unlimited() -> Self {
149 Self { max_bytes: u64::MAX, used_bytes: 0 }
150 }
151
152 #[must_use]
154 pub const fn remaining(self) -> u64 {
155 self.max_bytes.saturating_sub(self.used_bytes)
156 }
157
158 #[must_use]
160 pub const fn can_admit(self, additional: u64) -> bool {
161 self.used_bytes.saturating_add(additional) <= self.max_bytes
162 }
163}
164
165#[derive(Clone, Debug, Default, PartialEq)]
167pub struct MonteCarloBudget {
168 pub evaluations: u64,
170 pub samples: u64,
172 pub exact_enumerations: u64,
174}
175
176#[derive(Clone, Copy, Debug, Default, PartialEq)]
178pub struct MonteCarloError {
179 pub stderr: f64,
181 pub samples: u64,
183}
184
185#[derive(Clone, Copy, Debug, PartialEq)]
191pub struct AdaptiveBootstrapBudget {
192 pub enabled: bool,
194 pub min_replicates: u32,
196 pub se_rel_epsilon: f64,
198}
199
200impl AdaptiveBootstrapBudget {
201 #[must_use]
203 pub const fn enabled_default() -> Self {
204 Self { enabled: true, min_replicates: 10, se_rel_epsilon: 0.01 }
205 }
206
207 #[must_use]
209 pub const fn disabled() -> Self {
210 Self { enabled: false, min_replicates: 0, se_rel_epsilon: 0.0 }
211 }
212}
213
214impl Default for AdaptiveBootstrapBudget {
215 fn default() -> Self {
216 Self::enabled_default()
217 }
218}
219
220#[derive(Clone, Copy, Debug, PartialEq)]
225pub struct AdaptiveDrawBudget {
226 pub enabled: bool,
228 pub min_draws: usize,
230 pub quantile_width_rel_epsilon: f64,
232 pub ess_target: f64,
234}
235
236impl AdaptiveDrawBudget {
237 #[must_use]
239 pub const fn enabled_default() -> Self {
240 Self {
241 enabled: true,
242 min_draws: 32,
243 quantile_width_rel_epsilon: 0.01,
244 ess_target: 10_000.0,
245 }
246 }
247
248 #[must_use]
250 pub const fn disabled() -> Self {
251 Self { enabled: false, min_draws: 0, quantile_width_rel_epsilon: 0.0, ess_target: 0.0 }
252 }
253}
254
255impl Default for AdaptiveDrawBudget {
256 fn default() -> Self {
257 Self::enabled_default()
258 }
259}
260
261#[derive(Clone, Debug, Default)]
263pub struct CancellationToken {
264 cancelled: Arc<AtomicBool>,
265}
266
267impl CancellationToken {
268 #[must_use]
270 pub fn new() -> Self {
271 Self::default()
272 }
273
274 pub fn cancel(&self) {
276 self.cancelled.store(true, Ordering::SeqCst);
277 }
278
279 #[must_use]
281 pub fn is_cancelled(&self) -> bool {
282 self.cancelled.load(Ordering::SeqCst)
283 }
284}
285
286pub trait ProgressSink: Send + Sync {
288 fn report(&self, fraction: f64, stage: &str);
290}
291
292#[derive(Clone, Debug, Eq, PartialEq)]
297pub struct RngFactory {
298 master_seed: u64,
299}
300
301impl RngFactory {
302 #[must_use]
304 pub const fn from_seed(master_seed: u64) -> Self {
305 Self { master_seed }
306 }
307
308 #[must_use]
310 pub const fn master_seed(&self) -> u64 {
311 self.master_seed
312 }
313
314 #[must_use]
316 pub fn stream(&self, stream_id: u64) -> CausalRng {
317 let seed = mix_seed(self.master_seed, stream_id);
318 CausalRng::from_seed(seed)
319 }
320}
321
322#[derive(Clone, Debug)]
324pub struct CausalRng {
325 state: u64,
326}
327
328impl CausalRng {
329 #[must_use]
331 pub const fn from_seed(seed: u64) -> Self {
332 Self { state: seed ^ 0x9E37_79B9_7F4A_7C15 }
334 }
335
336 #[must_use]
338 pub const fn from_state(state: u64) -> Self {
339 Self { state }
340 }
341
342 #[must_use]
344 pub const fn state(&self) -> u64 {
345 self.state
346 }
347
348 #[must_use]
350 pub fn next_u64(&mut self) -> u64 {
351 self.state = self.state.wrapping_add(0x9E37_79B9_7F4A_7C15);
352 let mut z = self.state;
353 z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
354 z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
355 z ^ (z >> 31)
356 }
357
358 #[must_use]
360 pub fn next_f64(&mut self) -> f64 {
361 let bits = self.next_u64() >> 11;
363 #[allow(clippy::cast_precision_loss)]
364 {
365 bits as f64 * (1.0 / ((1u64 << 53) as f64))
366 }
367 }
368}
369
370fn mix_seed(master: u64, stream_id: u64) -> u64 {
371 let mut z = master.wrapping_add(stream_id).wrapping_mul(0xD6E8_FEB8_6659_FD93);
372 z = (z ^ (z >> 32)).wrapping_mul(0xD6E8_FEB8_6659_FD93);
373 z ^ (z >> 32)
374}
375
376pub struct ExecutionContext {
378 pub parallelism: Parallelism,
380 pub determinism: Determinism,
382 pub rng: RngFactory,
384 pub memory: MemoryBudget,
386 pub cancellation: CancellationToken,
388 pub progress: Option<Arc<dyn ProgressSink>>,
390 pub kernel_policy: KernelPolicy,
392 pub cache_policy: CachePolicy,
394 pub adaptive_bootstrap: AdaptiveBootstrapBudget,
396 pub adaptive_draws: AdaptiveDrawBudget,
398}
399
400impl ExecutionContext {
401 #[must_use]
412 pub fn for_tests(seed: u64) -> Self {
413 Self {
414 parallelism: Parallelism::serial(),
415 determinism: Determinism::Strict,
416 rng: RngFactory::from_seed(seed),
417 memory: MemoryBudget::unlimited(),
418 cancellation: CancellationToken::new(),
419 progress: None,
420 kernel_policy: KernelPolicy::scalar_only(),
421 cache_policy: CachePolicy::disabled(),
422 adaptive_bootstrap: AdaptiveBootstrapBudget::disabled(),
424 adaptive_draws: AdaptiveDrawBudget::disabled(),
425 }
426 }
427
428 #[must_use]
430 pub fn production(seed: u64, max_threads: u32) -> Self {
431 let threads =
432 NonZeroThreadCount::new(max_threads.max(1)).unwrap_or_else(NonZeroThreadCount::one);
433 Self {
434 parallelism: Parallelism::bounded(threads),
435 determinism: Determinism::Strict,
436 rng: RngFactory::from_seed(seed),
437 memory: MemoryBudget::unlimited(),
438 cancellation: CancellationToken::new(),
439 progress: None,
440 kernel_policy: KernelPolicy::default_policy(),
441 cache_policy: CachePolicy::enabled(None),
442 adaptive_bootstrap: AdaptiveBootstrapBudget::enabled_default(),
443 adaptive_draws: AdaptiveDrawBudget::enabled_default(),
444 }
445 }
446}
447
448impl core::fmt::Debug for ExecutionContext {
449 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
450 f.debug_struct("ExecutionContext")
451 .field("parallelism", &self.parallelism)
452 .field("determinism", &self.determinism)
453 .field("rng", &self.rng)
454 .field("memory", &self.memory)
455 .field("cancellation_cancelled", &self.cancellation.is_cancelled())
456 .field("progress_is_some", &self.progress.is_some())
457 .field("kernel_policy", &self.kernel_policy)
458 .field("cache_policy", &self.cache_policy)
459 .field("adaptive_bootstrap", &self.adaptive_bootstrap)
460 .field("adaptive_draws", &self.adaptive_draws)
461 .finish()
462 }
463}
464
465#[cfg(test)]
466mod tests {
467 use super::*;
468
469 #[test]
470 fn rng_streams_are_deterministic() {
471 let factory = RngFactory::from_seed(42);
472 let mut a1 = factory.stream(0);
473 let mut a2 = factory.stream(0);
474 let mut b = factory.stream(1);
475 let seq_a1: Vec<u64> = (0..8).map(|_| a1.next_u64()).collect();
476 let seq_a2: Vec<u64> = (0..8).map(|_| a2.next_u64()).collect();
477 let seq_b: Vec<u64> = (0..8).map(|_| b.next_u64()).collect();
478 assert_eq!(seq_a1, seq_a2);
479 assert_ne!(seq_a1, seq_b);
480 }
481
482 #[test]
483 fn independent_factories_same_seed_match() {
484 let f1 = RngFactory::from_seed(7);
485 let f2 = RngFactory::from_seed(7);
486 let mut s1 = f1.stream(99);
487 let mut s2 = f2.stream(99);
488 for _ in 0..32 {
489 assert_eq!(s1.next_u64(), s2.next_u64());
490 }
491 }
492
493 #[test]
494 fn cancellation_is_shared_across_clones() {
495 let token = CancellationToken::new();
496 let clone = token.clone();
497 assert!(!token.is_cancelled());
498 clone.cancel();
499 assert!(token.is_cancelled());
500 }
501
502 #[test]
503 fn f64_draws_are_in_unit_interval() {
504 let mut rng = CausalRng::from_seed(123);
505 for _ in 0..1000 {
506 let x = rng.next_f64();
507 assert!((0.0..1.0).contains(&x));
508 }
509 }
510}