fugue_evo/checkpoint/
rng.rs

1//! Serializable RNG snapshots for reproducible checkpoint resume (EV-02)
2//!
3//! Reproducible resume requires an RNG whose entire internal state can be
4//! captured and restored byte-for-byte. `rand`'s generic `Rng`/`RngCore`
5//! traits provide no serialization hook, so a resumed run seeded from a fresh
6//! (or differently-seeded) generator silently diverges from the trajectory a
7//! continuous run would have taken.
8//!
9//! The [`SnapshotRng`] trait fills that gap for the ChaCha family of counter
10//! based RNGs (`ChaCha8Rng`, `ChaCha12Rng`, `ChaCha20Rng`), whose state is
11//! fully serializable via `serde`. Capturing the RNG alongside the population
12//! in a [`Checkpoint`](crate::checkpoint::Checkpoint) makes resume
13//! bit-identical to a run that was never interrupted.
14//!
15//! # Reproducibility contract
16//!
17//! Bit-identical resume is only guaranteed when the algorithm draws all of its
18//! randomness from a `SnapshotRng` and that RNG is captured into the checkpoint
19//! (via [`Checkpoint::with_rng`](crate::checkpoint::Checkpoint::with_rng)) and
20//! restored on resume (via
21//! [`Checkpoint::restore_rng`](crate::checkpoint::Checkpoint::restore_rng)).
22//! Non-ChaCha generators (e.g. `StdRng`, `ThreadRng`) cannot be snapshotted and
23//! therefore cannot provide reproducible resume.
24
25use rand_chacha::{ChaCha12Rng, ChaCha20Rng, ChaCha8Rng};
26
27use crate::error::CheckpointError;
28
29/// An RNG whose complete internal state can be captured to bytes and restored,
30/// enabling bit-identical checkpoint resume.
31///
32/// Implemented for the ChaCha family (`ChaCha8Rng`, `ChaCha12Rng`,
33/// `ChaCha20Rng`), all of which expose fully serializable state.
34pub trait SnapshotRng: Sized {
35    /// Serialize the full RNG state to a byte buffer.
36    fn capture(&self) -> Result<Vec<u8>, CheckpointError>;
37
38    /// Reconstruct an RNG from bytes previously produced by [`capture`].
39    ///
40    /// [`capture`]: SnapshotRng::capture
41    fn restore(bytes: &[u8]) -> Result<Self, CheckpointError>;
42}
43
44macro_rules! impl_snapshot_rng {
45    ($rng:ty) => {
46        impl SnapshotRng for $rng {
47            fn capture(&self) -> Result<Vec<u8>, CheckpointError> {
48                bincode::serialize(self).map_err(|e| CheckpointError::SerializeError(Box::new(e)))
49            }
50
51            fn restore(bytes: &[u8]) -> Result<Self, CheckpointError> {
52                bincode::deserialize(bytes)
53                    .map_err(|e| CheckpointError::DeserializeError(Box::new(e)))
54            }
55        }
56    };
57}
58
59impl_snapshot_rng!(ChaCha8Rng);
60impl_snapshot_rng!(ChaCha12Rng);
61impl_snapshot_rng!(ChaCha20Rng);
62
63#[cfg(test)]
64mod tests {
65    use super::*;
66    use rand::{Rng, SeedableRng};
67
68    #[test]
69    fn test_snapshot_rng_round_trip_is_bit_identical() {
70        // regression: EV-02 - a captured RNG must restore to the exact same
71        // state, producing an identical sequence of draws afterwards.
72        let mut rng = ChaCha8Rng::seed_from_u64(42);
73        // Advance the stream so we are not just testing the seeded state.
74        for _ in 0..37 {
75            let _: u64 = rng.gen();
76        }
77
78        let bytes = rng.capture().unwrap();
79        let mut restored = ChaCha8Rng::restore(&bytes).unwrap();
80
81        // Both generators must now produce an identical sequence.
82        for _ in 0..1000 {
83            let a: u64 = rng.gen();
84            let b: u64 = restored.gen();
85            assert_eq!(a, b, "restored RNG diverged from the original");
86        }
87    }
88
89    #[test]
90    fn test_snapshot_rng_variants() {
91        let mut c12 = ChaCha12Rng::seed_from_u64(7);
92        let _: u32 = c12.gen();
93        let bytes = c12.capture().unwrap();
94        let mut r12 = ChaCha12Rng::restore(&bytes).unwrap();
95        assert_eq!(c12.gen::<u64>(), r12.gen::<u64>());
96
97        let mut c20 = ChaCha20Rng::seed_from_u64(9);
98        let _: u32 = c20.gen();
99        let bytes = c20.capture().unwrap();
100        let mut r20 = ChaCha20Rng::restore(&bytes).unwrap();
101        assert_eq!(c20.gen::<u64>(), r20.gen::<u64>());
102    }
103}
fugue_evo/checkpoint/rng.rs

fugue_evo/checkpoint/
rng.rs
