Optimized opening db lookup

dded transposition table

Author: shashankp28

Cargo.toml

@@ -26,4 +26,5 @@ clap = { version = "*", features = ["derive"] }
 flate2 = { version = "1.0", features = ["rust_backend"] }
 once_cell = "1.21.3"
 tar = "0.4.44"
-rand = "*"
+rand = "*"
+lru = "*"

notebooks/benchmark.ipynb

@@ -1,8 +1,13 @@
 use axum::{ extract::{ State, Json }, response::IntoResponse, http::StatusCode };
 use serde::{ Deserialize, Serialize };
-use std::{ str::FromStr, collections::HashMap, sync::Arc };
+use std::{ collections::HashMap, str::FromStr, sync::{ Arc, Mutex } };
 use chess::Board;
-use crate::bot::{ algorithm::root::search, include::types::{ EngineState, ServerState } };
+use crate::bot::{
+    algorithm::root::search,
+    include::types::{ EngineState, ServerState, TT_TABLE_SIZE },
+};
+use lru::LruCache;
+use std::num::NonZeroUsize;
 
 #[derive(Debug, Deserialize)]
 pub struct EvalRequest {
@@ -49,12 +54,15 @@ pub async fn eval_position_handler(
         }
     }
 
+    let capacity = NonZeroUsize::new(TT_TABLE_SIZE).unwrap();
+    let transposition_table = Arc::new(Mutex::new(LruCache::new(capacity)));
     let mut engine = EngineState {
         game_id: "eval_temp".to_string(),
         current_board,
         history,
         statistics: HashMap::new(),
         global_map: Arc::clone(&state.global_map),
+        transposition_table,
     };
 
     let board = engine.current_board.clone();

src/api/get/get_eval.rs

@@ -1,7 +1,8 @@
 use axum::{ extract::State, response::IntoResponse, Json, http::StatusCode };
+use lru::LruCache;
 use serde::{ Deserialize, Serialize };
-use std::{ collections::HashMap, str::FromStr, sync::Arc };
-use crate::bot::include::types::{ EngineState, ServerState };
+use std::{ collections::HashMap, num::NonZeroUsize, str::FromStr, sync::{ Arc, Mutex } };
+use crate::bot::include::types::{ EngineState, ServerState, TT_TABLE_SIZE };
 use chess::Board;
 
 #[derive(Debug, Deserialize)]
@@ -50,12 +51,15 @@ pub async fn new_game_handler(
         }
     }
 
+    let capacity = NonZeroUsize::new(TT_TABLE_SIZE).unwrap();
+    let transposition_table = Arc::new(Mutex::new(LruCache::new(capacity)));
     let engine = EngineState {
         game_id: payload.game_id.clone(),
         current_board,
         history,
         statistics: HashMap::new(),
         global_map: Arc::clone(&state.global_map),
+        transposition_table,
     };
 
     state.engines.insert(payload.game_id.clone(), engine);

src/api/post/add_game.rs

@@ -52,25 +52,38 @@ pub fn alpha_beta(
     }
 
     *nodes += 1;
-
-    // Track maximum depth reached
     *max_depth_reached = (*max_depth_reached).max(current_depth);
 
     let board_hash = board.get_hash();
     let repetition_count = engine_state.history.get(&board_hash).copied().unwrap_or(0);
     let prioritized_moves = get_prioritized_moves(board, is_noisy);
 
+    {
+        let mut tt = engine_state.transposition_table.lock().unwrap();
+        if let Some(cached_eval) = tt.get(&(board_hash, depth)) {
+            return (None, *cached_eval);
+        }
+    }
+
     match board.status() {
         chess::BoardStatus::Checkmate => {
             let eval = evaluate_board(board);
             let depth_weight = 20 - current_depth.min(20);
             let score = (eval * ((depth_weight as i32) + 1)).clamp(-500_000, 500_000);
+
+            // ✅ Save to TT
+            engine_state.transposition_table.lock().unwrap().put((board_hash, depth), score);
+
             return (None, score);
         }
         chess::BoardStatus::Stalemate => {
+            engine_state.transposition_table.lock().unwrap().put((board_hash, depth), 0);
+
             return (None, 0);
         }
         _ if repetition_count >= 3 => {
+            engine_state.transposition_table.lock().unwrap().put((board_hash, depth), 0);
+
             return (None, 0);
         }
         _ if !is_noisy && depth == 0 => {
@@ -82,14 +95,16 @@ pub fn alpha_beta(
                 nodes,
                 deadline,
                 engine_state,
-                current_depth-1,
+                current_depth - 1,
                 true,
                 current_depth + 1,
                 max_depth_reached
             );
         }
         _ if is_noisy && (prioritized_moves.is_empty() || depth == 0) => {
             let eval = evaluate_board(board);
+            engine_state.transposition_table.lock().unwrap().put((board_hash, depth), eval);
+
             return (None, eval);
         }
         _ => {}
@@ -148,5 +163,7 @@ pub fn alpha_beta(
         }
     }
 
+    engine_state.transposition_table.lock().unwrap().put((board_hash, depth), best_eval);
+
     (best_move, best_eval)
 }

src/bot/algorithm/ab.rs

@@ -48,7 +48,7 @@ pub fn search(
             engine_state,
             depth,
             false,
-            1,
+            0,
             &mut max_depth
         );
 

src/bot/algorithm/root.rs

@@ -2,56 +2,49 @@ use std::{ fs, io::{ self, Write }, path::Path, sync::Arc };
 
 use flate2::read::GzDecoder;
 use once_cell::sync::Lazy;
-use serde_json::Value;
 use tar::Archive;
 
-use crate::bot::include::types::GlobalMap;
+use crate::bot::include::types::{ GlobalMap, OpeningBook };
 
 const COMPRESSED_OPENING_DB: &[u8] = include_bytes!("../../data/openings.tar.gz");
 
-fn read_opening_db() -> Result<Value, io::Error> {
+pub fn read_opening_db() -> Result<OpeningBook, io::Error> {
     let output_dir = Path::new("./db");
     let compressed_path = output_dir.join("openings.tar.gz");
     let file_path = output_dir.join("openingDB.json");
 
-    // If file doesn't exist, extract from embedded archive
     if !file_path.exists() {
         println!("OpeningDB not found, extracting...");
 
         fs::create_dir_all(output_dir)?;
 
-        // Write embedded tar.gz to disk
         {
             let mut file = fs::File::create(&compressed_path)?;
             file.write_all(COMPRESSED_OPENING_DB)?;
         }
 
-        // Extract tar.gz
         let tar_file = fs::File::open(&compressed_path)?;
         let tar = GzDecoder::new(tar_file);
         let mut archive = Archive::new(tar);
         archive.unpack(output_dir)?;
 
-        // Clean up
         fs::remove_file(&compressed_path)?;
         println!("OpeningDB extracted to {:?}", file_path);
     }
 
-    // Read and parse JSON
     let file_content = fs::read_to_string(&file_path)?;
-    let json_data: Value = serde_json
+    let db: OpeningBook = serde_json
         ::from_str(&file_content)
         .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
 
-    Ok(json_data)
+    Ok(db)
 }
-
-pub const OPENING_DB: Lazy<Arc<Value>> = Lazy::new(|| {
+pub static OPENING_DB: Lazy<Arc<OpeningBook>> = Lazy::new(|| {
     Arc::new(read_opening_db().expect("Failed to load opening DB"))
 });
 
 impl GlobalMap {
-    pub fn opening_db() -> Arc<Value> {
+    pub fn opening_db() -> Arc<OpeningBook> {
         Arc::clone(&OPENING_DB)
     }
 

src/bot/include/map.rs

@@ -1,7 +1,9 @@
 use chess::Board;
+use serde::Deserialize;
 use std::collections::HashMap;
-use std::sync::Arc;
+use std::sync::{ Arc, Mutex };
 use dashmap::DashMap;
+use lru::LruCache;
 
 #[derive(Debug, Clone)]
 pub struct Statistics {
@@ -16,6 +18,7 @@ pub struct EngineState {
     pub history: HashMap<u64, u32>,
     pub statistics: HashMap<u64, Statistics>,
     pub global_map: Arc<GlobalMap>,
+    pub transposition_table: Arc<Mutex<LruCache<(u64, u8), i32>>>,
 }
 
 #[derive(Debug)]
@@ -35,3 +38,9 @@ pub enum SpecialMove {
     Promotion,
     EnPassant,
 }
+
+pub const TT_TABLE_SIZE: usize = 100_1000;
+
+#[derive(Debug, Clone, Deserialize)]
+pub struct OpeningEntry(pub String, pub u32);
+pub type OpeningBook = HashMap<u64, Vec<OpeningEntry>>;

src/bot/include/types.rs

@@ -124,7 +124,13 @@ impl BoardExt for Board {
             tactical_bonus += 40;
         }
 
-        return tactical_bonus + capture_value_sum / 10;
+        if let Some(piece) = self.piece_on(mv.get_source()) {
+            if piece == Piece::Pawn {
+                tactical_bonus += 10;
+            }
+        }
+
+        tactical_bonus + capture_value_sum / 10
     }
 
     fn halfmove_clock(&self) -> u32 {

src/bot/util/board.rs

@@ -2,23 +2,20 @@ use chess::{ Board, ChessMove };
 use rand::seq::SliceRandom;
 use rand::thread_rng;
 
-use crate::bot::include::types::GlobalMap;
+use crate::bot::include::map::OPENING_DB;
+use crate::bot::include::types::OpeningEntry;
 use crate::bot::util::moves::parse_uci_move;
 
 pub fn lookup_opening_db(board: &Board) -> Option<ChessMove> {
-    let board_hash_str = board.get_hash().to_string();
-    if let Some(opening_db) = GlobalMap::opening_db().as_object() {
-        if let Some(entry_array) = opening_db.get(&board_hash_str).and_then(|v| v.as_array()) {
-            let mut rng = thread_rng();
-            if let Some(random_entry) = entry_array.choose(&mut rng) {
-                if let Some(uci_str) = random_entry.get(0).and_then(|v| v.as_str()) {
-                    if let Some(chess_move) = parse_uci_move(uci_str, board) {
-                        println!("{}", chess_move);
-                        return Some(chess_move);
-                    }
-                }
-            }
+    let board_hash = board.get_hash();
+    let db = &OPENING_DB;
+
+    if let Some(entries) = db.get(&board_hash) {
+        let mut rng = thread_rng();
+        if let Some(OpeningEntry(uci_str, _weight)) = entries.choose(&mut rng) {
+            return parse_uci_move(uci_str, board);
         }
     }
+
     None
 }