Skip to main content

dupehell_core/
rng.rs

1// DupeHell -- MIT License
2//
3// Synthetic multi-domain dataset generator for record linkage benchmarking.
4// No liability for misuse.
5
6use rand::Rng as _; // `next_u64` (named `RngCore` before rand 0.10)
7use rand::SeedableRng;
8use rand_pcg::Pcg64Mcg;
9
10/// A thin wrapper around PCG64 for deterministic generation.
11///
12/// Python passes a master seed; this RNG can produce sub-seeds for
13/// parallel batches via `fork()`.
14#[derive(Clone, Debug)]
15pub struct Rng {
16    inner: Pcg64Mcg,
17}
18
19impl Rng {
20    /// Create a new RNG from a 64-bit seed.
21    pub fn new(seed: u64) -> Self {
22        Self {
23            inner: Pcg64Mcg::seed_from_u64(seed),
24        }
25    }
26
27    /// Return a random `u64`.
28    pub fn next_u64(&mut self) -> u64 {
29        self.inner.next_u64()
30    }
31
32    /// Return a random `usize` in `0..bound`.
33    ///
34    /// Uses Lemire's multiply-shift method instead of `% bound`, which would
35    /// otherwise skew low values for any `bound` that doesn't divide 2^64.
36    pub fn next_usize(&mut self, bound: usize) -> usize {
37        if bound <= 1 {
38            return 0;
39        }
40        ((self.next_u64() as u128 * bound as u128) >> 64) as usize
41    }
42
43    /// Return a random `f64` in `[0, 1)`.
44    pub fn next_f64(&mut self) -> f64 {
45        // convert to [0, 1) — 53 bits of precision
46        (self.next_u64() >> 11) as f64 * (1.0 / 9007199254740992.0)
47    }
48
49    /// Fork: create a fresh RNG seeded with a value from this one.
50    /// Useful for spawning per-batch or per-column RNGs.
51    pub fn fork(&mut self) -> Self {
52        Self::new(self.next_u64())
53    }
54}
55
56/// SplitMix64's output function: a bijective avalanche mix of a `u64`.
57fn mix64(mut z: u64) -> u64 {
58    z = z.wrapping_add(0x9E37_79B9_7F4A_7C15);
59    z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
60    z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
61    z ^ (z >> 31)
62}
63
64/// Seed for one independent random stream of a run, identified by `path`
65/// (a [`stream`] tag first, then e.g. an entity key and a batch index).
66///
67/// Replaces offset arithmetic such as `seed + batch_idx * 1000 + 1000`,
68/// which handed every entity the same seed for the same batch index (so
69/// column k of two entities drew identical values row by row) and made
70/// nearby seeds overlap (seed 42 batch 1 == seed 1042 batch 0). Every
71/// distinct `(seed, path)` here maps to an unrelated stream.
72pub fn derive_seed(seed: u64, path: &[u64]) -> u64 {
73    path.iter().fold(mix64(seed), |h, &p| mix64(h ^ p))
74}
75
76/// Stable 64-bit key for a name (FNV-1a), to give each entity its own
77/// [`derive_seed`] streams independently of its position in the schema.
78pub fn name_key(name: &str) -> u64 {
79    name.bytes().fold(0xcbf2_9ce4_8422_2325, |h, b| {
80        (h ^ u64::from(b)).wrapping_mul(0x0000_0100_0000_01b3)
81    })
82}
83
84/// Stream tags for [`derive_seed`]: one per independent use of randomness.
85pub mod stream {
86    /// Column values of one batch of an entity's base rows.
87    pub const ENTITY_BATCH: u64 = 1;
88    /// Which masters get duplicated, and the noise applied to them.
89    pub const DUPLICATES: u64 = 2;
90    /// FK remap targets of an entity's rows.
91    pub const FK_REMAP: u64 = 3;
92    /// Pairs drawn for one `hard_neg_types` entry.
93    pub const HARD_NEG: u64 = 4;
94    /// Seed of one `--chunk-size` chunk.
95    pub const CHUNK: u64 = 5;
96    /// Reservoir sampling of an entity's FK identifier pool.
97    pub const FK_POOL: u64 = 6;
98}
99
100#[cfg(test)]
101mod tests {
102    use super::*;
103
104    #[test]
105    fn test_derive_seed_separates_paths_and_seeds() {
106        use std::collections::HashSet;
107        let mut seen = HashSet::new();
108        for seed in 0..50u64 {
109            for entity in 0..5u64 {
110                for batch in 0..20u64 {
111                    assert!(seen.insert(derive_seed(seed, &[stream::ENTITY_BATCH, entity, batch])));
112                }
113            }
114        }
115        // The old `seed + batch_idx * 1000 + 1000` scheme collided here.
116        assert_ne!(
117            derive_seed(42, &[stream::ENTITY_BATCH, 0, 1]),
118            derive_seed(1042, &[stream::ENTITY_BATCH, 0, 0])
119        );
120        assert_ne!(name_key("contact"), name_key("lead"));
121    }
122
123    #[test]
124    fn test_deterministic() {
125        let mut a = Rng::new(42);
126        let mut b = Rng::new(42);
127        for _ in 0..100 {
128            assert_eq!(a.next_u64(), b.next_u64());
129        }
130    }
131
132    #[test]
133    fn test_fork_differs() {
134        let mut parent = Rng::new(42);
135        let mut child = parent.fork();
136        // parent and child should produce different sequences
137        assert_ne!(parent.next_u64(), child.next_u64());
138    }
139
140    #[test]
141    fn test_next_usize_bounded() {
142        let mut rng = Rng::new(0);
143        for _ in 0..1000 {
144            let v = rng.next_usize(10);
145            assert!(v < 10, "value {v} out of range");
146        }
147    }
148
149    #[test]
150    fn test_next_f64_range() {
151        let mut rng = Rng::new(1);
152        for _ in 0..1000 {
153            let v = rng.next_f64();
154            assert!((0.0..1.0).contains(&v), "value {v} out of [0,1)");
155        }
156    }
157}