From 942c758e151450576d98f362ad29a7075974c8d5 Mon Sep 17 00:00:00 2001
From: Eryn Wells <eryn@erynwells.me>
Date: Fri, 2 Feb 2024 08:05:37 -0800
Subject: [PATCH] [position] Calculate legal moves based on whether the king is
 in check

Use CheckingPieces to determine which and how many pieces check the king.

- If there are no checks, proceed with move generation as normal.
- If there is one checking piece, calculate push and capture masks and use those
   to generate legal moves.
- If there are more than one checking pieces, the only legal moves are king
  moves. Indicate this by setting the push and capture masks to EMPTY.
---
 position/src/position/position.rs | 22 ++++++++++++++++++----
 1 file changed, 18 insertions(+), 4 deletions(-)

diff --git a/position/src/position/position.rs b/position/src/position/position.rs
index 774fc26..f0f8414 100644
--- a/position/src/position/position.rs
+++ b/position/src/position/position.rs
@@ -2,11 +2,11 @@
 
 use super::{flags::Flags, piece_sets::PieceBitBoards, Pieces};
 use crate::{
-    move_generator::{MoveSet, Moves},
+    check::{self, CheckingPieces},
+    move_generator::Moves,
     position::DiagramFormatter,
     r#move::Castle,
     sight::SightExt,
-    Move,
 };
 use chessfriend_bitboard::BitBoard;
 use chessfriend_core::{Color, Piece, PlacedPiece, Shape, Square};
@@ -121,8 +121,22 @@ impl Position {
     }
 
     pub fn moves(&self) -> &Moves {
-        self.moves
-            .get_or_init(|| Moves::new(self, self.color_to_move, BitBoard::FULL, BitBoard::FULL))
+        self.moves.get_or_init(|| {
+            let checking_pieces = self.checking_pieces();
+            match checking_pieces.count() {
+                // Normal, unrestricted move generation
+                0 => Moves::new(self, self.color_to_move, BitBoard::FULL, BitBoard::FULL),
+                1 => {
+                    // Calculate push and capture masks for checking piece. Moves are restricted to those that intersect those masks.
+                    let capture_mask = checking_pieces.capture_mask();
+                    let push_mask =
+                        checking_pieces.push_mask(self.king_bitboard(self.color_to_move));
+                    Moves::new(self, self.color_to_move, capture_mask, push_mask)
+                }
+                // With more than one checking piece, the only legal moves are king moves.
+                _ => Moves::new(self, self.color_to_move, BitBoard::EMPTY, BitBoard::EMPTY),
+            }
+        })
     }
 
     /// Return a BitBoard representing the set of squares containing a piece.