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

@@ -657,25 +657,36 @@ 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)
 
         model, is_valid, index, result = m.model, m.is_valid, m.index, m.result
         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: