Merge pull request #602 from wowus/master

Cleaned up critical section code.

2 files changed, 34 insertions, 9 deletions

diff --git a/src/test/util_tests.cpp b/src/test/util_tests.cpp
index 8afc85c507..8c8b99e1b2 100644
--- a/src/test/util_tests.cpp
+++ b/src/test/util_tests.cpp
@@ -8,6 +8,25 @@ using namespace std;
 
 BOOST_AUTO_TEST_SUITE(util_tests)
 
+BOOST_AUTO_TEST_CASE(util_criticalsection)
+{
+    CCriticalSection cs;
+
+    do {
+        CRITICAL_BLOCK(cs)
+            break;
+
+        BOOST_ERROR("break was swallowed!");
+    } while(0);
+
+    do {
+        TRY_CRITICAL_BLOCK(cs)
+            break;
+
+        BOOST_ERROR("break was swallowed!");
+    } while(0);
+}
+
 BOOST_AUTO_TEST_CASE(util_MedianFilter)
 {    
     CMedianFilter<int> filter(5, 15);
diff --git a/src/util.h b/src/util.h
index 178923727a..4c966486f7 100644
--- a/src/util.h
+++ b/src/util.h
@@ -243,19 +243,20 @@ public:
         pcs = &csIn;
         pcs->Enter(pszName, pszFile, nLine);
     }
+
+    operator bool() const
+    {
+        return true;
+    }
+
     ~CCriticalBlock()
     {
         pcs->Leave();
     }
 };
 
-// WARNING: This will catch continue and break!
-// break is caught with an assertion, but there's no way to detect continue.
-// I'd rather be careful than suffer the other more error prone syntax.
-// The compiler will optimise away all this loop junk.
 #define CRITICAL_BLOCK(cs)     \
-    for (bool fcriticalblockonce=true; fcriticalblockonce; assert(("break caught by CRITICAL_BLOCK!" && !fcriticalblockonce)), fcriticalblockonce=false) \
-        for (CCriticalBlock criticalblock(cs, #cs, __FILE__, __LINE__); fcriticalblockonce; fcriticalblockonce=false)
+    if (CCriticalBlock criticalblock = CCriticalBlock(cs, #cs, __FILE__, __LINE__))
 
 class CTryCriticalBlock
 {
@@ -267,6 +268,12 @@ public:
     {
         pcs = (csIn.TryEnter(pszName, pszFile, nLine) ? &csIn : NULL);
     }
+
+    operator bool() const
+    {
+        return Entered();
+    }
+
     ~CTryCriticalBlock()
     {
         if (pcs)
@@ -274,12 +281,11 @@ public:
             pcs->Leave();
         }
     }
-    bool Entered() { return pcs != NULL; }
+    bool Entered() const { return pcs != NULL; }
 };
 
 #define TRY_CRITICAL_BLOCK(cs)     \
-    for (bool fcriticalblockonce=true; fcriticalblockonce; assert(("break caught by TRY_CRITICAL_BLOCK!" && !fcriticalblockonce)), fcriticalblockonce=false) \
-        for (CTryCriticalBlock criticalblock(cs, #cs, __FILE__, __LINE__); fcriticalblockonce && (fcriticalblockonce = criticalblock.Entered()); fcriticalblockonce=false)
+    if (CTryCriticalBlock criticalblock = CTryCriticalBlock(cs, #cs, __FILE__, __LINE__))