Building a Chess AI, Part 2 – Move Generation

In my last post I wrote about modelling game state, the first milestone of a project I’ve been working on to build a chess engine in Rust. In this post I’m going to write an overview of the second milestone: move generation.

If you’re interested in seeing the full code, check out tomcant/chess-rs on GitHub.

So far we have a few types to represent some of the game’s basic concepts: colours, pieces and squares, and a model of the board that uses bitboards to track the locations of pieces.

The approach I’ve taken to generating moves is to calculate a bitboard of all the squares under attack for each of the colour-to-move’s pieces, and convert each of these to a list of moves. Combining these into one list produces the majority of moves in a position, leaving just a handful of “special” cases like castling and en-passant.

The main topics to cover are:

  1. using bitboards to calculate “attack squares”,
  2. handling “special” cases with bitboards,
  3. converting bitboards to moves, and
  4. verifying correctness with perft.

Calculating Attack Squares

At the heart of the move generator is a simple looking function for calculating attacks:

fn get_attacks(piece: Piece, square: Square, board: &Board) -> u64 {
    match piece {
        WP | BP => get_pawn_attacks(...),
        WN | BN => get_knight_attacks(...),
        WB | BB => get_bishop_attacks(...),
        WR | BR => get_rook_attacks(...),
        WQ | BQ => get_bishop_attacks(...) | get_rook_attacks(...),
        WK | BK => get_king_attacks(...),
    }
}

Given a piece and a square, this returns a bitboard containing all the squares under attack by that piece. We can then convert this bitboard into a list of moves and continue with the next piece.

Sliding Piece Attacks

Bishops, rooks and queens all slide along their paths for as long as there isn’t another piece blocking the way. This is different from non-sliding pieces where the squares under attack are always known upfront, and so we require different approaches.

There are lots of popular methods for calculating sliding piece attacks using bitboards. As a starting point I implemented the simplest method I could find: the classical approach.

Consider a bishop on c3 with no blocking pieces:

We can construct a bitboard from the highlighted squares and use it any time we need the attacks for a bishop on c3. However, if pieces block the path then we’d need to figure out which squares are actually reachable first. Take the following position for example:

We need to determine that g7 and h8 should be excluded from the bishop’s attacks. One approach is to iterate over the bishop’s path until we hit a piece, but this would be slow because we’d be querying the state of the board a lot.

Instead, if we take the bishop’s path and perform a bitwise AND operation with all the pieces on the board, we get the following:

AND
=

This tells us where the blocking pieces are. Next, we perform bitwise XOR with just the section of path after the first blocking square:

XOR
=

This gives the attack squares for the bishop taking into account the pieces that block its path. Repeating for each direction yields the complete bitboard of attack squares.

Here’s what it looks like in code:

enum BishopDirection {
  NW, NE, SW, SE,
}

fn get_bishop_attacks(square: Square, board: &Board) -> u64 {
    let mut attacks = 0;

    for direction in [NW, NE, SW, SE] {
        let path = BISHOP_ATTACKS[square][direction];
        let blockers = path & board.occupancy();

        if blockers == 0 {
            // No blockers in this direction so
            // add the whole path and move on.
            attacks |= path;
            continue;
        }

        let first_blocking_square = match direction {
            NW | NE => Square::first(blockers),
            SW | SE => Square::last(blockers),
        };

        attacks |= path ^ BISHOP_ATTACKS[first_blocking_square][direction];
    }

    attacks
}

Notice the use of BISHOP_ATTACKS, a pre-calculated array of 64 x 4 bitboards indexed by square and direction. This is used to look up the squares a bishop can attack on an empty board and means we only have to calculate this once.

Finding the first blocking square

I want to highlight how the first blocking square calculation works because it’s not obvious from the code above. When the bishop’s path points north-east or north-west, the first blocker is always the first 1-bit in the blockers bitboard, which is just the number of trailing zeros:

impl Square {
    fn first(squares: u64) -> Self {
        Self(squares.trailing_zeros())
    }
}
let first_blocking_square = match direction {
    NW | NE => Square::first(blockers),
    SW | SE => Square::last(blockers),
};

Conversely, when there are blockers on a south-east or south-west path then the first blocker is always the last 1-bit, so we identify it by the number of leading zeros instead. A simple calculation is then used to flip its orientation:

impl Square {
    fn last(squares: u64) -> Self {
        Self(63 - squares.leading_zeros())
    }
}
let first_blocking_square = match direction {
    NW | NE => Square::first(blockers),
    SW | SE => Square::last(blockers),
};

The engine uses the same logic for generating rook attacks, and by combining this and bishop attacks together we get queen attacks.

This is the simplest method for calculating sliding piece attacks. There are several much more efficient methods, but I wanted to keep things simple to start with.

Non-sliding Piece Attacks

In contrast, pawns, knights and kings only attack a fixed set of their surrounding squares, and determining which ones is much easier because we don’t have to deal with blocking pieces.

We can use something like this to generate a bitboard of squares attacked by the king:

let king_attacks: u64 =
      king_square << 8  // north
    | king_square >> 8  // south

    | king_square << 1  // east
    | king_square >> 1  // west

    | king_square << 7  // north-west
    | king_square << 9  // north-east

    | king_square >> 7  // south-east
    | king_square >> 9; // south-west

Bear in mind that left-shifting moves a square towards the top of the board and right-shifting moves it towards the bottom:

By putting this calculation in a loop we can create a bitboard of king attacks for every square. However, we have to be careful when the king is not in the middle of the board. Consider what might happen if the king was on the edge:

It looks like the king attacks several squares on the other side of the board, but this is of course not true. We can avoid pieces wrapping around the edges by applying a bit-mask to the king square before shifting it:

const NOT_FILE_A: u64 = !0x0101_0101_0101_0101;
const NOT_FILE_H: u64 = !0x8080_8080_8080_8080;
NOT_FILE_A mask
NOT_FILE_H mask

Now, when applying a shift that moves the king towards the left edge we should first mask it with NOT_FILE_A. Similarly, we mask it with NOT_FILE_H when shifting towards the right edge.

Here’s how the logic above is modified to use these masks and populate a lookup table:

let mut KING_ATTACKS = [0; 64];

for square_index in 0..64 {
    let square = 1 << square_index;

    KING_ATTACKS[square_index] =
          square << 8
        | square >> 8

        | (square & NOT_FILE_H) << 1
        | (square & NOT_FILE_A) >> 1

        | (square & NOT_FILE_A) << 7
        | (square & NOT_FILE_H) << 9

        | (square & NOT_FILE_H) >> 7
        | (square & NOT_FILE_A) >> 9;
}

Notice that we don’t need a mask when shifting by 8 bits. This is because wrapping around the top or bottom edges of the board doesn’t actually occur with this bitboard layout. Right-shifting the e1 square by 8 simply means the bit disappears and no wrapping occurs.

We can now trivially write a function for finding king attacks:

fn get_king_attacks(square: Square) -> u64 {
    KING_ATTACKS[square]
}

Pawn and knight attack tables are mostly the same, except with different bitwise shifts and masks. Check out the source to see how knight attacks are calculated, for example.

At the beginning of this section we saw the get_attacks() function, which lives at the heart of the move generator. Here it is again:

fn get_attacks(piece: Piece, square: Square, board: &Board) -> u64 {
    match piece {
        WP | BP => get_pawn_attacks(...),
        WN | BN => get_knight_attacks(...),
        WB | BB => get_bishop_attacks(...),
        WR | BR => get_rook_attacks(...),
        WQ | BQ => get_bishop_attacks(...) | get_rook_attacks(...),
        WK | BK => get_king_attacks(...),
    }
}

With the methods described above we can implement all the “piece attack” functions and this gets us most of the way to generating all the moves in a position.

Other types of move

There’s a few types of move that require a little more than just the attack squares to find:

Pawn advances

Depending on which colour is advancing, the pawn’s square index shifts by +/– 8:

impl Square {
    fn advance(&self, colour: Colour) -> Self {
        match colour {
            White => Self(self.0 + 8),
            Black => Self(self.0 - 8),
        }
    }
}

Expanding on this, given a square and a colour, the following function checks if the square one ahead is empty. If the square is on the pawn’s start rank and the square one ahead is empty, it also checks the square two ahead. The result is a bitboard with at most two 1-bits.

fn get_pawn_advances(square: Square, colour: Colour, board: &Board) -> u64 {
    let one_ahead = square.advance(colour);

    if board.has_piece_at(one_ahead) {
        return 0;
    }

    if square.rank() != PAWN_START_RANKS[colour] {
        return one_ahead.u64();
    }

    let two_ahead = one_ahead.advance(colour);

    if board.has_piece_at(two_ahead) {
        return one_ahead.u64();
    }

    one_ahead.u64() | two_ahead.u64()
}

Since the function produces a bitboard we can just merge it with the attacks bitboard for the pawn to leverage the behaviour of converting the bitboard to moves later.

let mut to_squares = get_attacks(piece, from_square, board);

if piece.is_pawn() {
    to_squares |= get_pawn_advances(from_square, colour_to_move, board);
}
OR
=

Pawn promotions

If a pawn advance would move the pawn to the 1st or 8th ranks then there are four moves available, one for each promotion piece. Checking if a square is on the back rank is a simple bitwise AND operation:

const BACK_RANKS: u64 = 0xFF00_0000_0000_00FF;

impl Square {
    fn is_back_rank(&self) -> bool {
        *self & BACK_RANKS != 0
    }
}

Now, when converting a pawn’s bitboard of available squares into a list of moves, we can determine whether the four promotion moves should be added:

for to_square in to_squares {
    // ...

    if piece.is_pawn() && to_square.is_back_rank() {
        moves.push(...);
    }

    // ...
}

En-passant captures

Recall that the Position structure has an en_passant_square property. This gets set when applying a double pawn advance and reset when any other move occurs. En-passant is only possible if the property is set and an opponent pawn attacks that square. We can see if this is the case by finding the pawn attacks from the perspective of the en-passant square and checking for any overlap with the colour-to-move’s pawns:

fn get_en_passant_attacks(
    en_passant_square: Square,
    colour: Colour,
    board: &Board
) -> u64 {
    get_pawn_attacks(en_passant_square, colour.flip(), board)
        & board.pieces(Piece::pawn(colour))
}

We’ll see how to convert the result of this function to a list of moves later.

Castling

The last type of special move is castling. The colour-to-move is allowed to castle if:

  1. the king and rook have not moved from their start squares,
  2. the squares between the king and rook are not occupied,
  3. the square the king will pass through is not attacked, and
  4. the colour-to-move is not in check.

1. The king and rook have not moved from their start squares

The Position structure has a castling_rights property. Starting from a value that indicates full rights to castle, this property is updated to reflect which rights are lost as the pieces move:

impl Position {
    fn do_move(&mut self, mv: &Move) {
        let piece = self.board.piece_at(mv.from);

        if piece.is_king() {
            self.castling_rights
                .remove_for_colour(self.colour_to_move);
        }

        // Maybe a rook left its start square?
        if mv.from.is_corner() {
            self.castling_rights
                .remove_for_square(mv.from);
        }

        // Maybe a rook was captured?
        if mv.to.is_corner() {
            self.castling_rights
                .remove_for_square(mv.to);
        }

        // ...
    }
}

2. The squares between the king and rook are not occupied

The squares to check depend on which colour is castling and in which direction:

const WHITE_KING_SIDE_PATH: u64 = F1 | G1;
const WHITE_QUEEN_SIDE_PATH: u64 = B1 | C1 | D1;

const BLACK_KING_SIDE_PATH: u64 = F8 | G8;
const BLACK_QUEEN_SIDE_PATH: u64 = B8 | C8 | D8;

With these constants, we can use familiar bitwise logic to determine if there are pieces on these squares:

impl Board {
    fn has_occupancy_at(&self, squares: u64) -> bool {
        self.occupancy() & squares != 0
    }
}

3. The square the king will pass through is not attacked

The king can’t castle through check, so we need to determine if any of the opponent’s pieces attack the square next to the king. For example, black may still have the right to castle in the following position, but it would not be legal because the white bishop on c5 blocks it:

In this scenario, if we calculate the attacks for a bishop on f8, we would find an overlap with the c5 square:

This must mean the bishop on c5 attacks f8 and black can’t castle. We can use this idea to write a function that takes a square and calculates a bitboard of all its attacking pieces:

fn get_attackers(square: Square, colour: Colour, board: &Board) -> u64 {
    let bishop_attacks = get_bishop_attacks(...);
    let rook_attacks = get_rook_attacks(...);
    let queen_attacks = bishop_attacks | rook_attacks;

      (board.pieces(Piece::pawn(colour))   & get_pawn_attacks(...))
    | (board.pieces(Piece::knight(colour)) & get_knight_attacks(...))
    | (board.pieces(Piece::bishop(colour)) & bishop_attacks)
    | (board.pieces(Piece::rook(colour))   & rook_attacks)
    | (board.pieces(Piece::queen(colour))  & queen_attacks)
    | (board.pieces(Piece::king(colour))   & get_king_attacks(...))
}

Finally, we can use this to check if a square is attacked by a given colour:

fn is_attacked(square: Square, colour: Colour, board: &Board) -> bool {
    get_attackers(square, colour, board) != 0
}

4. The colour-to-move is not in check

This simply reuses the is_attacked() function to see if the king square is attacked:

fn is_in_check(colour: Colour, board: &Board) -> bool {
    let king_square = Square::first(board.pieces(Piece::king(colour)));

    is_attacked(king_square, colour.flip(), board)
}

All that’s left to do is to put these rules into a function that returns a bitboard of squares the king can move to when castling:

fn get_castling(rights: CastlingRights, colour: Colour, board: &Board) -> u64 {
    if is_in_check(colour, board) {
        return 0;
    }

    match colour {
        White => get_white_castling(rights, board),
        Black => get_black_castling(rights, board),
    }
}

fn get_white_castling(rights: CastlingRights, board: &Board) -> u64 {
    let mut white_castling = 0;

    if rights.has(WhiteKing)
        && !board.has_occupancy_at(WHITE_KING_SIDE_PATH)
        && !is_attacked(F1, Black, board)
    {
        white_castling |= G1;
    }

    if rights.has(WhiteQueen)
        && !board.has_occupancy_at(WHITE_QUEEN_SIDE_PATH)
        && !is_attacked(D1, Black, board)
    {
        white_castling |= C1;
    }

    white_castling
}

fn get_black_castling(rights: CastlingRights, board: &Board) -> u64 {
    // ...
}

Similarly to pawn advances, the resultant bitboard for castling moves can be merged with the attack squares for the king:

let mut to_squares = get_attacks(piece, from_square, board);

if piece.is_king() {
    to_squares |= get_castling(pos.castling_rights, colour_to_move, board);
}
OR
=

Converting Bitboards to Moves

By this point we can generate a bitboard containing all the available squares for any piece on the board. Now we just need to convert these bitboards into a list of moves and we’re done!

Suppose we have the following bitboard for the bishop’s available moves:

A specialised Square constructor allows me to get the first square in the bitboard and unset its corresponding bit:

impl Square {
    fn next(squares: &mut u64) -> Self {
        let square = Self::first(*squares);
        *squares ^= square;
        square
    }
}

I wouldn’t usually be a fan of such an impurity*, but I’ve found this to be extremely useful for iterating over all the 1-bits in a loop:

let mut moves = vec![];
let mut to_squares = ...;

while to_squares != 0 {
    let to_square = Square::next(&mut to_squares);
    moves.push(Move::new(from_square, to_square));
}

Of course, this isn’t only used to convert bitboards to moves. This technique can be used wherever there’s a need to iterate over a bitboard, as we’ll see below…

* Perhaps the only saving grace is that at least Rust makes clear what’s going on with its &mut notation.

Putting Everything Together

For completeness, here’s the whole move generation function:

fn generate_moves(pos: &Position) -> Vec<Move> {
    let mut moves = vec![];
    let board = pos.board;
    let colour = pos.colour_to_move;

    for piece in Piece::pieces_by_colour(colour) {
        let mut pieces = board.pieces(piece);

        while pieces != 0 {
            let from_square = Square::next(&mut pieces);

            let mut to_squares = !board.pieces_by_colour(colour)
                & get_attacks(piece, from_square, board);

            if piece.is_pawn() {
                to_squares |= get_pawn_advances(from_square, colour, board);
            } else if piece.is_king() {
                to_squares |= get_castling(pos.castling_rights, colour, board);
            }

            while to_squares != 0 {
                let to_square = Square::next(&mut to_squares);

                if piece.is_pawn() && to_square.is_back_rank() {
                    for piece in Piece::promotions(colour) {
                        moves.push(
                            Move::promotion(from_square, to_square, piece)
                        );
                    }

                    continue;
                }

                moves.push(Move::new(from_square, to_square));
            }
        }
    }

    if let Some(en_passant_square) = pos.en_passant_square {
        let mut from_squares =
            get_en_passant_attacks(en_passant_square, colour, board);

        while from_squares != 0 {
            let from_square = Square::next(&mut from_squares);
            moves.push(Move::new(from_square, en_passant_square));
        }
    }

    moves
}

Verifying Correctness

Writing a move generator has been one of the most error-prone tasks I’ve worked on. With so many nuances to the rules it can be tricky to implement a completely bug-free program, so having a reliable test suite is crucial.

Besides having a suite of TDD-style unit tests still hanging around from the early phases of development, my general approach to testing has been to rely on perft, a method used to give an indication of correctness for chess move generators.

The basic idea is to compare the number of generated moves with the known correct value for a given position. For example, there are 20 moves in the start position for white, and for each of those there are 20 replies for a total of 400 moves after 2 turns. The following table shows the number of possible moves as the number of turns increases:

Depth Move count
1 20
2 400
3 8,902
4 197,281
5 4,865,609
6 119,060,324

I wanted to be able to write a test like this:

#[test]
fn perft_start_position() {
    assert_eq!(perft(Position::startpos(), 6), 119_060_324);
}

To support this we need a function that counts moves in a position recursively:

fn perft(pos: &mut Position, depth: u8) -> u32 {
    if depth == 0 {
        return 1;
    }

    let mut count = 0;

    for mv in generate_moves(pos) {
        pos.do_move(mv);

        if pos.is_legal() {
            count += perft(pos, depth - 1);
        }

        pos.undo_move(mv);
    }

    count
}

Note that I’m checking the position is legal after applying each move because the engine only generates pseudo-legal moves, meaning that a move could leave the player in check, which would not be legal.

It was extremely useful to have tests like this during development. They helped me find so many unusual bugs that I never would have noticed otherwise. For example, here’s a position that highlights a common gotcha:

In this position, if black plays a double pawn advance on the c-file then technically en-passant is available, and at a glance it looks like white’s b-pawn could capture en-passant and move to c6. However, this would leave the white king in check by the rook on h5, making the move illegal.

I currently verify this position to depth 7:

#[test]
fn perft_position_3() {
    assert_perft_for_fen(
        "8/2p5/3p4/KP5r/1R3p1k/8/4P1P1/8 w - - 0 1",
        7,
        178_633_661
    );
}

Another position I found useful was this one:

White is in check so the number of possible moves should be limited, and after white’s turn black has 8 ways to promote the pawn on b2. Black has castling rights on both sides of the board but can only castle queen-side because of white’s knight and bishop blocking the king-side path. This position is great for testing lots of the subtle details in the rules.

I currently verify this position to depth 6:

#[test]
fn perft_position_4() {
    assert_perft_for_fen(
        "r3k2r/Pppp1ppp/1b3nbN/nP6/BBP1P3/q4N2/Pp1P2PP/R2Q1RK1 w kq - 0 1",
        6,
        706_045_033,
    );
}

More positions and their move counts are available on the Chess Programming Wiki.

I combine these positions with several others to gain a level of confidence I don’t think would be achievable with other forms of testing. The only drawback is that they’re slow, so I only run these on CI. For development, I run a reduced suite of perft tests that verify lower depths in order to keep the feedback loop fast.

Performance Considerations

Once I’ve got the engine working in its entirety I’ll revisit some of these areas to work on performance. Here’s a few ideas I have in mind:

It will be interesting to benchmark before and after these changes. However, bear in mind that the biggest performance gains come from improving positional evaluation so that the engine can cut off large branches of the search tree, since this avoids having to generate moves for so many positions in the first place.

Onwards: Evaluation

In the next post I’ll write about how the engine assigns each chess position a score, the fundamental metric used to find the best move.

♟️