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

quantik_core/bench/
dataset.rs
