From 75a468ec85fe4698ce515b27e80eb93e2a8f69e5 Mon Sep 17 00:00:00 2001
From: tqcq <unknown@example.com>
Date: Thu, 24 Oct 2024 04:36:04 +0000
Subject: [PATCH] fix: crash

---
 drama/src/rev-mc.c | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

diff --git a/drama/src/rev-mc.c b/drama/src/rev-mc.c
index 1feda2b..95adb6c 100644
--- a/drama/src/rev-mc.c
+++ b/drama/src/rev-mc.c
@@ -304,9 +304,10 @@ find_row_mask(std::vector<set_t> &sets,
     const uint64_t last_mask = (row_mask << (40 - 16));
     row_mask <<= CL_SHIFT;// skip the lowest 6 bits since they're used for CL addressing
 
-    auto resolve = [=](uint64_t row_mask) -> bool {
+    auto resolve = [=](uint64_t row_mask, std::atomic<bool> &stopped) -> bool {
         if (row_mask & LS_BITMASK(CL_SHIFT)) { return false; }
         for (auto addr_pool : same_row_sets) {
+            if (stopped.load(std::memory_order_relaxed)) { return false; }
             addr_tuple base_addr = addr_pool[0];
             for (int i = 1; i < addr_pool.size(); i++) {
                 addr_tuple tmp = addr_pool[i];
@@ -325,9 +326,9 @@ find_row_mask(std::vector<set_t> &sets,
         row_mask = next_bit_permutation(row_mask);
     }
 
-    std::vector<std::thread> worker;
+    std::vector<std::thread> workers;
     for (int i = 0; i < 8; ++i) {
-        worker.emplace_back([&] {
+        workers.emplace_back([&] {
             while (!found) {
                 uint64_t cur_mask;
                 {
@@ -338,16 +339,19 @@ find_row_mask(std::vector<set_t> &sets,
                     verbose_printerr("[LOG] - Row mask: 0x%0lx \t\t bits: %s\n", cur_mask, bit_string(cur_mask));
                 }
 
-                if (resolve(cur_mask)) {
+                if (resolve(cur_mask, found)) {
                     std::lock_guard<std::mutex> _(lock);
-                    found.store(true);
-                    row_mask = cur_mask;
-                    while (!mask_task.empty()) mask_task.pop();
+                    if (!found.exchange(true)) {
+                        row_mask = cur_mask;
+                        while (!mask_task.empty()) mask_task.pop();
+                    }
                 }
             }
         });
     }
 
+    for (auto &worker : workers) { worker.join(); }
+
     // super hackish way to recover the real row mask
     for (auto m : fn_masks) {
         uint64_t lsb = (1 << (__builtin_ctzl(m) + 1));