1 files changed, 72 insertions, 7 deletions

diff --git a/board.cpp b/board.cpp
index 0ac5348..b72ac5f 100644
--- a/board.cpp
+++ b/board.cpp
@@ -293,18 +293,76 @@ void Board::apply(const Move &mv) {
 	lastTarget = mv.to;
 }
 
-bool Board::isValid(const Move &mv) {
+bool Board::isValid(const Move &mv) const {
 	assert(mv.castle == 0);
 	if (colour(at(mv.from)) != onTurn) return false;
 	Board B(*this);
 	B.apply(mv);
-	// cerr << "after applying:" << endl << B << endl;
-	// cerr << subsequents()[27] << endl;
-	// cerr << (B == subsequents()[27]) << endl;
+	bool found = false;
 	for (const Board &B2 : subsequents()) {
-		if (B == B2) return true;
+		if (B == B2) {
+			found = true;
+			break;
+		}
+	}
+	if (!found) return false;
+	switch (B.checkCheck(onTurn)) {
+		case CheckResult::check:
+		case CheckResult::checkmate:
+			return false;
+		default:
+			return true;
+	}
+}
+
+CheckResult Board::checkCheck(int player) const {
+	// First check 'check' (p), then check if we have 'check' after any move (q).
+	// p | q | Result
+	// 0 | 0 | safe
+	// 0 | 1 | stalemate
+	// 1 | 0 | check
+	// 1 | 1 | checkmate
+
+
+	// Check 'check' by switching onTurn and seeing if the king can be taken
+	bool haveP = false;
+	Board B1(*this);
+	B1.onTurn = !player;
+	for (const Board &B2 : B1.subsequents()) {
+		if (B2.pieceScore[B1.onTurn] >= worth(KING)) {
+			haveP = true;
+			break;
+		}
+	}
+
+	// Then check if we have 'check' after any move
+	bool haveQ = true;
+	B1.onTurn = player;
+	for (const Board &B2 : subsequents()) {
+		bool checkNow = false;
+		for (const Board &B3 : B2.subsequents()) {
+			if (B3.pieceScore[B2.onTurn] >= worth(KING)) {
+				checkNow = true;
+				break;
+			}
+		}
+		if (!checkNow) {
+			haveQ = false;
+			break;
+		}
+	}
+
+	if (haveP) {
+		if (haveQ) return CheckResult::checkmate;
+		else return CheckResult::check;
+	} else {
+		if (haveQ) return CheckResult::stalemate;
+		else return CheckResult::safe;
 	}
-	return false;
+}
+
+static const char* playerName(int player) {
+	return &"WHITE\0BLACK"[player * 6];
 }
 
 ostream& operator<<(ostream &os, const Board &B) {
@@ -320,9 +378,16 @@ ostream& operator<<(ostream &os, const Board &B) {
 		}
 		if (y == 0) {
 			os << "  ps[W] = " << B.pieceScore[WHITE];
-			os << "  onTurn = " << (&"WHITE\0BLACK"[B.onTurn * 6]);
+			os << "  onTurn = " << playerName(B.onTurn);
 		} else if (y == 1) {
 			os << "  ps[B] = " << B.pieceScore[BLACK];
+		} else if (y == 2) {
+			switch (B.checkCheck(B.onTurn)) {
+				case CheckResult::safe: break;
+				case CheckResult::check: os << "  " << playerName(B.onTurn) << " IN CHECK"; break;
+				case CheckResult::checkmate: os << "  " << playerName(B.onTurn) << " CHECKMATED."; break;
+				case CheckResult::stalemate: os << "  " << playerName(B.onTurn) << " STALEMATED."; break;
+			}
 		}
 		os << '\n';
 	}