Use mask for ActionPool

Solver/ActionPool.cs

@@ -1,5 +1,6 @@
 using Craftimizer.Simulator.Actions;
 using System.Diagnostics.Contracts;
+using System.Numerics;
 using System.Runtime.CompilerServices;
 using System.Runtime.InteropServices;
 
@@ -39,77 +40,55 @@ public readonly struct ActionPool
     });
 
     public const int MaskSize = 32;
-    public const int EnumSize = 37;
 
-    private unsafe struct EnumBuffer
-    {
-        private fixed byte data[MaskSize];
+    // Bitmask of accepted actions
+    private readonly ulong acceptedActions;
 
-        public ActionType this[int index] => (ActionType)data[index];
-
-        public EnumBuffer(ReadOnlySpan<ActionType> actions)
-        {
-            fixed (byte* dataPtr = data)
-                actions.CopyTo(new Span<ActionType>(dataPtr, MaskSize));
-        }
-
-        public readonly ActionType[] ToArray(int size)
-        {
-            fixed (byte* dataPtr = data)
-                return new Span<ActionType>(dataPtr, size).ToArray();
-        }
-    }
-
-    private unsafe struct LUTBuffer
-    {
-        private fixed byte data[EnumSize];
-
-        public byte this[ActionType index] => data[(byte)index];
-
-        public LUTBuffer(ReadOnlySpan<ActionType> actions)
-        {
-            for (var i = 0; i < EnumSize; i++)
-                data[i] = 0xFF;
-            for (var i = 0; i < actions.Length; i++)
-                data[(byte)actions[i]] = (byte)i;
-        }
-    }
-
-    // List of accepted actions (max 32)
-    private readonly EnumBuffer acceptedActions;
-    // Lookup table for accepted actions (ActionType as idx -> idx in acceptedActions)
-    private readonly LUTBuffer acceptedActionsLUT;
-    private readonly byte size;
-
-    internal ActionType[] AcceptedActions => acceptedActions.ToArray(size);
+    internal ActionType[] AcceptedActions => GetActions();
+    internal int Count => BitOperations.PopCount(acceptedActions);
 
     public ActionPool(ReadOnlySpan<ActionType> actions)
     {
-        if (actions.Length > MaskSize)
-            throw new ArgumentOutOfRangeException(nameof(actions), actions.Length, $"ActionPool only supports up to {MaskSize} actions");
+        acceptedActions = 0;
+        foreach (var action in actions)
+            acceptedActions |= 1ul << ((byte)action);
 
-        acceptedActions = new(actions);
-        acceptedActionsLUT = new(actions);
-        size = (byte)actions.Length;
+        if (Count > MaskSize)
+            throw new ArgumentOutOfRangeException(nameof(actions), actions.Length, $"ActionPool only supports up to {MaskSize} actions");
     }
 
-    [Pure]
-    [MethodImpl(MethodImplOptions.AggressiveInlining)]
-    internal byte FromAction(ActionType action)
+    private ActionType[] GetActions()
     {
-        var ret = acceptedActionsLUT[action];
-        if (ret == 0xFF)
-            throw new ArgumentOutOfRangeException(nameof(action), action, $"Action {action} is unsupported in this pool.");
+        var ret = new ActionType[Count];
+        var i = 0;
+        foreach (var v in (byte[])Enum.GetValuesAsUnderlyingType<ActionType>())
+        {
+            if ((acceptedActions & (1ul << v)) != 0)
+                ret[i++] = (ActionType)v;
+        }
         return ret;
     }
 
     [Pure]
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
-    internal ActionType ToAction(byte index)
+    internal int FromAction(ActionType action)
     {
-        if (index < 0 || index >= size)
+        if ((acceptedActions & (1ul << (byte)action)) == 0)
+            throw new ArgumentOutOfRangeException(nameof(action), action, "Action is not accepted by this pool.");
+
+        // Get number of 1s before action
+        return BitOperations.PopCount(acceptedActions & ((1ul << (byte)action) - 1));
+    }
+
+    [Pure]
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    internal ActionType ToAction(int index)
+    {
+        if (index < 0 || index >= Count)
             throw new ArgumentOutOfRangeException(nameof(index), index, $"Index {index} is out of range for this pool.");
-        return acceptedActions[index];
+
+        // Get index of (index+1)th 1 in set
+        return (ActionType)Intrinsics.NthBitSet(acceptedActions, index);
     }
 
     [Pure]

Solver/ActionSet.cs

@@ -34,7 +34,7 @@ public struct ActionSet
     public readonly bool HasAction(in ActionPool pool, ActionType action) => (bits & pool.ToMask(action)) != 0;
     [Pure]
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
-    public readonly ActionType ElementAt(in ActionPool pool, int index) => pool.ToAction((byte)(Intrinsics.NthBitSet(bits, index) - 1));
+    public readonly ActionType ElementAt(in ActionPool pool, int index) => pool.ToAction(Intrinsics.NthBitSet(bits, index) - 1);
 
     [Pure]
     public readonly int Count => BitOperations.PopCount(bits);

Solver/Intrinsics.cs

@@ -92,11 +92,46 @@ internal static class Intrinsics
         return _base;
     }
 
+    [Pure]
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    private static int NthBitSetScalar(ulong value, int n)
+    {
+        var mask = 0x00000000FFFFFFFFul;
+        var size = 32;
+        var _base = 0;
+
+        if (n++ >= BitOperations.PopCount(value))
+            return 64;
+
+        while (size > 0)
+        {
+            var count = BitOperations.PopCount(value & mask);
+            if (n > count)
+            {
+                _base += size;
+                size >>= 1;
+                mask |= mask << size;
+            }
+            else
+            {
+                size >>= 1;
+                mask >>= size;
+            }
+        }
+
+        return _base;
+    }
+
     [Pure]
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
     private static int NthBitSetBMI2(uint value, int n) =>
         BitOperations.TrailingZeroCount(Bmi2.ParallelBitDeposit(1u << n, value));
 
+    [Pure]
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    private static int NthBitSetBMI2(ulong value, int n) =>
+        BitOperations.TrailingZeroCount(Bmi2.X64.ParallelBitDeposit(1ul << n, value));
+
     [Pure]
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
     public static int NthBitSet(uint value, int n)
@@ -109,6 +144,18 @@ internal static class Intrinsics
             NthBitSetScalar(value, n);
     }
 
+    [Pure]
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    public static int NthBitSet(ulong value, int n)
+    {
+        if (n >= BitOperations.PopCount(value))
+            return 64;
+
+        return Bmi2.X64.IsSupported ?
+            NthBitSetBMI2(value, n) :
+            NthBitSetScalar(value, n);
+    }
+
     [Pure]
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
     public static Vector<float> ReciprocalSqrt(Vector<float> data)