quantik_core/
opening_book.rs

1use crate::state::State;
2use rusqlite::{params, Connection, Result as SqlResult};
3
4#[derive(Clone, Copy, Debug, PartialEq, Eq)]
5pub enum TerminalStatus {
6    Interior = 0,
7    WinP0 = 1,
8    WinP1 = 2,
9    Stalemate = 3,
10}
11
12impl TerminalStatus {
13    pub fn from_i32(v: i32) -> Self {
14        match v {
15            1 => Self::WinP0,
16            2 => Self::WinP1,
17            3 => Self::Stalemate,
18            _ => Self::Interior,
19        }
20    }
21}
22
23#[derive(Clone, Debug)]
24pub struct OpeningBookEntry {
25    pub canonical_key: Vec<u8>,
26    pub qfen: String,
27    pub depth: i32,
28    pub evaluation: f64,
29    pub visit_count: i64,
30    pub win_count_p0: i64,
31    pub win_count_p1: i64,
32    pub draw_count: i64,
33    pub best_moves: Vec<(i32, i32)>, // (shape, position)
34    pub is_terminal: TerminalStatus,
35    pub symmetry_count: i32,
36    /// Whether this position carries an exactly-solved game value (see
37    /// [`OpeningBookDatabase::add_solved_position`]). `false`/default for
38    /// rows written before the `solved`/`game_value` columns existed.
39    pub solved: bool,
40    /// Exact game value for the side to move (+1/-1), when `solved`.
41    pub game_value: Option<i32>,
42}
43
44pub struct OpeningBookConfig {
45    pub database_path: String,
46    pub cache_size_mb: i32,
47    pub enable_wal: bool,
48}
49
50impl Default for OpeningBookConfig {
51    fn default() -> Self {
52        Self {
53            database_path: "quantik_opening_book.db".into(),
54            cache_size_mb: 100,
55            enable_wal: true,
56        }
57    }
58}
59
60pub struct OpeningBookDatabase {
61    conn: Connection,
62    db_path: String,
63}
64
65impl OpeningBookDatabase {
66    pub fn open(config: &OpeningBookConfig) -> SqlResult<Self> {
67        let conn = Connection::open(&config.database_path)?;
68
69        conn.execute_batch(&format!(
70            "PRAGMA cache_size = -{};",
71            config.cache_size_mb * 1024
72        ))?;
73        if config.enable_wal {
74            conn.execute_batch("PRAGMA journal_mode = WAL;")?;
75        }
76        conn.execute_batch("PRAGMA synchronous = NORMAL;")?;
77
78        conn.execute_batch(
79            "CREATE TABLE IF NOT EXISTS positions (
80                canonical_key BLOB PRIMARY KEY,
81                qfen TEXT NOT NULL,
82                depth INTEGER NOT NULL,
83                evaluation REAL NOT NULL,
84                visit_count INTEGER NOT NULL,
85                win_count_p0 INTEGER NOT NULL,
86                win_count_p1 INTEGER NOT NULL,
87                draw_count INTEGER NOT NULL,
88                is_terminal INTEGER NOT NULL DEFAULT 0,
89                symmetry_count INTEGER NOT NULL DEFAULT 0,
90                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
91            );
92            CREATE TABLE IF NOT EXISTS best_moves (
93                canonical_key BLOB NOT NULL,
94                move_rank INTEGER NOT NULL,
95                shape INTEGER NOT NULL,
96                position INTEGER NOT NULL,
97                FOREIGN KEY (canonical_key) REFERENCES positions(canonical_key),
98                PRIMARY KEY (canonical_key, move_rank)
99            );
100            CREATE TABLE IF NOT EXISTS position_edges (
101                parent_key BLOB NOT NULL,
102                child_key  BLOB NOT NULL,
103                PRIMARY KEY (parent_key, child_key),
104                FOREIGN KEY (parent_key) REFERENCES positions(canonical_key),
105                FOREIGN KEY (child_key)  REFERENCES positions(canonical_key)
106            );
107            CREATE INDEX IF NOT EXISTS idx_depth ON positions(depth);
108            CREATE INDEX IF NOT EXISTS idx_visit_count ON positions(visit_count DESC);
109            CREATE INDEX IF NOT EXISTS idx_edges_child ON position_edges(child_key);",
110        )?;
111
112        // Idempotent migration for DBs created before the `solved`/
113        // `game_value` columns existed: SQLite has no `ADD COLUMN IF NOT
114        // EXISTS`, so attempt the ALTER and swallow the "duplicate column
115        // name" error it raises when the column is already there.
116        for stmt in [
117            "ALTER TABLE positions ADD COLUMN solved INTEGER NOT NULL DEFAULT 0",
118            "ALTER TABLE positions ADD COLUMN game_value INTEGER",
119        ] {
120            if let Err(e) = conn.execute_batch(stmt) {
121                if !e.to_string().contains("duplicate column name") {
122                    return Err(e);
123                }
124            }
125        }
126
127        Ok(Self {
128            conn,
129            db_path: config.database_path.clone(),
130        })
131    }
132
133    #[allow(clippy::too_many_arguments)]
134    pub fn add_position(
135        &self,
136        state: &State,
137        evaluation: f64,
138        visit_count: i64,
139        win_count_p0: i64,
140        win_count_p1: i64,
141        draw_count: i64,
142        best_moves: &[(i32, i32)],
143        depth: i32,
144        is_terminal: TerminalStatus,
145        symmetry_count: i32,
146    ) -> SqlResult<()> {
147        let canonical_key = state.canonical_key().to_vec();
148        let qfen = state.to_qfen();
149
150        self.conn.execute(
151            "INSERT OR REPLACE INTO positions
152             (canonical_key, qfen, depth, evaluation, visit_count,
153              win_count_p0, win_count_p1, draw_count, is_terminal, symmetry_count)
154             VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)",
155            params![
156                canonical_key,
157                qfen,
158                depth,
159                evaluation,
160                visit_count,
161                win_count_p0,
162                win_count_p1,
163                draw_count,
164                is_terminal as i32,
165                symmetry_count,
166            ],
167        )?;
168
169        self.conn.execute(
170            "DELETE FROM best_moves WHERE canonical_key = ?1",
171            params![canonical_key],
172        )?;
173
174        for (rank, &(shape, position)) in best_moves.iter().take(5).enumerate() {
175            self.conn.execute(
176                "INSERT INTO best_moves (canonical_key, move_rank, shape, position)
177                 VALUES (?1, ?2, ?3, ?4)",
178                params![canonical_key, (rank + 1) as i32, shape, position],
179            )?;
180        }
181        Ok(())
182    }
183
184    /// Upsert an exactly-solved position: `evaluation` is `game_value` as
185    /// `f64`, `is_terminal` stays `Interior` (the position itself is not
186    /// terminal — its *game value* is exactly known), `depth` is pieces
187    /// placed (derived from `state.bb`), and `best_moves` records every
188    /// optimal move (not just a top-5 slice, unlike [`Self::add_position`])
189    /// as `(shape, position)` pairs in the given order.
190    ///
191    /// **Only canonical representatives are stored.** The optimal moves
192    /// are expressed in `state`'s own board orientation, but the row is
193    /// keyed by the canonical key shared by up to eight symmetric
194    /// orientations; storing a non-representative orientation would let a
195    /// later lookup on the *representative* serve moves that are wrong
196    /// (possibly illegal) for that board. If `state` is not its own
197    /// canonical representative (`canonical_payload() != bb.to_le_bytes()`),
198    /// nothing is written and `Ok(false)` is returned; `Ok(true)` means
199    /// the row was upserted. Translating moves across orientations via
200    /// the symmetry transform is a documented follow-up that would lift
201    /// this restriction.
202    ///
203    /// The position row and its best-move rows are written in one
204    /// transaction, so a mid-way failure can never leave a solved row
205    /// with partial `best_moves`.
206    ///
207    /// Visit/win/draw counters and `symmetry_count` are not meaningful for
208    /// a solved reference and are stored as `0`.
209    pub fn add_solved_position(
210        &self,
211        state: &State,
212        game_value: i32,
213        optimal_moves: &[(i32, i32)],
214    ) -> SqlResult<bool> {
215        if state.canonical_payload() != state.bb.to_le_bytes() {
216            return Ok(false);
217        }
218
219        let canonical_key = state.canonical_key().to_vec();
220        let qfen = state.to_qfen();
221        let depth = (state.bb.player_piece_count(0) + state.bb.player_piece_count(1)) as i32;
222        let evaluation = game_value as f64;
223
224        let tx = self.conn.unchecked_transaction()?;
225        tx.execute(
226            "INSERT OR REPLACE INTO positions
227             (canonical_key, qfen, depth, evaluation, visit_count,
228              win_count_p0, win_count_p1, draw_count, is_terminal, symmetry_count,
229              solved, game_value)
230             VALUES (?1, ?2, ?3, ?4, 0, 0, 0, 0, ?5, 0, 1, ?6)",
231            params![
232                canonical_key,
233                qfen,
234                depth,
235                evaluation,
236                TerminalStatus::Interior as i32,
237                game_value,
238            ],
239        )?;
240
241        tx.execute(
242            "DELETE FROM best_moves WHERE canonical_key = ?1",
243            params![canonical_key],
244        )?;
245
246        for (rank, &(shape, position)) in optimal_moves.iter().enumerate() {
247            tx.execute(
248                "INSERT INTO best_moves (canonical_key, move_rank, shape, position)
249                 VALUES (?1, ?2, ?3, ?4)",
250                params![canonical_key, (rank + 1) as i32, shape, position],
251            )?;
252        }
253        tx.commit()?;
254        Ok(true)
255    }
256
257    pub fn get_position(&self, state: &State) -> SqlResult<Option<OpeningBookEntry>> {
258        let canonical_key = state.canonical_key().to_vec();
259
260        let mut stmt = self.conn.prepare(
261            "SELECT qfen, depth, evaluation, visit_count,
262                    win_count_p0, win_count_p1, draw_count,
263                    is_terminal, symmetry_count, solved, game_value
264             FROM positions WHERE canonical_key = ?1",
265        )?;
266
267        let entry = stmt.query_row(params![canonical_key], |row| {
268            Ok(OpeningBookEntry {
269                canonical_key: canonical_key.clone(),
270                qfen: row.get(0)?,
271                depth: row.get(1)?,
272                evaluation: row.get(2)?,
273                visit_count: row.get(3)?,
274                win_count_p0: row.get(4)?,
275                win_count_p1: row.get(5)?,
276                draw_count: row.get(6)?,
277                best_moves: Vec::new(), // filled below
278                is_terminal: TerminalStatus::from_i32(row.get(7)?),
279                symmetry_count: row.get(8)?,
280                solved: row.get::<_, i64>(9)? != 0,
281                game_value: row.get(10)?,
282            })
283        });
284
285        match entry {
286            Ok(mut e) => {
287                let mut mv_stmt = self.conn.prepare(
288                    "SELECT shape, position FROM best_moves
289                     WHERE canonical_key = ?1 ORDER BY move_rank",
290                )?;
291                let moves = mv_stmt.query_map(params![e.canonical_key], |row| {
292                    Ok((row.get::<_, i32>(0)?, row.get::<_, i32>(1)?))
293                })?;
294                e.best_moves = moves.filter_map(|r| r.ok()).collect();
295                Ok(Some(e))
296            }
297            Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
298            Err(e) => Err(e),
299        }
300    }
301
302    pub fn query_by_depth(&self, depth: i32, limit: i64) -> SqlResult<Vec<OpeningBookEntry>> {
303        let mut stmt = self.conn.prepare(
304            "SELECT canonical_key, qfen, depth, evaluation, visit_count,
305                    win_count_p0, win_count_p1, draw_count,
306                    is_terminal, symmetry_count, solved, game_value
307             FROM positions WHERE depth = ?1
308             ORDER BY visit_count DESC LIMIT ?2",
309        )?;
310
311        let entries: Vec<OpeningBookEntry> = stmt
312            .query_map(params![depth, limit], |row| {
313                Ok(OpeningBookEntry {
314                    canonical_key: row.get(0)?,
315                    qfen: row.get(1)?,
316                    depth: row.get(2)?,
317                    evaluation: row.get(3)?,
318                    visit_count: row.get(4)?,
319                    win_count_p0: row.get(5)?,
320                    win_count_p1: row.get(6)?,
321                    draw_count: row.get(7)?,
322                    best_moves: Vec::new(),
323                    is_terminal: TerminalStatus::from_i32(row.get(8)?),
324                    symmetry_count: row.get(9)?,
325                    solved: row.get::<_, i64>(10)? != 0,
326                    game_value: row.get(11)?,
327                })
328            })?
329            .filter_map(|r| r.ok())
330            .collect();
331
332        Ok(entries)
333    }
334
335    // ── DAG edges ────────────────────────────────────────────────────
336
337    pub fn add_edges(&self, edges: &[(&[u8], &[u8])]) -> SqlResult<()> {
338        let tx = self.conn.unchecked_transaction()?;
339        {
340            let mut stmt = tx.prepare(
341                "INSERT OR IGNORE INTO position_edges (parent_key, child_key)
342                 VALUES (?1, ?2)",
343            )?;
344            for &(parent, child) in edges {
345                stmt.execute(params![parent, child])?;
346            }
347        }
348        tx.commit()
349    }
350
351    pub fn get_children(&self, canonical_key: &[u8]) -> SqlResult<Vec<Vec<u8>>> {
352        let mut stmt = self
353            .conn
354            .prepare("SELECT child_key FROM position_edges WHERE parent_key = ?1")?;
355        let rows = stmt.query_map(params![canonical_key], |row| row.get(0))?;
356        rows.collect()
357    }
358
359    pub fn get_parents(&self, canonical_key: &[u8]) -> SqlResult<Vec<Vec<u8>>> {
360        let mut stmt = self
361            .conn
362            .prepare("SELECT parent_key FROM position_edges WHERE child_key = ?1")?;
363        let rows = stmt.query_map(params![canonical_key], |row| row.get(0))?;
364        rows.collect()
365    }
366
367    pub fn get_edge_count(&self) -> SqlResult<i64> {
368        self.conn
369            .query_row("SELECT COUNT(*) FROM position_edges", [], |row| row.get(0))
370    }
371
372    // ── statistics ───────────────────────────────────────────────────
373
374    pub fn total_positions(&self) -> SqlResult<i64> {
375        self.conn
376            .query_row("SELECT COUNT(*) FROM positions", [], |row| row.get(0))
377    }
378
379    pub fn max_depth(&self) -> SqlResult<Option<i32>> {
380        self.conn
381            .query_row("SELECT MAX(depth) FROM positions", [], |row| row.get(0))
382    }
383
384    pub fn positions_by_depth(&self) -> SqlResult<Vec<(i32, i64)>> {
385        let mut stmt = self
386            .conn
387            .prepare("SELECT depth, COUNT(*) FROM positions GROUP BY depth ORDER BY depth")?;
388        let rows = stmt.query_map([], |row| Ok((row.get(0)?, row.get(1)?)))?;
389        rows.collect()
390    }
391
392    pub fn db_path(&self) -> &str {
393        &self.db_path
394    }
395
396    pub fn file_size(&self) -> u64 {
397        std::fs::metadata(&self.db_path)
398            .map(|m| m.len())
399            .unwrap_or(0)
400    }
401}
402
403impl Drop for OpeningBookDatabase {
404    fn drop(&mut self) {
405        // Connection is dropped automatically
406    }
407}
408
409#[cfg(test)]
410mod tests {
411    use super::*;
412    use std::fs;
413
414    fn temp_db_path() -> String {
415        let id = std::time::SystemTime::now()
416            .duration_since(std::time::UNIX_EPOCH)
417            .unwrap()
418            .as_nanos();
419        format!("/tmp/quantik_test_{}.db", id)
420    }
421
422    #[test]
423    fn open_and_add_position() {
424        let path = temp_db_path();
425        let config = OpeningBookConfig {
426            database_path: path.clone(),
427            ..Default::default()
428        };
429        let db = OpeningBookDatabase::open(&config).unwrap();
430
431        let state = State::empty();
432        db.add_position(
433            &state,
434            0.0,
435            100,
436            50,
437            40,
438            10,
439            &[(0, 0), (1, 5)],
440            0,
441            TerminalStatus::Interior,
442            1,
443        )
444        .unwrap();
445
446        let entry = db.get_position(&state).unwrap().unwrap();
447        assert_eq!(entry.visit_count, 100);
448        assert_eq!(entry.best_moves.len(), 2);
449
450        assert_eq!(db.total_positions().unwrap(), 1);
451
452        fs::remove_file(&path).ok();
453    }
454
455    #[test]
456    fn edge_operations() {
457        let path = temp_db_path();
458        let config = OpeningBookConfig {
459            database_path: path.clone(),
460            ..Default::default()
461        };
462        let db = OpeningBookDatabase::open(&config).unwrap();
463
464        let s1 = State::empty();
465        let s2 = State::from_qfen("A.../..../..../....").unwrap();
466
467        db.add_position(&s1, 0.0, 1, 0, 0, 0, &[], 0, TerminalStatus::Interior, 1)
468            .unwrap();
469        db.add_position(&s2, 0.1, 1, 0, 0, 0, &[], 1, TerminalStatus::Interior, 1)
470            .unwrap();
471
472        let k1 = s1.canonical_key();
473        let k2 = s2.canonical_key();
474        db.add_edges(&[(&k1, &k2)]).unwrap();
475
476        let children = db.get_children(&k1).unwrap();
477        assert_eq!(children.len(), 1);
478        assert_eq!(children[0], k2.to_vec());
479
480        let parents = db.get_parents(&k2).unwrap();
481        assert_eq!(parents.len(), 1);
482
483        assert_eq!(db.get_edge_count().unwrap(), 1);
484
485        fs::remove_file(&path).ok();
486    }
487
488    #[test]
489    fn query_by_depth_works() {
490        let path = temp_db_path();
491        let config = OpeningBookConfig {
492            database_path: path.clone(),
493            ..Default::default()
494        };
495        let db = OpeningBookDatabase::open(&config).unwrap();
496
497        let state = State::empty();
498        db.add_position(
499            &state,
500            0.0,
501            10,
502            5,
503            3,
504            2,
505            &[],
506            0,
507            TerminalStatus::Interior,
508            1,
509        )
510        .unwrap();
511
512        let entries = db.query_by_depth(0, 10).unwrap();
513        assert_eq!(entries.len(), 1);
514        assert_eq!(entries[0].depth, 0);
515
516        let entries = db.query_by_depth(1, 10).unwrap();
517        assert!(entries.is_empty());
518
519        fs::remove_file(&path).ok();
520    }
521
522    #[test]
523    fn add_solved_position_roundtrips() {
524        let path = temp_db_path();
525        let config = OpeningBookConfig {
526            database_path: path.clone(),
527            ..Default::default()
528        };
529        let db = OpeningBookDatabase::open(&config).unwrap();
530
531        // The empty board is trivially its own canonical representative.
532        let state = State::empty();
533        let written = db
534            .add_solved_position(&state, 1, &[(0, 0), (1, 5)])
535            .unwrap();
536        assert!(written);
537
538        let entry = db.get_position(&state).unwrap().unwrap();
539        assert!(entry.solved);
540        assert_eq!(entry.game_value, Some(1));
541        assert_eq!(entry.evaluation, 1.0);
542        assert_eq!(entry.is_terminal, TerminalStatus::Interior);
543        assert_eq!(entry.depth, 0);
544        assert_eq!(entry.best_moves, vec![(0, 0), (1, 5)]);
545
546        fs::remove_file(&path).ok();
547    }
548
549    /// Defense-in-depth guard: `add_solved_position` refuses (returns
550    /// `Ok(false)`, writes nothing) when the state is not its own
551    /// canonical representative — its moves would be stored in the wrong
552    /// orientation for the canonical key.
553    #[test]
554    fn add_solved_position_skips_non_canonical_orientations() {
555        let path = temp_db_path();
556        let config = OpeningBookConfig {
557            database_path: path.clone(),
558            ..Default::default()
559        };
560        let db = OpeningBookDatabase::open(&config).unwrap();
561
562        // Find a single-piece placement that is NOT its own canonical
563        // representative (most of the 16 cells aren't).
564        let state = (0..16)
565            .map(|pos| {
566                let mut qfen: Vec<char> = "..../..../..../....".chars().collect();
567                let index = pos + pos / 4; // account for '/' separators
568                qfen[index] = 'A';
569                State::from_qfen(&qfen.into_iter().collect::<String>()).unwrap()
570            })
571            .find(|s| s.canonical_payload() != s.bb.to_le_bytes())
572            .expect("some single-piece placement is non-canonical");
573
574        let written = db.add_solved_position(&state, 1, &[(0, 0)]).unwrap();
575        assert!(!written);
576        assert_eq!(db.total_positions().unwrap(), 0);
577        assert!(db.get_position(&state).unwrap().is_none());
578
579        fs::remove_file(&path).ok();
580    }
581
582    /// Opening a pre-existing DB created with the OLD schema (no
583    /// `solved`/`game_value` columns) must upgrade it in place: `open()`
584    /// succeeds and `get_position` returns the migrated defaults
585    /// (`solved: false`, `game_value: None`) for rows written before the
586    /// migration existed.
587    #[test]
588    fn migration_upgrades_pre_existing_db() {
589        let path = temp_db_path();
590
591        // Build the OLD schema by hand (mirrors the CREATE TABLE that
592        // predates the solved/game_value columns) and insert a row via raw
593        // SQL, bypassing OpeningBookDatabase entirely.
594        {
595            let conn = Connection::open(&path).unwrap();
596            conn.execute_batch(
597                "CREATE TABLE positions (
598                    canonical_key BLOB PRIMARY KEY,
599                    qfen TEXT NOT NULL,
600                    depth INTEGER NOT NULL,
601                    evaluation REAL NOT NULL,
602                    visit_count INTEGER NOT NULL,
603                    win_count_p0 INTEGER NOT NULL,
604                    win_count_p1 INTEGER NOT NULL,
605                    draw_count INTEGER NOT NULL,
606                    is_terminal INTEGER NOT NULL DEFAULT 0,
607                    symmetry_count INTEGER NOT NULL DEFAULT 0,
608                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
609                );
610                CREATE TABLE best_moves (
611                    canonical_key BLOB NOT NULL,
612                    move_rank INTEGER NOT NULL,
613                    shape INTEGER NOT NULL,
614                    position INTEGER NOT NULL,
615                    PRIMARY KEY (canonical_key, move_rank)
616                );",
617            )
618            .unwrap();
619
620            let state = State::empty();
621            let canonical_key = state.canonical_key().to_vec();
622            conn.execute(
623                "INSERT INTO positions
624                 (canonical_key, qfen, depth, evaluation, visit_count,
625                  win_count_p0, win_count_p1, draw_count, is_terminal, symmetry_count)
626                 VALUES (?1, ?2, 0, 0.5, 7, 3, 2, 1, 0, 1)",
627                params![canonical_key, state.to_qfen()],
628            )
629            .unwrap();
630        }
631
632        // Opening via OpeningBookDatabase must succeed (idempotent ALTER
633        // TABLE) and see the pre-existing row with migrated defaults.
634        let config = OpeningBookConfig {
635            database_path: path.clone(),
636            ..Default::default()
637        };
638        let db = OpeningBookDatabase::open(&config).unwrap();
639        let entry = db.get_position(&State::empty()).unwrap().unwrap();
640        assert!(!entry.solved);
641        assert_eq!(entry.game_value, None);
642        assert_eq!(entry.visit_count, 7);
643
644        // Re-opening again (columns already present) must also succeed.
645        OpeningBookDatabase::open(&config).unwrap();
646
647        fs::remove_file(&path).ok();
648    }
649}
quantik_core/opening_book.rs

quantik_core/
opening_book.rs
