bitstackchess 0.1.1

// src/engine/transposition_table.rs

//! A simple fixed‐size transposition table using Zobrist keys.
//! Each entry stores: (key, depth, value, flag, best_move).
//! 
//! On store: if table index is vacant, place it. If occupied, replace if new depth >= existing depth.
//! 
//! Lookup returns Option<(value, flag, best_move)> if the stored key matches.

use crate::core::Move10;

/// Flags for TTEntries.
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum TTFlag {
    Exact,
    LowerBound,
    UpperBound,
}

/// A single transposition‐table entry.
#[derive(Copy, Clone, Debug)]
pub struct TTEntry {
    pub key: u64,
    pub depth: u8,           // search depth at which this entry was stored
    pub value: i32,          // evaluation or score
    pub flag: TTFlag,        // exact/lower/upper bound
    pub best_move: Option<Move10>, // best move found at this node
}

/// A simple fixed‐size transposition table using direct‐mapped indexing.
pub struct TranspositionTable {
    size_mask: usize,              // mask = (capacity − 1); capacity must be power of two
    entries: Vec<Option<TTEntry>>, // length = capacity
}

impl TranspositionTable {
    /// Create a new table with capacity = next power‐of‐two ≥ `requested_size`.
    pub fn new(requested_size: usize) -> TranspositionTable {
        let mut cap = 1;
        while cap < requested_size {
            cap <<= 1;
        }
        TranspositionTable {
            size_mask: cap - 1,
            entries: vec![None; cap],
        }
    }

    /// Given a Zobrist `key`, compute its index: lower bits of key & size_mask.
    #[inline]
    fn index_of(&self, key: u64) -> usize {
        (key as usize) & self.size_mask
    }

    /// Store an entry. If a different key is present at the same index, replace only if new depth >= old depth.
    pub fn store(
        &mut self,
        key: u64,
        depth: u8,
        value: i32,
        flag: TTFlag,
        best_move: Option<Move10>,
    ) {
        let idx = self.index_of(key);
        match self.entries[idx] {
            Some(existing) => {
                if depth >= existing.depth {
                    self.entries[idx] = Some(TTEntry {
                        key,
                        depth,
                        value,
                        flag,
                        best_move,
                    });
                }
            }
            None => {
                self.entries[idx] = Some(TTEntry {
                    key,
                    depth,
                    value,
                    flag,
                    best_move,
                });
            }
        }
    }

    /// Look up a `key`. Returns `Some((value, flag, best_move, stored_depth))` if found and keys match.
    pub fn lookup(&self, key: u64) -> Option<(i32, TTFlag, Option<Move10>, u8)> {
        let idx = self.index_of(key);
        if let Some(entry) = self.entries[idx] {
            if entry.key == key {
                return Some((entry.value, entry.flag, entry.best_move, entry.depth));
            }
        }
        None
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::engine::zobrist::compute_hash;
    use crate::board::{init_chess_positions, PieceMapping};

    #[test]
    fn basic_store_and_lookup() {
        // Build a sample position:
        let mut mapping = PieceMapping::new_empty();
        for &(pid, sq) in &init_chess_positions() {
            mapping.place_piece(pid, sq);
        }
        let key = compute_hash(&mapping, 0, None);
        let mut tt = TranspositionTable::new(16);
        // Store a value:
        tt.store(key, 4, 42, TTFlag::Exact, None);
        // Lookup should find it:
        if let Some((val, flag, mv_opt, depth)) = tt.lookup(key) {
            assert_eq!(val, 42);
            assert_eq!(flag, TTFlag::Exact);
            assert!(mv_opt.is_none());
            assert_eq!(depth, 4);
        } else {
            panic!("Expected to find key in table");
        }

        // If we store a lower‐depth entry at same key, it should not overwrite:
        tt.store(key, 2, 100, TTFlag::LowerBound, Some(Move10::new(0, 0)));
        if let Some((val, flag, _, depth)) = tt.lookup(key) {
            assert_eq!(val, 42);
            assert_eq!(flag, TTFlag::Exact);
            assert_eq!(depth, 4);
        } else {
            panic!("Expected to find key in table");
        }

        // If we store a higher‐depth entry at same key, it should overwrite:
        tt.store(key, 6, -13, TTFlag::UpperBound, Some(Move10::new(1, 1)));
        if let Some((val, flag, mv_opt, depth)) = tt.lookup(key) {
            assert_eq!(val, -13);
            assert_eq!(flag, TTFlag::UpperBound);
            assert_eq!(depth, 6);
            assert_eq!(mv_opt.unwrap(), Move10::new(1, 1));
        } else {
            panic!("Expected to find key in table");
        }
    }

    #[test]
    fn collision_overwrite_based_on_depth() {
        let mut mapping = PieceMapping::new_empty();
        for &(pid, sq) in &init_chess_positions() {
            mapping.place_piece(pid, sq);
        }
        let key1 = compute_hash(&mapping, 0, None);

        // Make a trivial change so we get a different key (flip side):
        let key2 = compute_hash(&mapping, 1, None);

        // Force them to collide by using a table of size=1 (mask=0).
        let mut tt = TranspositionTable::new(1);

        tt.store(key1, 3, 10, TTFlag::Exact, None);
        // Next insert different key at same index; because 4 >= 3, it should overwrite.
        tt.store(key2, 4, 20, TTFlag::LowerBound, None);

        // Only key2 remains:
        assert!(tt.lookup(key1).is_none());
        if let Some((val, _, _, depth)) = tt.lookup(key2) {
            assert_eq!(val, 20);
            assert_eq!(depth, 4);
        } else {
            panic!("Expected key2 in table");
        }
    }
}