chessfriend/explorer/src/main.rs

// Eryn Wells <eryn@erynwells.me>

use chessfriend_board::{Board, ZobristState, fen::FromFenStr};
use chessfriend_core::{Color, Piece, Shape, Square, Wing, random::RandomNumberGenerator};
use chessfriend_moves::{GeneratedMove, ValidateMove, algebraic::AlgebraicMoveComponents};
use chessfriend_position::{PlacePieceStrategy, Position, fen::ToFenStr};
use clap::{Arg, Command, value_parser};
use rustyline::{DefaultEditor, error::ReadlineError};
use std::sync::Arc;
use thiserror::Error;

struct CommandResult {
    should_continue: bool,
    should_print_position: bool,
}

impl Default for CommandResult {
    fn default() -> Self {
        CommandResult {
            should_continue: true,
            should_print_position: true,
        }
    }
}

struct State {
    position: Position,
    zobrist: Arc<ZobristState>,
}

fn command_line() -> Command {
    // strip out usage
    const PARSER_TEMPLATE: &str = "\
        {all-args}
    ";

    Command::new("explorer")
        .multicall(true)
        .arg_required_else_help(true)
        .subcommand_required(true)
        .subcommand_value_name("CMD")
        .subcommand_help_heading("COMMANDS")
        .help_template(PARSER_TEMPLATE)
        .subcommand(Command::new("fen").about("Print the current position as a FEN string"))
        .subcommand(Command::new("flags").about("Print flags for the current position"))
        .subcommand(
            Command::new("load")
                .arg(Arg::new("fen").required(true))
                .alias("l")
                .about("Load a board position from a FEN string"),
        )
        .subcommand(
            Command::new("make")
                .arg(Arg::new("move").required(true))
                .alias("m")
                .about("Make a move"),
        )
        .subcommand(
            Command::new("place")
                .arg(Arg::new("color").required(true))
                .arg(Arg::new("piece").required(true))
                .arg(Arg::new("square").required(true))
                .alias("p")
                .about("Place a piece on the board"),
        )
        .subcommand(
            Command::new("sight")
                .arg(Arg::new("square").required(true))
                .about("Show sight of a piece on a square"),
        )
        .subcommand(
            Command::new("moves")
                .arg(Arg::new("square").required(false))
                .about("Show moves of a piece on a square. With no argument, show all moves for the active color."),
        )
        .subcommand(
            Command::new("movement")
                .arg(Arg::new("square").required(true))
                .about("Show moves of a piece on a square"),
        )
        .subcommand(
            Command::new("perft")
                .arg(Arg::new("depth")
                    .required(true)
                    .value_parser(value_parser!(usize))
                )
                .about("Run Perft on the current position to the given depth")
        )
        .subcommand(
            Command::new("reset")
                .subcommand(Command::new("clear").about("Reset to a cleared board"))
                .subcommand(Command::new("starting").about("Reset to the starting position"))
                .subcommand(
                    Command::new("fen")
                        .arg(Arg::new("fen").required(true))
                        .about("Reset to a position described by a FEN string"),
                )
                .about("Reset the board"),
        )
        .subcommand(Command::new("print").about("Print the board"))
        .subcommand(Command::new("quit").alias("exit").about("Quit the program"))
        .subcommand(Command::new("zobrist").about("Print the Zobrist hash of the current board"))
}

#[derive(Clone, Debug, Error, Eq, PartialEq)]
enum CommandHandlingError<'a> {
    #[error("lexer error")]
    LexerError,

    #[error("missing {0} argument")]
    MissingArgument(&'a str),

    #[error("no piece on {0}")]
    NoPiece(Square),

    #[error("{value:?} is not a valid value for {argument_name:?}")]
    ValueError {
        argument_name: &'static str,
        value: String,
    },
}

fn respond(line: &str, state: &mut State) -> anyhow::Result<CommandResult> {
    let args = shlex::split(line).ok_or(CommandHandlingError::LexerError)?;
    let matches = command_line().try_get_matches_from(args)?;

    let mut result = CommandResult::default();

    match matches.subcommand() {
        Some(("flags", matches)) => result = do_flags_command(state, matches),
        Some(("load", matches)) => result = do_load_command(state, matches)?,
        Some(("print", _matches)) => {}
        Some(("quit", _matches)) => {
            result.should_continue = false;
            result.should_print_position = false;
        }
        Some(("fen", _matches)) => {
            println!("{}", state.position.to_fen_str()?);
            result.should_print_position = false;
        }
        Some(("make", matches)) => result = do_make_command(state, matches)?,
        Some(("perft", matches)) => result = do_perft_command(state, matches)?,
        Some(("place", matches)) => {
            let color = matches
                .get_one::<String>("color")
                .ok_or(CommandHandlingError::MissingArgument("color"))?;
            let color = color.parse::<Color>()?;

            let shape = matches
                .get_one::<String>("piece")
                .ok_or(CommandHandlingError::MissingArgument("piece"))?;
            let shape = shape.parse::<Shape>()?;

            let square = matches
                .get_one::<String>("square")
                .ok_or(CommandHandlingError::MissingArgument("square"))?;
            let square = Square::from_algebraic_str(square)?;

            let piece = Piece::new(color, shape);

            state
                .position
                .place_piece(piece, square, PlacePieceStrategy::default())?;
        }
        Some(("sight", matches)) => {
            let square = matches
                .get_one::<String>("square")
                .ok_or(CommandHandlingError::MissingArgument("square"))?;
            let square = square.parse::<Square>()?;

            let sight = state.position.sight(square);

            let display = state.position.display().highlight(sight);
            println!("\n{display}");

            result.should_print_position = false;
        }
        Some(("moves", matches)) => result = do_moves_command(state, matches)?,
        Some(("movement", matches)) => result = do_movement_command(state, matches)?,
        Some(("reset", matches)) => result = do_reset_command(state, matches)?,
        Some(("zobrist", matches)) => result = do_zobrist_command(state, matches),
        Some((name, _matches)) => unimplemented!("{name}"),
        None => unreachable!("Subcommand required"),
    }

    Ok(result)
}

fn do_flags_command(state: &mut State, _matches: &clap::ArgMatches) -> CommandResult {
    let board = &state.position.board;

    println!("Castling:");

    for (color, wing) in [
        (Color::White, Wing::KingSide),
        (Color::White, Wing::QueenSide),
        (Color::Black, Wing::KingSide),
        (Color::Black, Wing::QueenSide),
    ] {
        let has_right = board.has_castling_right_unwrapped(color, wing.into());
        let can_castle = board.color_can_castle(wing, Some(color));

        let can_castle_message = match can_castle {
            Ok(_) => "ok".to_string(),
            Err(error) => format!("{error}"),
        };

        println!("  {color} {wing}: {has_right}, {can_castle_message}");
    }

    CommandResult {
        should_continue: true,
        should_print_position: false,
    }
}

fn do_load_command(state: &mut State, matches: &clap::ArgMatches) -> anyhow::Result<CommandResult> {
    let fen_string = matches
        .get_one::<String>("fen")
        .ok_or(CommandHandlingError::MissingArgument("fen"))?;

    let mut board = Board::from_fen_str(fen_string.as_str())?;
    board.set_zobrist_state(state.zobrist.clone());

    state.position = Position::new(board);

    Ok(CommandResult {
        should_continue: true,
        should_print_position: true,
    })
}

fn do_make_command(state: &mut State, matches: &clap::ArgMatches) -> anyhow::Result<CommandResult> {
    let move_string = matches
        .get_one::<String>("move")
        .ok_or(CommandHandlingError::MissingArgument("move"))?;

    let algebraic_move: AlgebraicMoveComponents = move_string.parse()?;

    let encoded_move = state
        .position
        .move_from_algebraic_components(algebraic_move)
        .ok_or(CommandHandlingError::ValueError {
            argument_name: "move",
            value: move_string.to_string(),
        })?;

    state.position.make_move(encoded_move, ValidateMove::Yes)?;

    Ok(CommandResult::default())
}

fn do_reset_command(
    state: &mut State,
    matches: &clap::ArgMatches,
) -> anyhow::Result<CommandResult> {
    match matches.subcommand() {
        None | Some(("clear", _)) => state.position = Position::empty(Some(state.zobrist.clone())),
        Some(("starting", _)) => state.position = Position::starting(Some(state.zobrist.clone())),
        Some(("fen", matches)) => {
            let fen = matches
                .get_one::<String>("fen")
                .ok_or(CommandHandlingError::MissingArgument("fen"))?;
            let board = Board::from_fen_str(fen)?;
            state.position = Position::new(board);
        }
        Some((name, _matches)) => unimplemented!("{name}"),
    }

    Ok(CommandResult::default())
}

fn do_moves_command(
    state: &mut State,
    matches: &clap::ArgMatches,
) -> anyhow::Result<CommandResult> {
    let moves: Vec<GeneratedMove> = if let Some(square) = matches
        .get_one::<String>("square")
        .and_then(|square| square.parse::<Square>().ok())
    {
        state
            .position
            .moves_for_piece(square)
            .map(|it| it.filter(|ply| ply.origin() == square))
            .map(Iterator::collect)
            .ok_or(CommandHandlingError::NoPiece(square))?
    } else {
        state.position.all_moves(None).collect()
    };

    let formatted_moves: Vec<String> = moves
        .iter()
        .map(|ply| {
            let origin = ply.origin();
            if let Some(piece) = state.position.get_piece(origin) {
                format!("{piece}{ply}")
            } else {
                format!("{ply}??")
            }
        })
        .collect();

    if !formatted_moves.is_empty() {
        let max_length = formatted_moves
            .iter()
            .map(|s| s.chars().count())
            .max()
            .unwrap_or(8)
            + 2;

        let columns_count = 80 / max_length;
        for row in formatted_moves.chunks(columns_count) {
            for ply in row {
                print!("{ply:<max_length$}");
            }
            println!();
        }
    }

    Ok(CommandResult {
        should_continue: true,
        should_print_position: false,
    })
}

fn do_movement_command(
    state: &mut State,
    matches: &clap::ArgMatches,
) -> anyhow::Result<CommandResult> {
    let square = *matches
        .get_one::<Square>("square")
        .ok_or(CommandHandlingError::MissingArgument("square"))?;

    let movement = state.position.movement(square);
    let display = state.position.display().highlight(movement);
    println!("\n{display}");

    Ok(CommandResult {
        should_continue: true,
        should_print_position: false,
    })
}

fn do_perft_command(
    state: &mut State,
    matches: &clap::ArgMatches,
) -> anyhow::Result<CommandResult> {
    let depth = *matches
        .get_one::<usize>("depth")
        .ok_or(CommandHandlingError::MissingArgument("depth"))?;

    let mut position = state.position.clone();
    let nodes_count = position.perft(depth);

    println!("nodes {nodes_count}");

    Ok(CommandResult {
        should_continue: true,
        should_print_position: false,
    })
}

fn do_zobrist_command(state: &mut State, _matches: &clap::ArgMatches) -> CommandResult {
    if let Some(hash) = state.position.zobrist_hash() {
        println!("hash:{hash}");
    } else {
        println!("No Zobrist hash available");
    }

    CommandResult {
        should_continue: true,
        should_print_position: false,
    }
}

fn main() -> Result<(), String> {
    let mut editor = DefaultEditor::new().map_err(|err| format!("Error: {err}"))?;

    let mut rng = RandomNumberGenerator::default();
    let zobrist_state = Arc::new(ZobristState::new(&mut rng));

    let starting_position = Position::starting(Some(zobrist_state.clone()));
    let mut state = State {
        position: starting_position,
        zobrist: zobrist_state,
    };

    let mut should_print_position = true;

    loop {
        if should_print_position {
            println!("{}", &state.position);
            println!("{} to move.", state.position.board.active_color());
        }

        let readline = editor.readline("\n? ");
        match readline {
            Ok(line) => {
                let line = line.trim();
                if line.is_empty() {
                    continue;
                }

                match respond(line, &mut state) {
                    Ok(result) => {
                        should_print_position = result.should_print_position;
                        if !result.should_continue {
                            break;
                        }
                    }
                    Err(message) => println!("{message}"),
                }
            }
            Err(ReadlineError::Interrupted) => {
                println!("CTRL-C");
                break;
            }
            Err(ReadlineError::Eof) => {
                println!("CTRL-D");
                break;
            }
            Err(err) => {
                println!("Error: {err}");
                break;
            }
        }
    }

    Ok(())
}