quantik_core/bench/
dataset.rs

1//! Shared, versioned position dataset for cross-engine benchmarks.
2//!
3//! Port of `benchmarks/dataset.py`. Artifacts are JSON with a sha256
4//! checksum over the canonical (sorted-key, compact-separator) encoding of
5//! the payload minus the `checksum` field — byte-compatible with the Python
6//! implementation, so datasets interoperate across languages. RNG streams
7//! differ from CPython, so `generate` with the same seed produces a
8//! *different but equally valid* dataset; the committed artifact is the
9//! shared one.
10
11use crate::game::{current_player, has_winning_line};
12use crate::moves::{apply_move, generate_legal_moves};
13use crate::state::State;
14use rand::prelude::*;
15use serde_json::{json, Map, Value};
16use sha2::{Digest, Sha256};
17use std::collections::BTreeMap;
18use std::path::Path;
19
20use super::contracts::CONTRACT_VERSION;
21
22pub const SCHEMA_VERSION: u64 = 1;
23pub const GENERATOR: &str = "benchmarks.dataset.generate/v1";
24
25/// Phase buckets by pieces placed (== plies from the empty board).
26pub const PHASES: [(&str, (u32, u32)); 4] = [
27    ("opening", (0, 4)),
28    ("early_mid", (5, 7)),
29    ("late_mid", (8, 11)),
30    ("endgame", (12, 16)),
31];
32
33/// Return the phase bucket for a piece count.
34pub fn phase_of(pieces: u32) -> Result<&'static str, String> {
35    for (phase, (lo, hi)) in PHASES {
36        if (lo..=hi).contains(&pieces) {
37            return Ok(phase);
38        }
39    }
40    Err(format!("no benchmark phase for {pieces} pieces"))
41}
42
43/// sha256 over canonical JSON excluding the `checksum` field.
44///
45/// The canonical encoding is byte-identical to Python's
46/// `json.dumps(sort_keys=True, separators=(",", ":"))` — see
47/// [`super::canonical::canonical_json`] for the float/string rules.
48pub fn checksum(payload: &Value) -> String {
49    let mut stripped = payload.as_object().cloned().unwrap_or_default();
50    stripped.remove("checksum");
51    let blob = super::canonical::canonical_json(&Value::Object(stripped));
52    let mut hasher = Sha256::new();
53    hasher.update(blob.as_bytes());
54    format!("{:x}", hasher.finalize())
55}
56
57/// Write a checksum-bearing JSON dataset artifact and return its checksum.
58pub fn save(payload: &Value, path: &Path) -> Result<String, String> {
59    let digest = checksum(payload);
60    let mut output = payload.as_object().cloned().unwrap_or_default();
61    output.insert("checksum".into(), Value::String(digest.clone()));
62    let text = serde_json::to_string_pretty(&Value::Object(output))
63        .map_err(|e| format!("serialize: {e}"))?;
64    if let Some(parent) = path.parent() {
65        std::fs::create_dir_all(parent).map_err(|e| format!("mkdir {parent:?}: {e}"))?;
66    }
67    std::fs::write(path, text + "\n").map_err(|e| format!("write {path:?}: {e}"))?;
68    Ok(digest)
69}
70
71/// Load a dataset artifact and verify its checksum.
72pub fn load(path: &Path) -> Result<Value, String> {
73    let text = std::fs::read_to_string(path).map_err(|e| format!("read {path:?}: {e}"))?;
74    let payload: Value = serde_json::from_str(&text).map_err(|e| format!("parse: {e}"))?;
75    let expected = payload
76        .get("checksum")
77        .and_then(Value::as_str)
78        .map(str::to_string);
79    let actual = checksum(&payload);
80    if expected.as_deref() != Some(actual.as_str()) {
81        return Err(format!(
82            "dataset checksum mismatch: expected {expected:?}, actual {actual}"
83        ));
84    }
85    Ok(payload)
86}
87
88/// Play `plies` random legal moves; `None` if the line hits a win or a
89/// dead end (including a final state with no legal continuation).
90fn random_position(rng: &mut StdRng, plies: u32) -> Option<crate::bitboard::Bitboard> {
91    let mut bb = crate::bitboard::Bitboard::EMPTY;
92    for _ in 0..plies {
93        let moves = generate_legal_moves(&bb);
94        if moves.is_empty() {
95            return None;
96        }
97        bb = apply_move(&bb, &moves[rng.gen_range(0..moves.len())]);
98        if has_winning_line(&bb) {
99            return None;
100        }
101    }
102    if generate_legal_moves(&bb).is_empty() {
103        return None;
104    }
105    Some(bb)
106}
107
108fn position_payload(position_id: usize, bb: &crate::bitboard::Bitboard, phase: &str) -> Value {
109    let pieces = bb.player_piece_count(0) + bb.player_piece_count(1);
110    json!({
111        "id": format!("p{position_id:04}"),
112        "qfen": State::new(*bb).to_qfen(),
113        "phase": phase,
114        "pieces": pieces,
115        "side_to_move": current_player(bb).expect("valid position"),
116        "legal_moves": generate_legal_moves(bb).len(),
117        "reference": Value::Null,
118    })
119}
120
121/// Generate a deterministic benchmark dataset for requested phase counts.
122pub fn generate(requested: &BTreeMap<String, u32>, seed: u64) -> Result<Value, String> {
123    let known: Vec<&str> = PHASES.iter().map(|(name, _)| *name).collect();
124    let unknown: Vec<&String> = requested
125        .keys()
126        .filter(|k| !known.contains(&k.as_str()))
127        .collect();
128    if !unknown.is_empty() {
129        return Err(format!("unknown phase(s): {unknown:?}"));
130    }
131
132    let mut rng = StdRng::seed_from_u64(seed);
133    let mut seen: std::collections::HashSet<[u8; 18]> = std::collections::HashSet::new();
134    let mut positions: Vec<Value> = Vec::new();
135
136    for (phase, (lo, hi)) in PHASES {
137        let want = *requested.get(phase).unwrap_or(&0);
138        let mut found = 0u32;
139        let mut attempts = 0u32;
140        let max_attempts = want * 500;
141
142        while found < want && attempts < max_attempts {
143            attempts += 1;
144            let target_plies = rng.gen_range(lo..=hi.min(15));
145            let Some(bb) = random_position(&mut rng, target_plies) else {
146                continue;
147            };
148            let key = State::new(bb).canonical_key();
149            if !seen.insert(key) {
150                continue;
151            }
152            positions.push(position_payload(positions.len(), &bb, phase));
153            found += 1;
154        }
155    }
156
157    let requested_map: Map<String, Value> = requested
158        .iter()
159        .map(|(k, v)| (k.clone(), json!(v)))
160        .collect();
161
162    Ok(json!({
163        "contract_version": CONTRACT_VERSION,
164        "schema_version": SCHEMA_VERSION,
165        "generator": GENERATOR,
166        "seed": seed,
167        "requested": requested_map,
168        "positions": positions,
169    }))
170}
171
172#[cfg(test)]
173mod tests {
174    use super::*;
175    use crate::game::check_winner;
176    use crate::game::WinStatus;
177
178    fn golden_path() -> std::path::PathBuf {
179        Path::new(env!("CARGO_MANIFEST_DIR")).join("../../benchmarks/positions-v1.json")
180    }
181
182    /// The linchpin of cross-language interop: the checksum recomputed by
183    /// the Rust canonical-JSON encoder over the Python-generated artifact
184    /// must equal the stored value (load() verifies it).
185    #[test]
186    fn golden_dataset_checksum_verifies() {
187        let payload = load(&golden_path()).unwrap();
188        assert_eq!(payload["schema_version"], json!(1));
189        assert_eq!(payload["generator"], json!(GENERATOR));
190        assert_eq!(payload["positions"].as_array().unwrap().len(), 36);
191    }
192
193    #[test]
194    fn golden_dataset_positions_are_valid() {
195        let payload = load(&golden_path()).unwrap();
196        for position in payload["positions"].as_array().unwrap() {
197            let qfen = position["qfen"].as_str().unwrap();
198            let state = State::from_qfen(qfen).unwrap();
199            let bb = state.bb;
200            assert_eq!(check_winner(&bb), WinStatus::NoWin, "{qfen} is terminal");
201            let legal = generate_legal_moves(&bb);
202            assert!(!legal.is_empty(), "{qfen} has no legal moves");
203            assert_eq!(
204                position["legal_moves"].as_u64().unwrap() as usize,
205                legal.len(),
206                "{qfen} legal move count"
207            );
208            let pieces = bb.player_piece_count(0) + bb.player_piece_count(1);
209            assert_eq!(position["pieces"].as_u64().unwrap() as u32, pieces);
210            assert_eq!(
211                position["side_to_move"].as_u64().unwrap() as u8,
212                current_player(&bb).unwrap(),
213                "{qfen} side to move"
214            );
215            assert_eq!(
216                position["phase"].as_str().unwrap(),
217                phase_of(pieces).unwrap()
218            );
219        }
220    }
221
222    #[test]
223    fn save_load_roundtrip_and_corruption_detection() {
224        let dir = std::env::temp_dir().join(format!("quantik-bench-{}", std::process::id()));
225        std::fs::create_dir_all(&dir).unwrap();
226        let path = dir.join("tiny.json");
227
228        let requested = BTreeMap::from([("opening".to_string(), 2u32)]);
229        let payload = generate(&requested, 7).unwrap();
230        let digest = save(&payload, &path).unwrap();
231        let loaded = load(&path).unwrap();
232        assert_eq!(loaded["checksum"].as_str().unwrap(), digest);
233
234        // Corrupt one byte inside the positions array.
235        let text = std::fs::read_to_string(&path).unwrap();
236        let corrupted = text.replacen("\"pieces\": 3", "\"pieces\": 4", 1);
237        let corrupted = if corrupted == text {
238            text.replacen("\"pieces\": 4", "\"pieces\": 5", 1)
239        } else {
240            corrupted
241        };
242        std::fs::write(&path, corrupted).unwrap();
243        assert!(load(&path).is_err());
244        std::fs::remove_file(&path).ok();
245    }
246
247    #[test]
248    fn generate_is_deterministic_and_deduped() {
249        let requested = BTreeMap::from([
250            ("opening".to_string(), 3u32),
251            ("early_mid".to_string(), 2u32),
252        ]);
253        let a = generate(&requested, 123).unwrap();
254        let b = generate(&requested, 123).unwrap();
255        assert_eq!(a, b);
256
257        let positions = a["positions"].as_array().unwrap();
258        assert_eq!(positions.len(), 5);
259        let mut keys = std::collections::HashSet::new();
260        for position in positions {
261            let state = State::from_qfen(position["qfen"].as_str().unwrap()).unwrap();
262            assert!(keys.insert(state.canonical_key()), "duplicate canonical");
263            let pieces = position["pieces"].as_u64().unwrap() as u32;
264            let phase = position["phase"].as_str().unwrap();
265            assert_eq!(phase_of(pieces).unwrap(), phase);
266        }
267    }
268
269    #[test]
270    fn unknown_phase_rejected() {
271        let requested = BTreeMap::from([("blitz".to_string(), 1u32)]);
272        assert!(generate(&requested, 1).is_err());
273    }
274}
quantik_core/bench/dataset.rs

quantik_core/bench/
dataset.rs
