zeal/src/search/negamax.rs

use anyhow::{anyhow, bail, Result};

use crate::{
    board::game::Game,
    evaluation::{MATE_SCORE, MIN_SCORE},
};

use super::{
    move_ordering::score_move,
    quiescence::quiescence,
    time::TimeInfo,
    transposition_table::{NodeType, TTEntry},
};

pub fn negamax(
    game: &mut Game,
    mut alpha: i32,
    beta: i32,
    mut depth: u8,
    plies: u8,
    time: &TimeInfo,
    nodes: &mut u64,
) -> Result<i32> {
    if time.exceed_hard_limit() {
        bail!("Hard limit exceeded in negamax");
    }

    if plies != 0 && game.in_repetition() {
        return Ok(0);
    }

    let color = game.current_player();
    let in_check = game.board.king_under_check(color);

    if in_check {
        depth += 1;
    }

    if depth == 0 {
        let q_score = quiescence(game, alpha, beta, time, nodes).map_err(|e| anyhow!("{e}"))?;
        return Ok(q_score);
    }

    let mut legal_moves = 0;
    let mut best_score = MIN_SCORE;
    let mut best_move = None;
    let mut moves = game.board.pseudo_moves_all();
    let tt_entry = game.tt.lookup(game.hash);

    if let Some(tt_entry) = tt_entry {
        if plies > 0 && tt_entry.hash == game.hash && tt_entry.depth >= depth {
            match tt_entry.node_type {
                NodeType::Exact => return Ok(tt_entry.score),
                NodeType::LowerBound if tt_entry.score >= beta => return Ok(tt_entry.score),
                NodeType::UpperBound if tt_entry.score <= alpha => return Ok(tt_entry.score),
                _ => (),
            }
        }
    }

    moves.sort_unstable_by_key(|mv| {
        score_move(
            &game.mailbox,
            *mv,
            game.killer[plies as usize],
            tt_entry.and_then(|entry| entry.mv),
        )
    });

    for mv in moves {
        game.make_move(&mv);

        if game.board.king_under_check(color) {
            game.unmake_move();
            continue;
        }
        legal_moves += 1;
        *nodes += 1;

        let score = if legal_moves == 1 {
            -negamax(game, -beta, -alpha, depth - 1, plies + 1, time, nodes)?
        } else {
            let mut score = -negamax(game, -alpha - 1, -alpha, depth - 1, plies + 1, time, nodes)?;
            if score > alpha && score < beta {
                score = -negamax(game, -beta, -alpha, depth - 1, plies + 1, time, nodes)?;
            }
            score
        };

        game.unmake_move();

        if score > best_score {
            best_score = score;
        }

        if score > alpha {
            alpha = score;
            best_move = Some(mv);
        }

        if score >= beta {
            if !mv.is_capture() {
                game.killer[plies as usize] = Some(mv);
            }
            break;
        }
    }

    if legal_moves == 0 {
        if in_check {
            return Ok(-MATE_SCORE + plies as i32);
        }
        return Ok(0);
    }

    let node_type = match best_score {
        s if s >= beta => NodeType::LowerBound,
        s if s <= alpha => NodeType::UpperBound,
        _ => NodeType::Exact,
    };

    game.tt.insert(TTEntry::new(
        game.hash, best_move, depth, best_score, node_type,
    ));
    Ok(best_score)
}

#[cfg(test)]
mod tests {
    use crate::board::{fen::from_fen, square::Square};
    use crate::evaluation::{MAX_SCORE, MIN_SCORE};
    use crate::movegen::attack_generator::init_attacks;
    use crate::movegen::r#move::Move;
    use crate::search::negamax::negamax;
    use crate::search::time::TimeInfo;
    use crate::search::{INC, TIME};

    const FEN_MATE_IN_1: [&str; 2] = [
        "8/8/8/8/8/4q1k1/8/5K2 b - - 0 1",
        "8/8/8/8/8/4Q1K1/8/5k2 w - - 0 1",
    ];

    #[test]
    fn test_negamax() -> anyhow::Result<()> {
        init_attacks();
        let mut game = from_fen(FEN_MATE_IN_1[0]).unwrap();

        let e3f2 = Move::new(Square::E3, Square::F2);
        let time_info = TimeInfo::new(std::time::Instant::now(), TIME, INC);
        negamax(&mut game, MIN_SCORE, MAX_SCORE, 2, 0, &time_info, &mut 0)
            .expect("Expected a search result");

        assert_eq!(e3f2, game.tt.lookup(game.hash).unwrap().mv.unwrap());

        let mut game = from_fen(FEN_MATE_IN_1[1]).unwrap();

        let e3f2 = Move::new(Square::E3, Square::F2);
        negamax(&mut game, MIN_SCORE, MAX_SCORE, 2, 0, &time_info, &mut 0)
            .expect("Expected a search result");

        assert_eq!(e3f2, game.tt.lookup(game.hash).unwrap().mv.unwrap());

        Ok(())
    }
}