From 37260b4d38bc0f4c939aa309ecb429e4e15b3aa8 Mon Sep 17 00:00:00 2001
From: dogeystamp <dogeystamp@disroot.org>
Date: Sun, 20 Oct 2024 13:09:04 -0400
Subject: [PATCH] feat: king queen rook bishop slider movegen

---
 src/lib.rs     | 11 +++++++
 src/movegen.rs | 84 +++++++++++++++++++++++++++++++++++++++++++++-----
 2 files changed, 87 insertions(+), 8 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 1a90d16..bf44feb 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -6,6 +6,8 @@ use std::str::FromStr;
 pub mod fen;
 pub mod movegen;
 
+use fen::FromFen;
+
 const BOARD_WIDTH: usize = 8;
 const BOARD_HEIGHT: usize = 8;
 const N_SQUARES: usize = BOARD_WIDTH * BOARD_HEIGHT;
@@ -418,6 +420,11 @@ impl BoardState {
         &mut self.players[col as usize]
     }
 
+    /// Get immutable reference to a player.
+    fn pl(&self, col: Color) -> &Player {
+        &self.players[col as usize]
+    }
+
     /// Create a new piece in a location.
     fn set_piece(&mut self, idx: Square, pc: ColPiece) {
         let pl = self.pl_mut(pc.col);
@@ -539,5 +546,9 @@ mod tests {
                 assert_eq!(squares[i], sq)
             }
         }
+
+        let board = BoardState::from_fen("8/4p3/1q1Q1p2/4p3/1p1r4/8/8/8 w - - 0 1").unwrap();
+        let white_queens = board.pl(Color::White).board(Piece::Queen).into_iter().collect::<Vec<Square>>();
+        assert_eq!(white_queens, vec![Square::from_str("d6").unwrap()])
     }
 }
diff --git a/src/movegen.rs b/src/movegen.rs
index 2bd134b..84dcf57 100644
--- a/src/movegen.rs
+++ b/src/movegen.rs
@@ -354,7 +354,7 @@ fn move_slider(
                     move_type: MoveType::Normal,
                 });
 
-                // Stop at other pieces.
+                // stop at other pieces.
                 if let Some(_cap_pc) = board.get_piece(dest) {
                     break;
                 }
@@ -374,10 +374,23 @@ impl PseudoMoveGen for BoardState {
 
     fn gen_pseudo_moves(self) -> Self::MoveIterable {
         let mut ret = Vec::new();
-        for pl in self.players {
-            for sq in pl.board(Piece::Rook).into_iter() {
-                move_slider(&self, sq, &mut ret, SliderDirection::Straight, true);
-            }
+        let pl = self.pl(self.turn);
+        macro_rules! squares {
+            ($pc: ident) => {
+                pl.board(Piece::$pc).into_iter()
+            };
+        }
+        for sq in squares!(Rook) {
+            move_slider(&self, sq, &mut ret, SliderDirection::Straight, true);
+        }
+        for sq in squares!(Bishop) {
+            move_slider(&self, sq, &mut ret, SliderDirection::Diagonal, true);
+        }
+        for sq in squares!(Queen) {
+            move_slider(&self, sq, &mut ret, SliderDirection::Star, true);
+        }
+        for sq in squares!(King) {
+            move_slider(&self, sq, &mut ret, SliderDirection::Star, false);
         }
         ret
     }
@@ -398,6 +411,7 @@ mod tests {
     #[test]
     fn test_slider_movegen() {
         let test_cases = [
+            // rook test
             (
                 // start position
                 "8/8/8/8/8/8/8/R7 w - - 0 1",
@@ -413,16 +427,70 @@ mod tests {
                     MoveType::Normal,
                 )],
             ),
+            // king against the boundary
+            (
+                "3K4/4p3/1q3p2/4p3/1p1r4/8/8/8 w - - 0 1",
+                vec![("d8", vec!["c8", "c7", "d7", "e7", "e8"], MoveType::Normal)],
+            ),
+            // king test
+            (
+                "8/4p3/1q1K1p2/4p3/1p1r4/8/8/8 w - - 0 1",
+                vec![(
+                    "d6",
+                    vec!["c7", "c6", "c5", "d7", "d5", "e7", "e6", "e5"],
+                    MoveType::Normal,
+                )],
+            ),
+            // queen test
+            (
+                "8/4p3/1q1Q1p2/4p3/1p1r4/8/8/8 w - - 0 1",
+                vec![(
+                    "d6",
+                    vec![
+                        "d5", "d4", "d7", "d8", "e7", "c5", "b4", "e6", "f6", "c6", "b6", "e5",
+                        "c7", "b8",
+                    ],
+                    MoveType::Normal,
+                )],
+            ),
+            // rook test (again)
             (
                 "8/1p6/8/1R2p3/1p6/8/8/8 w - - 0 1",
                 vec![(
                     "b5",
-                    vec![
-                        "b6", "b7", "b4", "a5", "c5", "d5", "e5",
-                    ],
+                    vec!["b6", "b7", "b4", "a5", "c5", "d5", "e5"],
                     MoveType::Normal,
                 )],
             ),
+            // bishop test
+            (
+                "8/4p3/3B4/4p3/1p6/8/8/8 w - - 0 1",
+                vec![(
+                    "d6",
+                    vec!["e5", "e7", "c5", "b4", "c7", "b8"],
+                    MoveType::Normal,
+                )],
+            ),
+            // black test
+            (
+                "8/3b4/2R1R3/1Q6/1RqRrR2/1QQ5/2R1R3/k7 b - - 0 1",
+                vec![
+                    ("a1", vec!["a2", "b2", "b1"], MoveType::Normal),
+                    (
+                        "c4",
+                        vec![
+                            "b3", "b4", "b5", "c3", "c5", "c6", "d3", "e2", "d4", "d5", "e6",
+                        ],
+                        MoveType::Normal,
+                    ),
+                    (
+                        "e4",
+                        vec!["d4", "f4", "e3", "e2", "e5", "e6"],
+                        MoveType::Normal,
+                    ),
+                    ("d7", vec!["c6", "e6", "c8", "e8"], MoveType::Normal),
+                ],
+            ),
         ];
 
         for (fen, expected) in test_cases {