Spell-chess: persist potion state in FEN

Author: Belzedar94

src/position.cpp

@@ -40,14 +40,61 @@ namespace {
 
   inline Variant::PotionType potion_type_from_piece(const Variant* var, PieceType pt) {
     if (!var || !var->potions)
-        return static_cast<Variant::PotionType>(Variant::POTION_TYPE_NB);
+        return static_cast<Variant::PotionType>(Variant::POTION_TYPE_NB);       
     if (pt == var->potionPiece[Variant::POTION_FREEZE])
         return Variant::POTION_FREEZE;
     if (pt == var->potionPiece[Variant::POTION_JUMP])
         return Variant::POTION_JUMP;
     return static_cast<Variant::PotionType>(Variant::POTION_TYPE_NB);
   }
 
+  inline Square potion_zone_center(const Position& pos, Variant::PotionType potion,
+                                   Bitboard zone) {
+    if (zone == Bitboard(0))
+        return SQ_NONE;
+    if (potion == Variant::POTION_JUMP)
+        return lsb(zone);
+    Bitboard candidates = zone;
+    while (candidates)
+    {
+        Square s = pop_lsb(candidates);
+        if (pos.freeze_zone_from_square(s) == zone)
+            return s;
+    }
+    return lsb(zone);
+  }
+
+  inline Bitboard potion_zone_from_center(const Position& pos, Variant::PotionType potion,
+                                          Square s) {
+    if (s == SQ_NONE)
+        return Bitboard(0);
+    if (potion == Variant::POTION_FREEZE)
+        return pos.freeze_zone_from_square(s);
+    return square_bb(s);
+  }
+
+  inline Square parse_fen_square(const Position& pos, const std::string& token) {
+    if (token.size() < 2)
+        return SQ_NONE;
+    char fileChar = char(tolower(token[0]));
+    if (fileChar < 'a' || fileChar > char('a' + pos.max_file()))
+        return SQ_NONE;
+    int file = fileChar - 'a';
+    int rank = 0;
+    for (size_t i = 1; i < token.size(); ++i)
+    {
+        if (!isdigit(token[i]))
+            return SQ_NONE;
+        rank = rank * 10 + (token[i] - '0');
+    }
+    if (rank <= 0)
+        return SQ_NONE;
+    int rankIndex = rank - 1;
+    if (rankIndex > pos.max_rank())
+        return SQ_NONE;
+    return make_square(File(file), Rank(rankIndex));
+  }
+
   struct SpellContextScope {
     Position& pos;
     bool active;
@@ -402,7 +449,87 @@ Position& Position::set(const Variant* v, const string& fenStr, bool isChess960,
               add_to_hand(Piece(idx));
       }
 
+  if (!sfen && potions_enabled())
+  {
+      std::string potionState;
+      ss >> std::ws;
+      if (ss.peek() == '{')
+      {
+          ss >> token;
+          while (ss >> token)
+          {
+              if (token == '}')
+                  break;
+              if (!isspace(token))
+                  potionState.push_back(char(token));
+          }
+      }
+
+      auto parse_potion_state = [&](const std::string& state) {
+          size_t start = 0;
+          const int maxCooldown = (1u << POTION_COOLDOWN_BITS) - 1;
+          while (start < state.size())
+          {
+              size_t end = state.find(',', start);
+              std::string entry = state.substr(start, end == std::string::npos ? end : end - start);
+              if (entry.size() >= 4)
+              {
+                  char pieceChar = entry[0];
+                  Color c = islower(pieceChar) ? BLACK : WHITE;
+                  Variant::PotionType potion = static_cast<Variant::PotionType>(Variant::POTION_TYPE_NB);
+                  for (int pt = 0; pt < Variant::POTION_TYPE_NB; ++pt)
+                  {
+                      Variant::PotionType ptEnum = static_cast<Variant::PotionType>(pt);
+                      PieceType potionPiece = potion_piece(ptEnum);
+                      if (potionPiece == NO_PIECE_TYPE)
+                          continue;
+                      char expected = piece_to_char()[make_piece(WHITE, potionPiece)];
+                      if (tolower(expected) == tolower(pieceChar))
+                      {
+                          potion = ptEnum;
+                          break;
+                      }
+                  }
+
+                  size_t at = entry.find('@', 1);
+                  size_t colon = entry.find(':', at == std::string::npos ? 0 : at + 1);
+                  if (potion != Variant::POTION_TYPE_NB && at != std::string::npos && colon != std::string::npos)
+                  {
+                      std::string squareToken = entry.substr(at + 1, colon - at - 1);
+                      std::string cooldownToken = entry.substr(colon + 1);
+                      Square center = SQ_NONE;
+                      if (!squareToken.empty() && squareToken != "-")
+                          center = parse_fen_square(*this, squareToken);
+
+                      int cooldown = 0;
+                      for (char ch : cooldownToken)
+                      {
+                          if (!isdigit(ch))
+                          {
+                              cooldown = 0;
+                              break;
+                          }
+                          cooldown = cooldown * 10 + (ch - '0');
+                      }
+                      if (cooldown > maxCooldown)
+                          cooldown = maxCooldown;
+
+                      st->potionCooldown[c][potion] = cooldown;
+                      st->potionZones[c][potion] = potion_zone_from_center(*this, potion, center);
+                  }
+              }
+              if (end == std::string::npos)
+                  break;
+              start = end + 1;
+          }
+      };
+
+      if (!potionState.empty())
+          parse_potion_state(potionState);
+  }
+
   // 2. Active color
+  ss >> std::ws;
   ss >> token;
   sideToMove = (token != (sfen ? 'w' : 'b') ? WHITE : BLACK);  // Invert colors for SFEN
   ss >> token;
@@ -836,6 +963,38 @@ string Position::fen(bool sfen, bool showPromoted, int countStarted, std::string
       ss << ']';
   }
 
+  if (potions_enabled())
+  {
+      std::string potionState;
+      const int maxCooldown = (1u << POTION_COOLDOWN_BITS) - 1;
+      for (Color c : {WHITE, BLACK})
+          for (int pt = 0; pt < Variant::POTION_TYPE_NB; ++pt)
+          {
+              Variant::PotionType potion = static_cast<Variant::PotionType>(pt);
+              PieceType potionPiece = potion_piece(potion);
+              if (potionPiece == NO_PIECE_TYPE)
+                  continue;
+              if (!potionState.empty())
+                  potionState += ",";
+              potionState += piece_to_char()[make_piece(c, potionPiece)];
+              potionState += "@";
+              Bitboard zone = potion_zone(c, potion);
+              if (zone)
+                  potionState += UCI::square(*this, potion_zone_center(*this, potion, zone));
+              else
+                  potionState += "-";
+              potionState += ":";
+              int cooldown = potion_cooldown(c, potion);
+              if (cooldown < 0)
+                  cooldown = 0;
+              if (cooldown > maxCooldown)
+                  cooldown = maxCooldown;
+              potionState += std::to_string(cooldown);
+          }
+      if (!potionState.empty())
+          ss << " {" << potionState << "}";
+  }
+
   ss << (sideToMove == WHITE ? " w " : " b ");
 
   // Disambiguation for chess960 "king" square
@@ -1352,8 +1511,8 @@ bool Position::legal(Move m) const {
       if (freeze_squares() & to_sq(m))
           return false;
 
-      // Non-royal pieces can not be impeded from castling
-      if (type_of(piece_on(from)) != KING)
+      // Only the castling king piece is subject to attack checks
+      if (type_of(piece_on(from)) != castling_king_piece(us))
           return true;
 
       for (Square s = to; s != from; s += step)

tests/perft.sh

@@ -157,6 +157,15 @@ if [[ $1 == "all" || $1 == "variant" ]]; then
   expect perft.exp extinction "fen rnbqb1kr/pppppppp/8/8/8/8/PPPPPPPP/RNBQB1KR w AHah - 0 1" 4 195286 true > /dev/null
   expect perft.exp seirawan "fen qbbrnkrn/pppppppp/8/8/8/8/PPPPPPPP/QBBRNKRN[HEhe] w ABCDEFGHabcdefgh - 0 1" 3 21170 true > /dev/null
   expect perft.exp spell-chess "fen 4k3/p7/8/8/8/8/8/4K2R[f] b K - 0 1 moves f@h1,a7a6" 1 5 > /dev/null
+  castling_output=$(printf "setoption name UCI_Variant value spell-chess\nposition fen rnb1k2r/pp1pnppp/2p1p3/q1b1P3/3P4/5N2/PPP1BPPP/RNBQK2R[JJFFFFjjffff] w KQkq - 1 6\ngo perft 1\nquit\n" | ./stockfish)
+  if echo "$castling_output" | grep -q "e1g1:"; then
+    echo "spell-chess castling-in-check test failed (e1g1)"
+    exit 1
+  fi
+  if echo "$castling_output" | grep -q "e1c1:"; then
+    echo "spell-chess castling-in-check test failed (e1c1)"
+    exit 1
+  fi
 fi
 
 # large-board variants