Add searchmoves functionality

lc0/src/chess/uciloop.cc

@@ -41,7 +41,7 @@ const std::unordered_map<std::string, std::unordered_set<std::string>>
         {{"position"}, {"fen", "startpos", "moves"}},
         {{"go"},
          {"infinite", "wtime", "btime", "winc", "binc", "movestogo", "depth",
-          "nodes", "movetime"}},
+          "nodes", "movetime", "searchmoves"}},
         {{"start"}, {}},
         {{"stop"}, {}},
         {{"quit"}, {}},
@@ -133,11 +133,12 @@ bool UciLoop::DispatchCommand(
   } else if (command == "go") {
     GoParams go_params;
     if (ContainsKey(params, "infinite")) {
-      if (!GetOrEmpty(params, "infinite").empty()) {
-        throw Exception("Unexpected token " + GetOrEmpty(params, "infinite"));
-      }
       go_params.infinite = true;
     }
+    if (ContainsKey(params, "searchmoves")) {
+      go_params.searchmoves =
+          StrSplitAtWhitespace(GetOrEmpty(params, "searchmoves"));
+    }
 #define OPTION(x)                                    \
   if (ContainsKey(params, #x)) {                     \
     go_params.x = std::stoi(GetOrEmpty(params, #x)); \
@@ -211,4 +212,4 @@ void UciLoop::SendInfo(const ThinkingInfo& info) {
   SendResponse(res);
 }
 
-}  // namespace lczero
+}  // namespace lczero

lc0/src/chess/uciloop.h

@@ -37,6 +37,7 @@ struct GoParams {
   int nodes = -1;
   std::int64_t movetime = -1;
   bool infinite = false;
+  std::vector<std::string> searchmoves;
 };
 
 class UciLoop {
@@ -78,4 +79,4 @@ class UciLoop {
   std::ofstream debug_log_;
 };
 
-}  // namespace lczero
+}  // namespace lczero

lc0/src/engine.cc

@@ -75,6 +75,7 @@ SearchLimits EngineController::PopulateSearchLimits(int /*ply*/, bool is_black,
   limits.visits = params.nodes;
   limits.time_ms = params.movetime;
   int64_t time = (is_black ? params.btime : params.wtime);
+  limits.searchmoves = params.searchmoves;
   limits.infinite = params.infinite;
   if (params.infinite || time < 0) return limits;
   int increment = std::max(int64_t(0), is_black ? params.binc : params.winc);

lc0/src/mcts/node.cc

@@ -382,4 +382,4 @@ void NodeTree::DeallocateTree() {
   current_head_ = nullptr;
 }
 
-}  // namespace lczero
+}  // namespace lczero

lc0/src/mcts/node.h

@@ -57,7 +57,7 @@ class Node {
   // Returns whether a node has children.
   bool HasChildren() const { return child_ != nullptr; }
 
-  // Returns move from the point of new of player BEFORE the position.
+  // Returns move from the point of view of player BEFORE the position.
   Move GetMove() const { return move_; }
 
   // Returns move, with optional flip (false == player BEFORE the position).
@@ -210,4 +210,4 @@ class NodeTree {
   Node* gamebegin_node_ = nullptr;
   PositionHistory history_;
 };
-}  // namespace lczero
+}  // namespace lczero

lc0/src/mcts/search.cc

@@ -641,6 +641,8 @@ Node* Search::PickNodeToExtend(Node* node, PositionHistory* history) {
             ? -node->GetQ(0, kExtraVirtualLoss)
             : -node->GetQ(0, kExtraVirtualLoss) -
                   kFpuReduction * std::sqrt(node->GetVisitedPolicy());
+
+    bool root_color = played_history_.IsBlackToMove();
     for (Node* iter : node->Children()) {
       if (is_root_node) {
         // If there's no chance to catch up the currently best node with
@@ -652,6 +654,14 @@ Node* Search::PickNodeToExtend(Node* node, PositionHistory* history) {
             remaining_playouts_ < best_node_n - iter->GetNStarted()) {
           continue;
         }
+        // If searchmoves was sent, restrict the search only in that moves
+        if (!limits_.searchmoves.empty() &&
+            std::find(limits_.searchmoves.begin(),
+                      limits_.searchmoves.end(),
+                      iter->GetMove(root_color).as_string()) ==
+            limits_.searchmoves.end()) {
+          continue;
+        }
         ++possible_moves;
       }
       float Q = iter->GetQ(parent_q, kExtraVirtualLoss);

lc0/src/mcts/search.h

@@ -37,6 +37,7 @@ struct SearchLimits {
   std::int64_t playouts = -1;
   std::int64_t time_ms = -1;
   bool infinite = false;
+  std::vector<std::string> searchmoves;
 };
 
 class Search {