From 92f2e21f9dafe3d497a3c4848d16936fd0eab818 Mon Sep 17 00:00:00 2001
From: Andrew Kane <andrew@ankane.org>
Date: Sat, 5 Oct 2024 23:01:52 -0700
Subject: [PATCH] Added stop_search method

---
 CHANGELOG.md                                | 4 ++++
 ext/or-tools/constraint.cpp                 | 8 +++++++-
 lib/or_tools/cp_solver_solution_callback.rb | 4 ++++
 3 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index fd894ee..754ade0 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,3 +1,7 @@
+## 0.13.2 (unreleased)
+
+- Added `stop_search` method
+
 ## 0.13.1 (2024-10-05)
 
 - Added binary installation for Debian 12
diff --git a/ext/or-tools/constraint.cpp b/ext/or-tools/constraint.cpp
index d015efc..7248c7d 100644
--- a/ext/or-tools/constraint.cpp
+++ b/ext/or-tools/constraint.cpp
@@ -1,9 +1,11 @@
 #include <google/protobuf/text_format.h>
 #include <ortools/sat/cp_model.h>
+#include <ortools/util/time_limit.h>
 
 #include "ext.h"
 
 using operations_research::Domain;
+using operations_research::TimeLimit;
 using operations_research::sat::BoolVar;
 using operations_research::sat::Constraint;
 using operations_research::sat::CpModelBuilder;
@@ -424,17 +426,21 @@ void init_constraint(Rice::Module& m) {
       [](Object self, CpModelBuilder& model, SatParameters& parameters, Object callback) {
         Model m;
 
+        std::atomic<bool> stopped(false);
+        m.GetOrCreate<TimeLimit>()->RegisterExternalBooleanAsLimit(&stopped);
+
         if (!callback.is_nil()) {
           // TODO figure out how to use callback with multiple cores
           parameters.set_num_search_workers(1);
 
           m.Add(NewFeasibleSolutionObserver(
-            [&callback](const CpSolverResponse& r) {
+            [&](const CpSolverResponse& r) {
               // ensure Ruby thread
               if (ruby_native_thread_p()) {
                 // TODO find a better way to do this
                 callback.call("response=", r);
                 callback.call("on_solution_callback");
+                stopped = callback.attr_get("@stopped");
               }
             })
           );
diff --git a/lib/or_tools/cp_solver_solution_callback.rb b/lib/or_tools/cp_solver_solution_callback.rb
index c71509a..ec2166c 100644
--- a/lib/or_tools/cp_solver_solution_callback.rb
+++ b/lib/or_tools/cp_solver_solution_callback.rb
@@ -16,5 +16,9 @@ def value(expr)
     def objective_value
       @response.objective_value
     end
+
+    def stop_search
+      @stopped = true
+    end
   end
 end