Optimized index selection

dotnet · Mar 15, 2024 · 3411920 · 3411920
1 parent 1b7bf85
commit 3411920
Showing 1 changed file with 13 additions and 10 deletions.
diff --git a/src/DotNext.Threading/Collections/Concurrent/IndexPool.cs b/src/DotNext.Threading/Collections/Concurrent/IndexPool.cs
@@ -3,6 +3,7 @@
 using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
 using System.Numerics;
+using System.Runtime.CompilerServices;
 using System.Runtime.InteropServices;
 
 namespace DotNext.Collections.Concurrent;
@@ -73,19 +74,21 @@ public readonly bool TryPeek(out int result)
     /// <seealso cref="Return(int)"/>
     public bool TryTake(out int result)
     {
-        ulong current, newValue = Volatile.Read(in bitmask);
-        do
-        {
-            result = BitOperations.TrailingZeroCount(current = newValue);
+        return TryTake(ref bitmask, maxValue, out result);
 
-            if (result > maxValue)
-                return false;
+        static bool TryTake(ref ulong bitmask, int maxValue, out int result)
+        {
+            var current = Volatile.Read(in bitmask);
+            for (ulong newValue; ; current = newValue)
+            {
+                newValue = current & (current - 1UL); // Reset lowest set bit, the same as BLSR instruction
+                newValue = Interlocked.CompareExchange(ref bitmask, newValue, current);
+                if (newValue == current)
+                    break;
+            }
 
-            newValue = current ^ (1UL << result);
+            return (result = BitOperations.TrailingZeroCount(current)) <= maxValue;
         }
-        while ((newValue = Interlocked.CompareExchange(ref bitmask, newValue, current)) != current);
-
-        return true;
     }
 
     /// <summary>