tak_rs/tak_lib/src/engine.rs

use crate::{
    Board, Direction, Move,
    piece::{ColoredPiece, Piece},
    tile::Tile,
};

impl<const N: usize> Board<N> {
    /// Gets an iterator for the moves that can be made by
    /// the active player on this board.
    pub fn moves(&self) -> impl Iterator<Item = Move<N>> {
        let to_move = self.to_move;
        let tiles = &self.tiles;
        let can_place_caps = if Self::N_CAPS == 0 {
            false
        } else {
            let caps_on_board = self
                .tiles
                .iter()
                .flat_map(|col| col.iter())
                .filter(|t| {
                    let piece = t.top_piece();
                    piece.color_opt() == Some(to_move) && piece.piece() == Piece::Caps
                })
                .count() as u8;
            caps_on_board < Self::N_CAPS
        };

        let chords = (0..(N as u8)).flat_map(|x| (0..(N as u8)).map(move |y| (x, y)));

        chords.flat_map(move |(x, y)| {
            let top_piece = tiles[x as usize][y as usize].top_piece();
            if top_piece == ColoredPiece::None {
                // possible moves are place moves
                get_placement_moves(x, y, can_place_caps)
            } else if top_piece.color() == to_move {
                get_movement_moves(x, y, tiles)
            } else {
                Vec::with_capacity(0)
            }
        })
    }
}

fn get_placement_moves<const N: usize>(x: u8, y: u8, can_place_caps: bool) -> Vec<Move<N>> {
    let mut moves = Vec::with_capacity(2 + can_place_caps as usize);
    moves.push(Move::Place {
        piece: Piece::Flat,
        at: (x, y),
    });
    moves.push(Move::Place {
        piece: Piece::Wall,
        at: (x, y),
    });
    if can_place_caps {
        moves.push(Move::Place {
            piece: Piece::Caps,
            at: (x, y),
        });
    }
    moves
}

pub(crate) fn get_movement_moves<const N: usize>(
    x: u8,
    y: u8,
    tiles: &[[Tile; N]; N],
) -> Vec<Move<N>> {
    let mut moves = Vec::new();
    let tile = tiles[x as usize][y as usize];
    let is_caps = tile.top_piece().piece() == Piece::Caps;
    let max_grab = u8::min(N as u8, tile.stack_len() + 1);

    for dir in Direction::ALL {
        get_moves_in_dir(x, y, dir, tiles, &mut moves, is_caps, max_grab);
    }

    moves
}

/// Gets the possible moves moving from (`x`, `y`) in the given direction `dir`.
///
/// This function does not check the colors of anything, so the caller must make sure
/// the active player owns the piece on this tile, if you want to get only valid moves
/// for the active player.
fn get_moves_in_dir<const N: usize>(
    x: u8,
    y: u8,
    dir: Direction,
    tiles: &[[Tile; N]; N],
    moves: &mut Vec<Move<N>>,
    is_caps: bool,
    max_grab: u8,
) {
    // first, get the max distance we can go
    let mut max_dist = max_grab.min(dir.dist_to_edge::<N>(x, y));
    let mut end_w_capstone = false;
    let (dx, dy) = dir.delta();

    // step through the tiles to see if out movement gets blocked
    for i in 1..=max_dist {
        let xi = x as isize + (dx * i as isize);
        let yi = y as isize + (dy * i as isize);
        // we already clamp `max_distance` to the board edge, so indexing should be safe
        let next_piece = tiles[xi as usize][yi as usize].top_piece().piece();

        // check for pieces that will block us
        if next_piece == Piece::Caps {
            max_dist = i - 1;
            break;
        } else if next_piece == Piece::Wall {
            // if we have a capstone, we can smash the wall
            if is_caps {
                end_w_capstone = true;
                max_dist = i;
            } else {
                max_dist = i - 1;
            }
            break;
        }
    }
    if max_dist == 0 {
        return;
    }

    // now that we have our max distance, get all the possible moves
    let mut insert_fn = |drops: [u8; N]| {
        moves.push(Move::Move {
            from: (x, y),
            count: drops.iter().copied().sum::<u8>(),
            dir,
            drop: drops,
        });
    };
    for grab in 1..=max_grab {
        let mut buf = [0; N];
        get_drop_amounts(&mut buf, 0, &mut insert_fn, grab, max_dist, end_w_capstone);
    }
}

/// Gets the possible amounts of pieces to drop at each tile.
///
/// `pieces` is the number of pieces that are picked up.
/// `len` is the max length that the stack can be moved.
///
/// If the top piece of the stack is a capstone, and the end of the line
/// is a wall, `end_w_capstone` should be `true` and `len` should include
/// that tile.
pub(crate) fn get_drop_amounts<const N: usize>(
    // shared state
    buf: &mut [u8; N],
    i: usize,
    insert_fn: &mut impl FnMut([u8; N]),
    // recursive parameters
    pieces: u8,
    len: u8,
    end_w_capstone: bool,
) {
    debug_assert!(len > 0 || pieces == 0);
    // base cases
    if pieces == 0 {
        insert_fn(buf.clone());
        return;
    }
    if pieces == 1 {
        buf[i] = 1;
        insert_fn(buf.clone());
        buf[i] = 0;
        return;
    }
    if len == 1 {
        if end_w_capstone {
            if pieces == 1 {
                buf[i] = 1;
                insert_fn(buf.clone());
                buf[i] = 0;
            }
        } else {
            buf[i] = pieces;
            insert_fn(buf.clone());
            buf[i] = 0;
        }
        return;
    }

    // recursive case
    for n in 1..=pieces {
        buf[i] = n;
        get_drop_amounts(buf, i + 1, insert_fn, pieces - n, len - 1, end_w_capstone);
    }
    buf[i] = 0;
}