From 2c94afde6aac9b390f0264b5a85d2fedf314d387 Mon Sep 17 00:00:00 2001
From: Ada Joule <ada.fulmina@gmail.com>
Date: Thu, 23 Dec 2021 06:31:55 +0700
Subject: [PATCH] Support checking for checks on pseudo-royal pieces

---
 src/apiutil.h    |    5 +++++
 src/ffishjs.cpp  |    2 +-
 src/position.cpp |   13 ++++++++++++-
 src/position.h   |    1 +
 src/pyffish.cpp  |    2 +-
 test.py          |   13 +++++++++++++
 tests/js/test.js |    9 +++++++++
 7 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/src/apiutil.h b/src/apiutil.h
index f316cca..11a599d 100644
--- a/src/apiutil.h
+++ b/src/apiutil.h
@@ -358,6 +358,11 @@ inline bool has_insufficient_material(Color c, const Position& pos) {
     return true;
 }
 
+inline bool is_check(const Position& pos) {
+    return pos.checkers()
+        || (pos.extinction_pseudo_royal() && pos.attackers_to_pseudo_royals(~pos.side_to_move()));
+}
+
 namespace FEN {
 
 enum FenValidation : int {
diff --git a/src/ffishjs.cpp b/src/ffishjs.cpp
index 54d8761..bed7cd1 100644
--- a/src/ffishjs.cpp
+++ b/src/ffishjs.cpp
@@ -325,7 +325,7 @@ public:
   }
 
   bool is_check() const {
-    return pos.checkers();
+    return Stockfish::is_check(pos);
   }
 
   bool is_bikjang() const {
diff --git a/src/position.cpp b/src/position.cpp
index d5dbd67..2e79820 100644
--- a/src/position.cpp
+++ b/src/position.cpp
@@ -918,6 +918,18 @@ Bitboard Position::attackers_to(Square s, Bitboard occupied) const {
   return attackers_to(s, occupied, WHITE) | attackers_to(s, occupied, BLACK);
 }
 
+/// Position::attackers_to_pseudo_royals computes a bitboard of all pieces
+/// of a particular color attacking at least one opposing pseudo-royal piece
+Bitboard Position::attackers_to_pseudo_royals(Color c) const {
+  Bitboard attackers = 0;
+  Bitboard pseudoRoyals = st->pseudoRoyals & pieces(~c);
+  while (pseudoRoyals) {
+      Square sr = pop_lsb(pseudoRoyals);
+      attackers |= attackers_to(sr, c);
+  }
+  return attackers;
+}
+
 
 /// Position::legal() tests whether a pseudo-legal move is legal
 
@@ -1316,7 +1328,6 @@ bool Position::gives_check(Move m) const {
   }
 }
 
-
 /// Position::do_move() makes a move, and saves all information necessary
 /// to a StateInfo object. The move is assumed to be legal. Pseudo-legal
 /// moves should be filtered out before this function is called.
diff --git a/src/position.h b/src/position.h
index 699d06c..0bb0d10 100644
--- a/src/position.h
+++ b/src/position.h
@@ -239,6 +239,7 @@ public:
   Bitboard blockers_for_king(Color c) const;
   Bitboard check_squares(PieceType pt) const;
   Bitboard pinners(Color c) const;
+  Bitboard attackers_to_pseudo_royals(Color c) const;
 
   // Attacks to/from a given square
   Bitboard attackers_to(Square s) const;
diff --git a/src/pyffish.cpp b/src/pyffish.cpp
index 71d6ec1..734922b 100644
--- a/src/pyffish.cpp
+++ b/src/pyffish.cpp
@@ -248,7 +248,7 @@ extern "C" PyObject* pyffish_givesCheck(PyObject* self, PyObject *args) {
 
     StateListPtr states(new std::deque<StateInfo>(1));
     buildPosition(pos, states, variant, fen, moveList, chess960);
-    return Py_BuildValue("O", pos.checkers() ? Py_True : Py_False);
+    return Py_BuildValue("O", Stockfish::is_check(pos) ? Py_True : Py_False);
 }
 
 // INPUT variant, fen, move list
diff --git a/test.py b/test.py
index 8ba1b72..fc5e8dd 100644
--- a/test.py
+++ b/test.py
@@ -707,6 +707,19 @@ class TestPyffish(unittest.TestCase):
         result = sf.gives_check("capablanca", CAPA, moves)
         self.assertTrue(result)
 
+        # Test giving check to pseudo royal piece
+        result = sf.gives_check("atomic", CHESS, [])
+        self.assertFalse(result)
+
+        result = sf.gives_check("atomic", CHESS, ["e2e4"])
+        self.assertFalse(result)
+
+        result = sf.gives_check("atomic", CHESS, ["e2e4", "d7d5", "f1b5"])
+        self.assertTrue(result)
+
+        result = sf.gives_check("atomic", "rnbqkbnr/ppp2ppp/8/8/8/8/PPP2PPP/RNBQKBNR w KQkq - 0 4", ["d1d7"])
+        self.assertTrue(result)
+
     def test_game_result(self):
         result = sf.game_result("chess", CHESS, ["f2f3", "e7e5", "g2g4", "d8h4"])
         self.assertEqual(result, -sf.VALUE_MATE)
diff --git a/tests/js/test.js b/tests/js/test.js
index f53052a..6b26297 100644
--- a/tests/js/test.js
+++ b/tests/js/test.js
@@ -476,6 +476,15 @@ describe('board.isCheck()', function () {
     board.pushSan("Qxf7#");
     chai.expect(board.isCheck()).to.equal(true);
     board.delete();
+
+    board = new ffish.Board("atomic");
+    chai.expect(board.isCheck()).to.equal(false);
+    board.setFen("rnbqkbnr/ppp1pppp/8/1B1p4/4P3/8/PPPP1PPP/RNBQK1NR b KQkq - 3 2");
+    chai.expect(board.isCheck()).to.equal(true);
+    board.setFen("rnbqkbnr/ppp2ppp/8/8/8/8/PPP2PPP/RNBQKBNR w KQkq - 0 4");
+    board.pushSan("Qd7");
+    chai.expect(board.isCheck()).to.equal(true);
+    board.delete();
   });
 });
 
-- 
1.7.0.4