From 7a41d7ee6a493077a2800a4e3bf9ebfae1df9301 Mon Sep 17 00:00:00 2001
From: Saurabh Chalke <saurabhchalke@gmail.com>
Date: Sun, 5 May 2024 02:47:29 +0000
Subject: [PATCH] fix: terminate early on first counterexample with
 --early-exit (#243)

The --early-exit flag is now properly handled to ensure that Halmos terminates
immediately after finding the first counterexample, avoiding the reporting of
multiple counterexamples. The previous implementation used a busy-waiting loop
with time.sleep(1) to wait for the first counterexample or all futures to
complete, which was inefficient and allowed multiple counterexamples to accumulate.

The updated code introduces an early_exit flag variable to keep track of whether
an early exit is requested. When a counterexample is found and --early-exit is set,
the early_exit flag is set to True. The future_callback function checks this flag
and immediately returns if it is True, preventing further processing of results.

The thread pool is shut down after all the submitted futures have completed or
been canceled, ensuring a clean termination.
---
 src/halmos/__main__.py | 30 +++++++++++++++++++++---------
 1 file changed, 21 insertions(+), 9 deletions(-)

diff --git a/src/halmos/__main__.py b/src/halmos/__main__.py
index 45c8331f6..b8b05821d 100644
--- a/src/halmos/__main__.py
+++ b/src/halmos/__main__.py
@@ -657,7 +657,14 @@ def run(
     counterexamples = []
     traces = {}
 
+    early_exit = False
+
     def future_callback(future_model):
+        nonlocal early_exit
+
+        if early_exit:
+            return
+
         m = future_model.result()
         models.append(m)
 
@@ -665,17 +672,21 @@ def future_callback(future_model):
         if result == unsat:
             return
 
+        counterexample_found = False
+
         # model could be an empty dict here
         if model is not None:
             if is_valid:
                 print(red(f"Counterexample: {render_model(model)}"))
                 counterexamples.append(model)
+                counterexample_found = True
             else:
                 warn(
                     COUNTEREXAMPLE_INVALID,
                     f"Counterexample (potentially invalid): {render_model(model)}",
                 )
                 counterexamples.append(model)
+                counterexample_found = True
         else:
             warn(COUNTEREXAMPLE_UNKNOWN, f"Counterexample: {result}")
 
@@ -691,6 +702,9 @@ def future_callback(future_model):
             )
             print(traces[index], end="")
 
+        if args.early_exit and counterexample_found:
+            early_exit = True
+
     for idx, ex in enumerate(exs):
         result_exs.append(ex)
 
@@ -739,16 +753,14 @@ def future_callback(future_model):
             f"# of potential paths involving assertion violations: {len(future_models)} / {len(result_exs)}  (--solver-threads {args.solver_threads})"
         )
 
-    if args.early_exit:
-        while not (
-            len(counterexamples) > 0 or all([fm.done() for fm in future_models])
-        ):
-            time.sleep(1)
-
-        thread_pool.shutdown(wait=False, cancel_futures=True)
+    for future in future_models:
+        try:
+            future.result()
+        except Exception as e:
+            if not isinstance(e, concurrent.futures.CancelledError):
+                raise
 
-    else:
-        thread_pool.shutdown(wait=True)
+    thread_pool.shutdown(wait=True)
 
     counter = Counter(str(m.result) for m in models)
     if counter["sat"] > 0: