From 96848371d64d6562e900e272862b7adf60311423 Mon Sep 17 00:00:00 2001
From: Tait Hoyem <44244401+TTWNO@users.noreply.github.com>
Date: Sun, 28 Apr 2019 03:02:19 +0000
Subject: [PATCH] Fix an error where castling is allowed even when it a) puts
 the king in check or b) puts the king in check on one of the squares on the
 way to the final position

---
 src/all_moves_functions.cpp | 50 +------------------------------------
 src/functions.cpp           | 18 ++++++++-----
 tests/main.cpp              |  6 +++--
 tests/valid_moves.h         | 18 +++++++++++++
 4 files changed, 35 insertions(+), 57 deletions(-)

diff --git a/src/all_moves_functions.cpp b/src/all_moves_functions.cpp
index 7ede64e..200e8ff 100644
--- a/src/all_moves_functions.cpp
+++ b/src/all_moves_functions.cpp
@@ -129,19 +129,6 @@ void _get_all_moves_rook(int pos, std::vector<int> *pns, std::array<PieceType, 1
 			_add_if_not_blocked(pos+(rk_off*offset), pos, pns, board, pc, rc, not_blocked);
 		}
 	}
-	
-	/*
-	bool* not_blocked_addx = new bool(true);
-	bool* not_blocked_minx = new bool(true);
-	bool* not_blocked_addy = new bool(true);
-	bool* not_blocked_miny = new bool(true);
-	for (int offset=1; offset<8; offset++){
-		_add_if_not_blocked(x+offset, y, pns, board, pc, rc, not_blocked_addx);
-		_add_if_not_blocked(x, y+offset, pns, board, pc, rc, not_blocked_addy);
-		_add_if_not_blocked(x-offset, y, pns, board, pc, rc, not_blocked_minx);
-		_add_if_not_blocked(x, y-offset, pns, board, pc, rc, not_blocked_miny);
-	}
-	*/
 }
 
 void _get_all_moves_bishop(int pos, std::vector<int> *pns, std::array<PieceType, 120>* board, Color pc, Color rc){
@@ -151,37 +138,12 @@ void _get_all_moves_bishop(int pos, std::vector<int> *pns, std::array<PieceType,
 			_add_if_not_blocked(pos+(bs_off*offset), pos, pns, board, pc, rc, not_blocked);
 		}
 	}
-	/*
-	bool* not_blocked_addx_addy = new bool(true);
-	bool* not_blocked_addx_miny = new bool(true);
-	bool* not_blocked_minx_addy = new bool(true);
-	bool* not_blocked_minx_miny = new bool(true);
-	
-	for (int offset=1; offset<8; offset++){
-		int xpoff = x+offset;
-		int ypoff = y+offset;
-		int xnoff = x-offset;
-		int ynoff = y-offset;
-		_add_if_not_blocked(xpoff, ypoff, pns, board, pc, rc, not_blocked_addx_addy);
-		_add_if_not_blocked(xpoff, ynoff, pns, board, pc, rc, not_blocked_addx_miny);
-		_add_if_not_blocked(xnoff, ypoff, pns, board, pc, rc, not_blocked_minx_addy);
-		_add_if_not_blocked(xnoff, ynoff, pns, board, pc, rc, not_blocked_minx_miny);	
-	}
-	*/
 }
+
 void _get_all_moves_knight(int pos, std::vector<int> *pns, std::array<PieceType, 120>* board, Color pc, Color rc){
 	for (int kn_off : KNIGHT_PIECE_OFFSETS){
 		_add_if_not_blocked(pos+kn_off, pos, pns, board, pc, rc);
 	}
-	/*
-	for (int xo=1;xo<=2;xo++){
-		int yo=(xo==1)?2:1;
-		_add_if_not_blocked(x+xo, y+yo, pns, board, pc, rc);
-		_add_if_not_blocked(x-xo, y+yo, pns, board, pc, rc);
-		_add_if_not_blocked(x+xo, y-yo, pns, board, pc, rc);
-		_add_if_not_blocked(x-xo, y-yo, pns, board, pc, rc);
-	}
-	*/
 }
 
 void _get_all_moves_king(int pos, std::vector<int> *pns, std::array<PieceType, 120>* board, Color pc, Color rc, int castle_perms){
@@ -212,16 +174,6 @@ void _get_all_moves_king(int pos, std::vector<int> *pns, std::array<PieceType, 1
 			_king_add_if_not_blocked(pos-2, pos, pns, board, pc, rc, left_castle, 1);
 		}
 	}
-	/*
-	_add_if_not_blocked(x+1, y+1, pns, board, pc, rc);
-	_add_if_not_blocked(x+1, y-1, pns, board, pc, rc);
-	_add_if_not_blocked(x-1, y+1, pns, board, pc, rc);
-	_add_if_not_blocked(x-1, y-1, pns, board, pc, rc);
-	_add_if_not_blocked(x, y+1, pns, board, pc, rc);
-	_add_if_not_blocked(x, y-1, pns, board, pc, rc);
-	_add_if_not_blocked(x+1, y, pns, board, pc, rc);
-	_add_if_not_blocked(x-1, y, pns, board, pc, rc);
-	*/
 }
 
 void _get_all_moves_pawn(int pos, std::vector<int> *pns, std::array<PieceType, 120>* board, Color pc, Color rc, int en_passant){
diff --git a/src/functions.cpp b/src/functions.cpp
index ee1fdc1..c275e44 100644
--- a/src/functions.cpp
+++ b/src/functions.cpp
@@ -133,25 +133,31 @@ void filter_checked_moves(PieceType pt, std::array<PieceType, 120> *board, std::
 	for (auto p_pn= pns->begin(); p_pn!=pns->end();){
 		if (get_castle_flag(*p_pn) == 1){
 			// If moved left
-			// B1 - A1 = -1
-			// A1 - B1 = +1
+			// Queenside
 			if ((get_from_sq(*p_pn) - get_to_sq(*p_pn)) > 0){
 				int right_move = make_move(get_from_sq(*p_pn), get_to_sq(*p_pn)+1);
+				int full_move = make_move(get_from_sq(*p_pn), get_to_sq(*p_pn));
 				auto right_board = dumb_move(right_move, *board);
-				if (is_checked(get_to_sq(*p_pn)+1, right_board)){
-					p_pn = pns->erase(p_pn);	
+				auto full_board = dumb_move(full_move, *board);
+				if (is_checked(get_to_sq(*p_pn)+1, right_board) ||
+					is_checked(get_to_sq(*p_pn), full_board)){
+					p_pn = pns->erase(p_pn);
 				} else {
 					++p_pn;
 					break;
 				}
+			// Kingside
 			} else {
 				int left_move = make_move(get_from_sq(*p_pn), get_to_sq(*p_pn)-1);
+				int full_move = make_move(get_from_sq(*p_pn), get_to_sq(*p_pn));
 				auto left_board = dumb_move(left_move, *board);
-				if (is_checked(get_to_sq(*p_pn)-1, left_board)){
+				auto full_board = dumb_move(full_move, *board);
+				if (is_checked(get_to_sq(*p_pn)-1, left_board) ||
+					is_checked(get_to_sq(*p_pn), full_board)){
 					p_pn = pns->erase(p_pn);
 				} else {
 					++p_pn;
-					continue;
+					break;
 				}
 			}
 		} else {
diff --git a/tests/main.cpp b/tests/main.cpp
index 08f9b82..7c28a2f 100644
--- a/tests/main.cpp
+++ b/tests/main.cpp
@@ -277,9 +277,11 @@ TEST_CASE("Test for promotions on empty squares", "[get_all_moves]"){
 }
 
 TEST_CASE("Tests for check on square of queenside capture", "[get_all_moves]"){
-	auto cannot_queenside = get_all_moves(CASTLE_CHECK1_POS, CASTLE_CHECK1_BOARD);
-	auto cannot_queenside2 = get_all_moves(CASTLE_CHECK2_POS, CASTLE_CHECK2_BOARD);
+	auto cannot_queenside = get_all_moves(CASTLE_CHECK1_POS, CASTLE_CHECK1_BOARD, true, 0, 0xF);
+	auto cannot_queenside2 = get_all_moves(CASTLE_CHECK2_POS, CASTLE_CHECK2_BOARD, true, 0, 0xF);
+	auto cannot_queenside3 = get_all_moves(CASTLE_CHECK3_POS, CASTLE_CHECK3_BOARD, true, 0, 0xF);
 
 	CHECK(get_notations(cannot_queenside, CASTLE_CHECK1_BOARD) == CASTLE_CHECK1_NOTATION);
 	CHECK(get_notations(cannot_queenside2, CASTLE_CHECK2_BOARD) == CASTLE_CHECK2_NOTATION);
+	CHECK(get_notations(cannot_queenside3, CASTLE_CHECK3_BOARD) == CASTLE_CHECK3_NOTATION);
 }
diff --git a/tests/valid_moves.h b/tests/valid_moves.h
index f515433..0b8fdfd 100644
--- a/tests/valid_moves.h
+++ b/tests/valid_moves.h
@@ -835,5 +835,23 @@ const std::array<PieceType, 120> CASTLE_CHECK2_BOARD = {
 const std::vector<std::string> CASTLE_CHECK2_NOTATION = {
 	"Kd8"
 };
+const int CASTLE_CHECK3_POS = E8;
+const std::array<PieceType, 120> CASTLE_CHECK3_BOARD = {
+	INV, INV, INV, INV, INV, INV, INV, INV, INV, INV, 	
+	INV, INV, INV, INV, INV, INV, INV, INV, INV, INV, 	
+	INV, B_ROOK, NONE, NONE, B_PAWN, B_KING, NONE, NONE, B_ROOK,  INV,	
+	INV, NONE, NONE, NONE, B_QUEEN, B_PAWN, B_PAWN, NONE, NONE,  INV,	
+	INV, NONE, NONE, NONE, NONE, NONE, NONE, NONE, NONE,  INV,	
+	INV, NONE, NONE, NONE, NONE, NONE, NONE, NONE, NONE, INV,
+	INV, NONE, NONE, NONE, NONE, NONE, NONE, W_ROOK, NONE,  INV,	
+	INV, NONE, NONE, NONE, NONE, NONE, NONE, NONE, NONE,  INV,	
+	INV, NONE, NONE, NONE, NONE, NONE, NONE, NONE, NONE,  INV,	
+	INV, NONE, NONE, NONE, NONE, NONE, NONE, W_KING, NONE, INV,
+	INV, INV, INV, INV, INV, INV, INV, INV, INV, INV, 	
+	INV, INV, INV, INV, INV, INV, INV, INV, INV, INV	
+};
+const std::vector<std::string> CASTLE_CHECK3_NOTATION = {
+	"Kf8"
+};