From 6e9c1600bb40d0d1103e2ef60aa57534bd409c7c Mon Sep 17 00:00:00 2001
From: Dima <dima@current.ai>
Date: Tue, 30 Apr 2024 18:43:50 +0200
Subject: [PATCH] Added the fix to `WaitableAtomic` that fixes the test two
 commits above. [ FIXED NOW ]

---
 bricks/sync/waitable_atomic.h | 17 +++++++++++++++--
 1 file changed, 15 insertions(+), 2 deletions(-)

diff --git a/bricks/sync/waitable_atomic.h b/bricks/sync/waitable_atomic.h
index ad6efa4e..3ec060bc 100644
--- a/bricks/sync/waitable_atomic.h
+++ b/bricks/sync/waitable_atomic.h
@@ -214,17 +214,30 @@ class WaitableAtomic {
   // TODO(dkorolev): The `.Wait()` above always returning `true` could use some TLC.
 
   template <typename F>
-  std::invoke_result_t<F, data_t&> Wait(std::function<bool(const data_t&)> wait_predicate, F&& retval_predicate) {
+  std::invoke_result_t<F, data_t&> DoWait(std::function<bool(const data_t&)> wait_predicate, F&& retval_predicate) {
     std::unique_lock<std::mutex> lock(data_mutex_);
     if (!wait_predicate(data_)) {
       const data_t& data = data_;
-      data_condition_variable_.wait(lock, [&wait_predicate, &data] { return wait_predicate(data); });
+      data_condition_variable_.wait(lock, [&wait_predicate, &data]() { return wait_predicate(data); });
       return retval_predicate(data_);
     } else {
       return retval_predicate(data_);
     }
   }
 
+  template <typename F, class = std::enable_if_t<std::is_same_v<std::invoke_result_t<F, data_t&>, void>>>
+  void Wait(std::function<bool(const data_t&)> wait_predicate, F&& retval_predicate) {
+    DoWait(wait_predicate, std::forward<F>(retval_predicate));
+    Notify();
+  }
+
+  template <typename F, class = std::enable_if_t<!std::is_same_v<std::invoke_result_t<F, data_t&>, void>>>
+  std::invoke_result_t<F, data_t&> Wait(std::function<bool(const data_t&)> wait_predicate, F&& retval_predicate) {
+    std::invoke_result_t<F, data_t&> retval = DoWait(wait_predicate, std::forward<F>(retval_predicate));
+    Notify();
+    return retval;
+  }
+
 #endif  // CURRENT_FOR_CPP14
 
   template <typename T>