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}