1 files changed, 146 insertions, 38 deletions

diff --git a/src/moves.cpp b/src/moves.cpp
index 8a7feb5..2663ea0 100644
--- a/src/moves.cpp
+++ b/src/moves.cpp
@@ -103,8 +103,8 @@ void GeneratePawnMoves(Game *b, int from, std::vector<Move> &moves) {
 
 std::vector<Move> GetPseudoLegalMoves(Game *b) {
   std::vector<Move> moves;
-  moves.reserve(20);
-  
+  moves.reserve(40);
+
   constexpr std::array<int, 4> rook_Moves{-1, 1, 8, -8};
   constexpr std::array<int, 4> bishop_Moves{-9, 9, -7, 7};
 
@@ -141,57 +141,59 @@ std::vector<Move> GetPseudoLegalMoves(Game *b) {
 
   return moves;
 }
-bool IsSquareAttacked(Game *board, Position square, bool white) {
-  Game temp = *board;
+std::vector<Move> GetLegalMoves(Game *g) {
+  std::vector<Move> moves = GetPseudoLegalMoves(g);
+  std::vector<Move> legalMoves;
 
-  // Generate moves for the attacking side
-  temp.turn = white;
+  for (Move move : moves) {
+    UndoMove undo = MakeMove(move, g);
 
-  auto moves = GetPseudoLegalMoves(&temp);
+    bool legal = true;
 
-  for (const Move &move : moves) {
-    if (move.To == square)
-      return true;
-  }
+    Position kingPosition = FindKing(g, !g->turn);
 
-  return false;
-}
-bool IsPieceTypeAttacked(Game *board, PieceType type, bool white) {
-  Game temp = *board;
+    // opponent attacks our king?
+    if (IsSquareAttacked(g, kingPosition, g->turn)) {
+      legal = false;
+    }
 
-  // Generate moves for the attacking side
-  temp.turn = white;
+    UnMakeMove(undo, g);
+
+    if (legal) {
+      legalMoves.push_back(move);
+    }
+  }
+
+  if (legalMoves.empty()) {
+    Position kingPosition = FindKing(g, g->turn);
 
-  auto moves = GetPseudoLegalMoves(&temp);
+    bool check = IsSquareAttacked(g, kingPosition, !g->turn);
 
-  for (const Move &move : moves) {
-    if (temp.pieces[PositionToIndex(move.To)].type == type)
-      return true;
+    if (check) {
+      g->state = g->turn ? BLACK_WON : WHITE_WON;
+    } else {
+      g->state = STALEMATE;
+    }
   }
 
-  return false;
-}
-std::vector<Move> GetLegalMoves(Game *b) {
-  std::vector<Move> moves = GetPseudoLegalMoves(b);
-  std::vector<Move> legalMove;
   for (Move move : moves) {
     bool isLegal = true;
-    UndoMove undo = MakeMove(move, b);
-    auto opponentResponse = GetPseudoLegalMoves(b);
+    UndoMove undo = MakeMove(move, g);
+    auto opponentResponse = GetPseudoLegalMoves(g);
     for (Move move : opponentResponse) {
-      if (b->pieces[PositionToIndex(move.To)].type == KING) {
+      if (g->pieces[PositionToIndex(move.To)].type == KING) {
         isLegal = false;
       }
     }
-    UnMakeMove(undo, b);
+    UnMakeMove(undo, g);
     if (isLegal) {
-      legalMove.push_back(move);
+      legalMoves.push_back(move);
     };
   }
-  if (legalMove.size() == 0) {
+  if (legalMoves.size() == 0) {
     bool isCheck = false;
 
-    Game testBoard = *b;
+    Game testBoard = *g;
     testBoard.turn = !testBoard.turn;
     auto opponentResponse = GetPseudoLegalMoves(&testBoard);
     for (Move move : opponentResponse) {
@@ -201,18 +203,124 @@ std::vector<Move> GetLegalMoves(Game *b) {
     }
 
     if (isCheck) {
-      if (b->turn) {
-        b->state = BLACK_WON;
+      if (g->turn) {
+        g->state = BLACK_WON;
       } else {
-        b->state = WHITE_WON;
+        g->state = WHITE_WON;
       };
     } else {
-      b->state = STALEMATE;
+      g->state = STALEMATE;
     }
   }
-  return legalMove;
+  return legalMoves;
 }
 
+bool IsSquareAttacked(Game *board, Position square, bool white) {
+  Game temp = *board;
+  temp.turn = white;
+
+  // Moves
+  constexpr std::array<int, 4> bishopMoves{-9, 9, -7, 7};
+  constexpr std::array<int, 4> rookMoves{-8, 8, -1, 1};
+
+  int target = PositionToIndex(square);
+
+  // Pawns (custom because pawn moves != pawn attacks)
+  int pawnDirection = white ? -1 : 1;
+
+  int pawnRank = square.rank - pawnDirection;
+
+  if (pawnRank >= 0 && pawnRank < 8) {
+    for (int fileOffset : {-1, 1}) {
+      int pawnFile = square.file + fileOffset;
+
+      if (pawnFile < 0 || pawnFile >= 8)
+        continue;
+
+      Piece p = temp.pieces[pawnRank * 8 + pawnFile];
+
+      if (p.type == PAWN && p.color == white)
+        return true;
+    }
+  }
+
+  // knights
+  for (int i = 0; i < 64; i++) {
+    Piece p = temp.pieces[i];
+
+    if (p.type == KNIGHT && p.color == white) {
+      std::vector<Move> moves;
+      GenerateKnightMoves(&temp, i, moves);
+
+      for (Move m : moves) {
+        if (PositionToIndex(m.To) == target)
+          return true;
+      }
+    }
+  }
+  for (int i = 0; i < 64; i++) {
+    Piece p = temp.pieces[i];
+
+    if (p.type == BISHOP && p.color == white) {
+      std::vector<Move> moves;
+      GenerateSlidingMoves(&temp, i, bishopMoves, moves);
+
+      for (Move m : moves) {
+        if (PositionToIndex(m.To) == target)
+          return true;
+      }
+    }
+  }
+
+  // rook
+  for (int i = 0; i < 64; i++) {
+    Piece p = temp.pieces[i];
+
+    if (p.type == ROOK && p.color == white) {
+      std::vector<Move> moves;
+      GenerateSlidingMoves(&temp, i, rookMoves, moves);
+
+      for (Move m : moves) {
+        if (PositionToIndex(m.To) == target)
+          return true;
+      }
+    }
+  }
+
+  // queen
+  for (int i = 0; i < 64; i++) {
+    Piece p = temp.pieces[i];
+
+    if (p.type == QUEEN && p.color == white) {
+      std::vector<Move> moves;
+
+      GenerateSlidingMoves(&temp, i, bishopMoves, moves);
+      GenerateSlidingMoves(&temp, i, rookMoves, moves);
+
+      for (Move m : moves) {
+        if (PositionToIndex(m.To) == target)
+          return true;
+      }
+    }
+  }
+
+  // --- Kings ---
+  for (int i = 0; i < 64; i++) {
+    Piece p = temp.pieces[i];
+
+    if (p.type == KING && p.color == white) {
+      std::vector<Move> moves;
+      GenerateKingMoves(&temp, i, moves);
+
+      for (Move m : moves) {
+        if (PositionToIndex(m.To) == target)
+          return true;
+      }
+    }
+  }
+
+  return false;
+}
 Position IndexToPosition(int i) {
   uint8_t rank = i / 8; // 0-7
   uint8_t file = i % 8; // 0-7