Skip to main content

quantik_core/bench/
contracts.rs

1//! Contract-oriented projections for benchmark artifacts.
2//!
3//! The benchmark harness keeps its rich JSON bundle shape for reports and
4//! checkpointing. This module projects those rows into the cross-repository
5//! contracts used by training and analytics pipelines.
6
7use crate::bench::canonical::canonical_json;
8use crate::bench::reference::parse_move_key;
9use crate::bitboard::Bitboard;
10use crate::game::current_player;
11use crate::moves::generate_legal_moves;
12use crate::state::State;
13use serde_json::{json, Value};
14use std::collections::BTreeMap;
15use std::io::Write;
16use std::path::Path;
17
18pub const CONTRACT_VERSION: &str = "1.0.0";
19pub const OPENING_BOOK_SCHEMA: &str = "opening-book.v1";
20pub const OBSERVATION_SCHEMA: &str = "observation.v1";
21pub const GAME_RESULT_SCHEMA: &str = "game-result.v1";
22pub const MODEL_CHECKPOINT_SCHEMA: &str = "model-checkpoint.v1";
23
24pub fn action_index(shape: u8, position: u8) -> u8 {
25    shape * 16 + position
26}
27
28pub fn legal_action_mask(bb: &Bitboard) -> u64 {
29    generate_legal_moves(bb).into_iter().fold(0u64, |mask, mv| {
30        mask | (1u64 << action_index(mv.shape, mv.position))
31    })
32}
33
34pub fn canonical_key_hex(state: &State) -> String {
35    state
36        .canonical_key()
37        .iter()
38        .map(|byte| format!("{byte:02x}"))
39        .collect()
40}
41
42pub fn position_lookup(dataset_payload: &Value) -> Result<BTreeMap<String, Value>, String> {
43    let positions = dataset_payload["positions"]
44        .as_array()
45        .ok_or("dataset has no positions")?;
46    let mut by_id = BTreeMap::new();
47    for position in positions {
48        let id = position["id"]
49            .as_str()
50            .ok_or("dataset position missing id")?
51            .to_string();
52        by_id.insert(id, position.clone());
53    }
54    Ok(by_id)
55}
56
57pub fn observation_v1_row(
58    row_id: u64,
59    run_id: &str,
60    benchmark_row: &Value,
61    position: &Value,
62) -> Result<Value, String> {
63    let qfen = position["qfen"]
64        .as_str()
65        .ok_or("dataset position missing qfen")?;
66    let state = State::from_qfen(qfen)?;
67    let bb = state.bb;
68    let (_, shape, position_index) = parse_move_key(
69        benchmark_row["move"]
70            .as_str()
71            .ok_or("benchmark observation missing move")?,
72    )?;
73    let selected_action = action_index(shape, position_index) as usize;
74    let mut policy_visits = vec![0u32; 64];
75    policy_visits[selected_action] = 1;
76    let mut root_q_values = vec![0.0f64; 64];
77    if let Some(score) = benchmark_row["score"].as_f64() {
78        root_q_values[selected_action] = score;
79    }
80
81    let source_confidence = if benchmark_row["exact"].as_bool() == Some(true) {
82        1.0
83    } else {
84        0.5
85    };
86
87    Ok(json!({
88        "schema": OBSERVATION_SCHEMA,
89        "contract_version": CONTRACT_VERSION,
90        "run_id": run_id,
91        "row_id": row_id,
92        "position_key": canonical_key_hex(&state),
93        "ply": bb.player_piece_count(0) + bb.player_piece_count(1),
94        "side_to_move": current_player(&bb).ok_or("inconsistent side to move")?,
95        "bitboards": bb.planes,
96        "qfen": qfen,
97        "legal_action_mask": legal_action_mask(&bb),
98        "engine_kind": benchmark_row["engine"].as_str().unwrap_or("unknown"),
99        "engine_checkpoint": Value::Null,
100        "engine_version": env!("CARGO_PKG_VERSION"),
101        "search_depth": benchmark_row["depth_reached"],
102        "rollouts": benchmark_row["iterations"],
103        "beam_width": Value::Null,
104        "node_budget": benchmark_row["nodes"],
105        "time_budget_ms": Value::Null,
106        "elapsed_ms": (benchmark_row["wall_time_s"].as_f64().unwrap_or(0.0) * 1000.0).round() as u64,
107        "seed": benchmark_row["seed"],
108        "policy_visits": policy_visits,
109        "policy_priors": Value::Null,
110        "root_q_values": root_q_values,
111        "value": benchmark_row["score"].as_f64().unwrap_or(0.0),
112        "value_source": if benchmark_row["exact"].as_bool() == Some(true) { "exact" } else { "heuristic" },
113        "source_confidence": source_confidence,
114        "principal_variation": Value::Null,
115    }))
116}
117
118pub fn game_result_v1_row(
119    row_id: u64,
120    run_id: &str,
121    started_at: &str,
122    h2h_record: &Value,
123    position: &Value,
124) -> Result<Value, String> {
125    let qfen = position["qfen"]
126        .as_str()
127        .ok_or("dataset position missing qfen")?;
128    let state = State::from_qfen(qfen)?;
129    let bb = state.bb;
130    let side_to_move = current_player(&bb).ok_or("inconsistent side to move")?;
131    let mover = h2h_record["mover"].as_str().ok_or("h2h missing mover")?;
132    let responder = h2h_record["responder"]
133        .as_str()
134        .ok_or("h2h missing responder")?;
135    let (p0_engine, p1_engine) = if side_to_move == 0 {
136        (mover, responder)
137    } else {
138        (responder, mover)
139    };
140    let winner_engine = h2h_record["winner"].as_str().ok_or("h2h missing winner")?;
141    let winner = if winner_engine == p0_engine { 0 } else { 1 };
142    let move_action_indices = h2h_record["move_action_indices"]
143        .as_array()
144        .cloned()
145        .unwrap_or_default();
146
147    Ok(json!({
148        "schema": GAME_RESULT_SCHEMA,
149        "contract_version": CONTRACT_VERSION,
150        "game_id": format!("{run_id}-{row_id:08}"),
151        "started_at": started_at,
152        "p0_engine_kind": p0_engine,
153        "p0_engine_version": env!("CARGO_PKG_VERSION"),
154        "p0_engine_checkpoint": Value::Null,
155        "p1_engine_kind": p1_engine,
156        "p1_engine_version": env!("CARGO_PKG_VERSION"),
157        "p1_engine_checkpoint": Value::Null,
158        "opening_book_id": Value::Null,
159        "initial_position_key": canonical_key_hex(&state),
160        "winner": winner,
161        "plies": h2h_record["plies"],
162        "terminal_reason": "win_condition_or_no_legal_moves",
163        "seed": h2h_record["seed"],
164        "time_budget_ms_per_move": Value::Null,
165        "node_budget_per_move": Value::Null,
166        "move_action_indices": move_action_indices,
167        "position_keys": Value::Null,
168        "hardware": Value::Null,
169        "run_id": run_id,
170    }))
171}
172
173pub fn export_observation_rows(
174    bundle: &Value,
175    dataset_payload: &Value,
176    output: &Path,
177) -> Result<usize, String> {
178    let positions = position_lookup(dataset_payload)?;
179    let run_id = bundle["started_at"].as_str().unwrap_or("benchmark-run");
180    let rows = bundle["observations"]
181        .as_array()
182        .ok_or("bundle has no observations array")?;
183    write_jsonl(
184        output,
185        rows.iter().enumerate().map(|(index, row)| {
186            let position_id = row["position_id"]
187                .as_str()
188                .ok_or("observation missing position_id")?;
189            let position = positions
190                .get(position_id)
191                .ok_or_else(|| format!("dataset missing position {position_id}"))?;
192            observation_v1_row(index as u64, run_id, row, position)
193        }),
194    )
195}
196
197pub fn export_game_result_rows(
198    bundle: &Value,
199    dataset_payload: &Value,
200    output: &Path,
201) -> Result<usize, String> {
202    let positions = position_lookup(dataset_payload)?;
203    let run_id = bundle["started_at"].as_str().unwrap_or("benchmark-run");
204    let started_at = bundle["started_at"].as_str().unwrap_or("unknown");
205    let rows = bundle["head_to_head"]["records"]
206        .as_array()
207        .ok_or("bundle has no head_to_head.records array")?;
208    write_jsonl(
209        output,
210        rows.iter().enumerate().map(|(index, row)| {
211            let position_id = row["position_id"]
212                .as_str()
213                .ok_or("h2h record missing position_id")?;
214            let position = positions
215                .get(position_id)
216                .ok_or_else(|| format!("dataset missing position {position_id}"))?;
217            game_result_v1_row(index as u64, run_id, started_at, row, position)
218        }),
219    )
220}
221
222fn write_jsonl<I>(output: &Path, rows: I) -> Result<usize, String>
223where
224    I: Iterator<Item = Result<Value, String>>,
225{
226    if let Some(parent) = output.parent() {
227        std::fs::create_dir_all(parent).map_err(|e| format!("mkdir {parent:?}: {e}"))?;
228    }
229    let mut file = std::fs::File::create(output).map_err(|e| format!("create {output:?}: {e}"))?;
230    let mut count = 0usize;
231    for row in rows {
232        writeln!(file, "{}", canonical_json(&row?)).map_err(|e| format!("write: {e}"))?;
233        count += 1;
234    }
235    Ok(count)
236}
237
238#[cfg(test)]
239mod tests {
240    use super::*;
241
242    fn dataset() -> Value {
243        json!({
244            "positions": [{
245                "id": "p0000",
246                "qfen": "..../..../..../....",
247                "phase": "opening"
248            }]
249        })
250    }
251
252    #[test]
253    fn legal_mask_empty_board_covers_all_actions() {
254        assert_eq!(legal_action_mask(&Bitboard::EMPTY), u64::MAX);
255    }
256
257    #[test]
258    fn observation_projection_has_contract_shape() {
259        let positions = position_lookup(&dataset()).unwrap();
260        let row = json!({
261            "position_id": "p0000",
262            "engine": "minimax",
263            "move": "0:1:2",
264            "wall_time_s": 0.25,
265            "exact": true,
266            "seed": null,
267            "nodes": 10,
268            "iterations": null,
269            "depth_reached": 4,
270            "score": 1.0
271        });
272        let projected = observation_v1_row(0, "run", &row, &positions["p0000"]).unwrap();
273        assert_eq!(projected["schema"], OBSERVATION_SCHEMA);
274        assert_eq!(projected["contract_version"], CONTRACT_VERSION);
275        assert_eq!(projected["policy_visits"][18], json!(1));
276        assert_eq!(projected["legal_action_mask"], json!(u64::MAX));
277        assert_eq!(projected["value_source"], json!("exact"));
278    }
279
280    #[test]
281    fn game_result_projection_maps_engines_to_players() {
282        let positions = position_lookup(&dataset()).unwrap();
283        let record = json!({
284            "position_id": "p0000",
285            "mover": "mcts",
286            "responder": "minimax",
287            "winner": "minimax",
288            "plies": 3,
289            "seed": 7,
290            "move_action_indices": [0, 17, 2]
291        });
292        let projected = game_result_v1_row(
293            0,
294            "run",
295            "2026-07-14T00:00:00+0200",
296            &record,
297            &positions["p0000"],
298        )
299        .unwrap();
300        assert_eq!(projected["schema"], GAME_RESULT_SCHEMA);
301        assert_eq!(projected["p0_engine_kind"], json!("mcts"));
302        assert_eq!(projected["p1_engine_kind"], json!("minimax"));
303        assert_eq!(projected["winner"], json!(1));
304        assert_eq!(
305            projected["move_action_indices"].as_array().unwrap().len(),
306            3
307        );
308    }
309}