dupehell_core/context.rs
1// DupeHell -- MIT License
2//
3// Synthetic multi-domain dataset generator for record linkage benchmarking.
4// No liability for misuse.
5
6use std::collections::HashMap;
7use std::path::Path;
8
9use rayon::prelude::*;
10
11/// A loaded pool entry — flat list of strings.
12pub type Pool = Vec<String>;
13
14/// All pools keyed by name.
15#[derive(Debug)]
16pub struct PoolStore {
17 pub pools: HashMap<String, Pool>,
18}
19
20/// Known locale codes supported by pool files.
21pub const LOCALES: &[&str] = &["en", "fr", "de", "es", "it", "pt"];
22
23fn is_locale_key(key: &str) -> bool {
24 LOCALES.contains(&key)
25}
26
27/// Load and parse a single pool file, returning `(name, pool)` unless the
28/// file yields no usable pool (an object with neither a matching locale key
29/// nor region-keyed arrays) — split out of `PoolStore::load` so it can run
30/// on `rayon`'s pool, one call per file, independently of every other file.
31fn load_one(entry: &std::fs::DirEntry, locale: &str) -> Result<Option<(String, Pool)>, String> {
32 let name = entry
33 .path()
34 .file_stem()
35 .unwrap()
36 .to_string_lossy()
37 .to_string();
38 // `from_reader` on a raw `File` reads via small unbuffered calls (no
39 // internal buffering in serde_json) — reading the whole small pool
40 // file into memory first and parsing with `from_str` is markedly
41 // faster, and this runs once per pool file (151 of them) at every
42 // startup.
43 let raw = std::fs::read_to_string(entry.path())
44 .map_err(|e| format!("cannot open {name}.json: {e}"))?;
45 let data: serde_json::Value =
46 serde_json::from_str(&raw).map_err(|e| format!("cannot parse {name}.json: {e}"))?;
47
48 let pool = match &data {
49 serde_json::Value::Array(arr) => arr
50 .iter()
51 .filter_map(|v| v.as_str())
52 .map(String::from)
53 .collect::<Pool>(),
54 serde_json::Value::Object(_map) => {
55 if let Some(arr) = data
56 .get(locale)
57 .or_else(|| data.get("en"))
58 .and_then(|v| v.as_array())
59 {
60 arr.iter()
61 .filter_map(|v| v.as_str())
62 .map(String::from)
63 .collect::<Pool>()
64 } else if let Some(obj) = data.as_object() {
65 // No locale key found — check if this is a region-keyed pool
66 // (e.g. french_cities.json with "ile_de_france" keys).
67 let has_locale_keys = obj.keys().any(|k| is_locale_key(k));
68 if !has_locale_keys {
69 obj.values()
70 .filter_map(|v| v.as_array())
71 .flatten()
72 .filter_map(|v| v.as_str().or_else(|| v.get(0).and_then(|s| s.as_str())))
73 .map(String::from)
74 .collect::<Pool>()
75 } else {
76 Vec::new()
77 }
78 } else {
79 Vec::new()
80 }
81 }
82 _ => return Ok(None),
83 };
84 if pool.is_empty() {
85 return Ok(None);
86 }
87 Ok(Some((name, pool)))
88}
89
90impl PoolStore {
91 /// Load all JSON files from `pools_dir` (non-recursive), selecting `locale` data.
92 /// Falls back to `"en"` if the requested locale is not found.
93 pub fn load(pools_dir: &str, locale: &str) -> Result<Self, String> {
94 let dir = Path::new(pools_dir);
95 if !dir.is_dir() {
96 return Err(format!("pools dir not found: {pools_dir}"));
97 }
98 let mut entries: Vec<_> = std::fs::read_dir(dir)
99 .map_err(|e| format!("cannot read {pools_dir}: {e}"))?
100 .filter_map(|r| r.ok())
101 .filter(|e| e.path().extension().is_some_and(|ext| ext == "json"))
102 .collect();
103 entries.sort_by_key(|e| e.file_name());
104
105 // Each file is read + parsed independently of every other
106 // (measured in isolation ~2.6-3.4x faster on the read+parse
107 // portion) — the previous sequential loop paid this startup cost
108 // once per file, one process at a time; a scripted campaign that spawns
109 // many short-lived processes back to back pays it many times over.
110 // `collect::<Result<Vec<_>, _>>()` (not `filter_map`+`ok()`)
111 // preserves the previous behavior of propagating the FIRST read/
112 // parse error instead of silently dropping it.
113 let loaded: Vec<Option<(String, Pool)>> = entries
114 .par_iter()
115 .map(|entry| load_one(entry, locale))
116 .collect::<Result<Vec<_>, String>>()?;
117 // A `HashMap`'s contents don't depend on insertion order, and pool
118 // file names are unique (drawn from distinct `file_stem`s of one
119 // directory listing), so collecting from `par_iter`'s
120 // possibly-reordered output is exactly as deterministic as the
121 // previous sequential-order insertion.
122 let pools: HashMap<String, Pool> = loaded.into_iter().flatten().collect();
123 Ok(Self { pools })
124 }
125
126 pub fn get(&self, name: &str) -> Option<&Pool> {
127 self.pools.get(name)
128 }
129}
130
131/// The engine context holding config and pool data.
132///
133/// Deliberately not `Clone`: `pool_store` holds
134/// ~2.8 MB across 151 pools, so a `Context::clone()` would be a silent deep
135/// copy of tens of thousands of `String`s. Nothing in the crate ever clones
136/// a `Context` — pass `&Context` (the pattern already used throughout
137/// `pipeline.rs`/`entity_gen.rs`/`column_gen.rs`) or wrap it in `Arc` if a
138/// genuine shared-ownership need comes up.
139#[derive(Debug)]
140pub struct Context {
141 pub pool_store: PoolStore,
142 pub locale: String,
143 pub watermark_map: std::collections::HashMap<u64, u64>, // col_tag → masked value
144}
145
146impl Context {
147 const WATERMARK_SECRET: &'static str = "DupeHell-WATERMARK-v0.4-educational-only-2026";
148
149 /// Build a Context from a domain name + locale + path to pools directory.
150 pub fn new(domain: &str, locale: &str, pools_dir: &str) -> Result<Self, String> {
151 let pool_store = PoolStore::load(pools_dir, locale)?;
152 log::info!(
153 "Context loaded: domain={domain}, locale={locale}, pools={}",
154 pool_store.pools.len()
155 );
156 Ok(Self {
157 pool_store,
158 locale: locale.to_string(),
159 watermark_map: std::collections::HashMap::new(),
160 })
161 }
162
163 /// Enable watermarking by computing per-column tags from the pipeline config.
164 pub fn enable_watermark(&mut self, domain: &str, size: usize, seed: u64) {
165 use sha2::{Digest, Sha256};
166 for &tag in &[
167 0x53534e, // "SSN"
168 0x50484f4e, // "PHONE"
169 0x50414e, // "PAN"
170 0x4d4544, // "MEDICARE"
171 0x4f4643, // "OFFICE_PHONE"
172 0x504153, // "PASSPORT"
173 0x414354, // "ACCOUNT"
174 0x424152, // "BARCODE"
175 0x494343, // "ICCID"
176 0x555043, // "UPC"
177 ] {
178 let input = format!(
179 "{}{}{}{}{}",
180 Self::WATERMARK_SECRET,
181 domain,
182 size,
183 seed,
184 tag
185 );
186 let hash = Sha256::digest(input.as_bytes());
187 let wm = u64::from_le_bytes(hash[..8].try_into().unwrap());
188 self.watermark_map.insert(tag, wm);
189 }
190 }
191
192 /// Return the watermark mask for a given column tag (last 3 digits, 0..999).
193 pub fn watermark_3digits(&self, tag: u64) -> u64 {
194 self.watermark_map.get(&tag).copied().unwrap_or(0) % 1000
195 }
196
197 /// Return the watermark mask for a given column tag (last 2 digits, 0..99).
198 pub fn watermark_2digits(&self, tag: u64) -> u64 {
199 self.watermark_map.get(&tag).copied().unwrap_or(0) % 100
200 }
201
202 /// Create a minimal context for testing (no pools loaded, watermark disabled).
203 /// Watermark helpers return 0, ensuring deterministic test output.
204 #[cfg(test)]
205 pub fn test() -> Self {
206 Self {
207 pool_store: PoolStore {
208 pools: HashMap::new(),
209 },
210 locale: "en".to_string(),
211 watermark_map: HashMap::new(),
212 }
213 }
214}