crabchess 0.1.15

//! This module allows you to create performant trees of chess openings nested by move.

use fxhash::FxHashMap;
use std::{fs, sync::OnceLock};

/// The `Opening` struct represents an ECO opening.
///
/// The `moves` field should be formatted like: `"1. c4 Nf6 2. Nf3 e6 3. g3 a6 4. Bg2 b5"`.
#[derive(Debug, PartialEq, Eq, Clone)]
pub struct Opening {
    pub eco: String,
    pub name: String,
    pub moves: String,
}

impl Opening {
    /// Iterate through moves in `moves` field as string slices.
    pub fn iter_moves(&self) -> impl Iterator<Item = &str> {
        self.moves.split(' ').filter(|s| !s.contains('.'))
    }
}

/// Opening book from [Lichess](https://github.com/lichess-org/chess-openings).
///
/// # Panics
///
/// This function will panic if the default book is corrupted or cannot be read.
pub fn default_book() -> &'static Vec<Opening> {
    static DEFAULT_BOOK: OnceLock<Vec<Opening>> = OnceLock::new();
    DEFAULT_BOOK.get_or_init(|| {
        fs::read_to_string("src/data/book.csv")
            .unwrap()
            .split('\n')
            .filter_map(|line| {
                let mut iter = line.split(';');
                Some(Opening {
                    eco: iter.next()?.into(),
                    name: iter.next()?.into(),
                    moves: iter.next()?.into(),
                })
            })
            .collect()
    })
}

/// Opening tree from [Lichess](https://github.com/lichess-org/chess-openings).
///
/// # Panics
///
/// This function will panic if the default book is corrupted or cannot be read.
pub fn default_tree() -> &'static OpeningTree {
    static DEFAULT_TREE: OnceLock<OpeningTree> = OnceLock::new();
    DEFAULT_TREE.get_or_init(|| OpeningTree::from(default_book().as_slice()))
}

/// A tree of chess openings nested by move for fast search and retrieval.
#[derive(Default, Debug, Clone, PartialEq, Eq)]
pub struct OpeningTree {
    pub element: Option<Opening>,
    pub children: Option<FxHashMap<String, OpeningTree>>,
}

impl OpeningTree {
    /// Get the corresponding opening for a sequence of moves in SAN notation.
    #[must_use]
    pub fn get_opening<'a>(
        &self,
        sans: impl Iterator<Item = &'a str>,
    ) -> Option<&Opening> {
        let mut cursor = self;
        let mut longest_match: Option<&Opening> = None;

        for san in sans {
            let Some(ref children) = cursor.children else {
                break;
            };
            match children.get(san) {
                Some(tree) => {
                    cursor = tree;
                    if let Some(opening) = &cursor.element {
                        longest_match = Some(opening);
                    }
                }
                None => break,
            }
        }

        longest_match
    }
}

impl From<&[Opening]> for OpeningTree {
    fn from(value: &[Opening]) -> Self {
        let mut root = Self::default();

        for opening in value {
            let mut cursor = &mut root;

            for san in opening.iter_moves() {
                let children = cursor.children.get_or_insert_with(FxHashMap::default);
                cursor = children.entry(san.into()).or_default();
            }

            cursor.element = Some(opening.clone());
        }

        root
    }
}

impl From<OpeningTree> for Vec<Opening> {
    fn from(value: OpeningTree) -> Self {
        let mut output = Self::new();

        if let Some(opening) = value.element {
            output.push(opening);
        }

        if let Some(children) = value.children {
            for child in children.into_values() {
                output.append(&mut child.into());
            }
        }

        output
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use rayon::prelude::*;

    #[test]
    fn test_default_book() {
        let book: Vec<Opening> = default_book().clone();
        assert_eq!(3455, book.len());
    }

    #[test]
    fn test_default_tree() {
        let tree = default_tree().clone();
        let moves = "d4 Nf6 c4 e6 Nc3 Bb4 Qb3 Bxc3+ Qxc3 O-O Bg5 c5 dxc5 Nc6 Nf3 Qa5 Bxf6 gxf6 Qxa5 Nxa5 e3 Rd8 Rd1 Kg7 Be2 b6 Rd4 bxc5 Rg4+ Kh6 Bd3 f5 Rh4+ Kg6 g4 Ba6 gxf5+ exf5 Ne5+ Kf6 Rh6+ Kxe5 f4#".split(' ');
        let opening = tree.get_opening(moves).unwrap();
        assert_eq!(opening.eco, "E22");
    }

    #[test]
    fn test_tree_perf() {
        const MOVE_SEQUENCES_PATH: &str = "test_data/bare_move_sequences.json";
        use std::time;

        // Time tree formation.
        let start_time = std::time::SystemTime::now()
            .duration_since(time::UNIX_EPOCH)
            .unwrap();
        let tree: OpeningTree = default_tree().clone();
        let end_time = std::time::SystemTime::now()
            .duration_since(time::UNIX_EPOCH)
            .unwrap();
        println!("File -> Tree: {:?} elapsed.", end_time - start_time);

        let json = fs::read_to_string(MOVE_SEQUENCES_PATH).unwrap();
        let games: Vec<String> = serde_json::from_str(&json).unwrap();
        let games: Vec<Vec<&str>> = games
            .par_iter()
            .map(|v: &String| v.split(' ').collect())
            .collect();

        // Time opening retrieval for 65k games, sourced from Lichess:
        // <https://www.kaggle.com/datasets/datasnaek/chess>
        let start_time = std::time::SystemTime::now()
            .duration_since(time::UNIX_EPOCH)
            .unwrap();
        games.into_par_iter().for_each(|game| {
            let _ = tree.get_opening(game.into_iter());
        });
        let end_time = std::time::SystemTime::now()
            .duration_since(time::UNIX_EPOCH)
            .unwrap();
        println!("65k games: {:?} elapsed.", end_time - start_time);
    }
}