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.
21const 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 // (perf-hunt hunt0109/H12, measured isolated ~2.6-3.4x on the
107 // read+parse portion) — the previous sequential loop paid this
108 // startup cost once per file, one process at a time; a `--force`
109 // rebuild or a `validate_all_domains.py`-style campaign that spawns
110 // many short-lived processes back to back pays it many times over.
111 // `collect::<Result<Vec<_>, _>>()` (not `filter_map`+`ok()`)
112 // preserves the previous behavior of propagating the FIRST read/
113 // parse error instead of silently dropping it.
114 let loaded: Vec<Option<(String, Pool)>> = entries
115 .par_iter()
116 .map(|entry| load_one(entry, locale))
117 .collect::<Result<Vec<_>, String>>()?;
118 // A `HashMap`'s contents don't depend on insertion order, and pool
119 // file names are unique (drawn from distinct `file_stem`s of one
120 // directory listing), so collecting from `par_iter`'s
121 // possibly-reordered output is exactly as deterministic as the
122 // previous sequential-order insertion.
123 let pools: HashMap<String, Pool> = loaded.into_iter().flatten().collect();
124 Ok(Self { pools })
125 }
126
127 pub fn get(&self, name: &str) -> Option<&Pool> {
128 self.pools.get(name)
129 }
130}
131
132/// The engine context holding config and pool data.
133///
134/// Deliberately not `Clone` (perf-hunt hunt0109/H13): `pool_store` holds
135/// ~2.8 MB across 151 pools, so a `Context::clone()` would be a silent deep
136/// copy of tens of thousands of `String`s. Nothing in the crate ever clones
137/// a `Context` — pass `&Context` (the pattern already used throughout
138/// `pipeline.rs`/`entity_gen.rs`/`column_gen.rs`) or wrap it in `Arc` if a
139/// genuine shared-ownership need comes up.
140#[derive(Debug)]
141pub struct Context {
142 pub pool_store: PoolStore,
143 pub locale: String,
144 pub watermark_map: std::collections::HashMap<u64, u64>, // col_tag → masked value
145}
146
147impl Context {
148 const WATERMARK_SECRET: &'static str = "DupeHell-WATERMARK-v0.4-educational-only-2026";
149
150 /// Build a Context from a domain name + locale + path to pools directory.
151 pub fn new(domain: &str, locale: &str, pools_dir: &str) -> Result<Self, String> {
152 let pool_store = PoolStore::load(pools_dir, locale)?;
153 log::info!(
154 "Context loaded: domain={domain}, locale={locale}, pools={}",
155 pool_store.pools.len()
156 );
157 Ok(Self {
158 pool_store,
159 locale: locale.to_string(),
160 watermark_map: std::collections::HashMap::new(),
161 })
162 }
163
164 /// Enable watermarking by computing per-column tags from the pipeline config.
165 pub fn enable_watermark(&mut self, domain: &str, size: usize, seed: u64) {
166 use sha2::{Digest, Sha256};
167 for &tag in &[
168 0x53534e, // "SSN"
169 0x50484f4e, // "PHONE"
170 0x50414e, // "PAN"
171 0x4d4544, // "MEDICARE"
172 0x4f4643, // "OFFICE_PHONE"
173 0x504153, // "PASSPORT"
174 0x414354, // "ACCOUNT"
175 0x424152, // "BARCODE"
176 0x494343, // "ICCID"
177 0x555043, // "UPC"
178 ] {
179 let input = format!(
180 "{}{}{}{}{}",
181 Self::WATERMARK_SECRET,
182 domain,
183 size,
184 seed,
185 tag
186 );
187 let hash = Sha256::digest(input.as_bytes());
188 let wm = u64::from_le_bytes(hash[..8].try_into().unwrap());
189 self.watermark_map.insert(tag, wm);
190 }
191 }
192
193 /// Return the watermark mask for a given column tag (last 3 digits, 0..999).
194 pub fn watermark_3digits(&self, tag: u64) -> u64 {
195 self.watermark_map.get(&tag).copied().unwrap_or(0) % 1000
196 }
197
198 /// Return the watermark mask for a given column tag (last 2 digits, 0..99).
199 pub fn watermark_2digits(&self, tag: u64) -> u64 {
200 self.watermark_map.get(&tag).copied().unwrap_or(0) % 100
201 }
202
203 /// Create a minimal context for testing (no pools loaded, watermark disabled).
204 /// Watermark helpers return 0, ensuring deterministic test output.
205 #[cfg(test)]
206 pub fn test() -> Self {
207 Self {
208 pool_store: PoolStore {
209 pools: HashMap::new(),
210 },
211 locale: "en".to_string(),
212 watermark_map: HashMap::new(),
213 }
214 }
215}