From 2538401fa2d654535b6825676d87f91c3c933646 Mon Sep 17 00:00:00 2001
From: Hao Wei <yuanhao1@cs.cmu.edu>
Date: Sun, 11 Feb 2024 16:13:04 -0500
Subject: [PATCH] minor optimization for hand-over-hand locking

---
 benchmarks/concurrentKNN/octTree/flock/spin_lock.h | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/benchmarks/concurrentKNN/octTree/flock/spin_lock.h b/benchmarks/concurrentKNN/octTree/flock/spin_lock.h
index 608d3a9b..80228727 100644
--- a/benchmarks/concurrentKNN/octTree/flock/spin_lock.h
+++ b/benchmarks/concurrentKNN/octTree/flock/spin_lock.h
@@ -7,7 +7,7 @@
 //   is_locked() -> bool
 
 #include <parlay/parallel.h> // needed for worker_id
-#include "epoch.h"  // for worker_id()
+#include <flock/epoch.h>  // for worker_id()
 
 #include <atomic>
 #include<chrono>
@@ -82,9 +82,10 @@ struct lock {
       lock_entry newl = current.take_lock();
       if (lck.compare_exchange_strong(current, newl)) {
         RT result = f();
-        if (no_release == nullptr)
-          lck.compare_exchange_strong(newl, newl.release_lock());
-          // lck = newl.release_lock();  // release lock
+        if (no_release == nullptr) {
+          if(lck.load().is_self_locked()) // lock might have been released early during f()
+            lck = newl.release_lock();  // release lock
+        }
         else *no_release = true;
         return std::optional<RT>(result); 
       } else {
@@ -124,7 +125,7 @@ struct lock {
 
   void unlock() {
     //std::cout << "released lock:" << this << ", " << is_locked() << std::endl;
-    if (!is_locked()) abort();
+    assert(lck.load().is_self_locked());
     lck = lck.load().release_lock();
   }