1use 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}