Refactor BITPOS Implementation (#1016)

* wip

* cleanup bitpos tests

* Fix bitpos bit search

* simplify bitpos byte search

* cleanup unit tests for bitpos

* fix masked invalidPayload matching

* add more tests

* fix boundary conditions for bit search and add tests

* change to portable intrinsics

* add more tests for bitpos

libs/server/Resp/Bitmap/BitmapManagerBitPos.cs

@@ -1,185 +1,174 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT license.
 
-using System.Diagnostics;
-using System.Runtime.Intrinsics.X86;
+using System.Buffers.Binary;
+using System.Numerics;
 
 namespace Garnet.server
 {
     public unsafe partial class BitmapManager
     {
         /// <summary>
-        /// Find pos of bit set/clear for given bit offsets within a single byte.
+        /// Main driver for BITPOS command
         /// </summary>
-        /// <param name="value">Byte value to search within.</param>
-        /// <param name="bSetVal">Bit value to search for (0|1).</param>
-        /// <param name="startBitOffset">Start most significant bit offset in byte value.</param>
-        /// <param name="endBitOffset">End most significant bit offset in bitmap.</param>
-        /// <returns></returns>
-        private static long BitPosIndexBitSingleByteSearch(byte value, byte bSetVal, int startBitOffset = 0, int endBitOffset = 8)
-        {
-            Debug.Assert(startBitOffset >= 0 && startBitOffset <= 8);
-            Debug.Assert(endBitOffset >= 0 && endBitOffset <= 8);
-            bool bflag = (bSetVal == 0);
-            long mask = bflag ? -1 : 0;
-
-            int leftBitIndex = 1 << (8 - startBitOffset);
-            int rightBitIndex = 1 << (8 - endBitOffset);
-
-            // Create extraction mask
-            long extract = leftBitIndex - rightBitIndex;
-
-            long payload = (long)(value & extract) << 56;
-            // Trim leading bits
-            payload = payload << startBitOffset;
-
-            // Transform to count leading zeros
-            payload = bflag ? ~payload : payload;
-
-            // Return not found
-            if (payload == mask) return -1;
-
-            return (long)Lzcnt.X64.LeadingZeroCount((ulong)payload);
-        }
-
-        /// <summary>
-        /// Find pos of bit set/clear for given bit offset.
-        /// </summary>
-        /// <param name="value">Pointer to start of bitmap.</param>
-        /// <param name="bSetVal">Bit value to search for (0|1).</param>
-        /// <param name="offset">Bit offset in bitmap.</param>
-        /// <returns></returns>
-        private static long BitPosIndexBitSearch(byte* value, byte bSetVal, long offset = 0)
-        {
-            bool bflag = (bSetVal == 0);
-            long mask = bflag ? -1 : 0;
-            long startByteOffset = (offset / 8);
-            int bitOffset = (int)(offset & 7);
-
-            long payload = (long)value[startByteOffset] << 56;
-            // Trim leading bits
-            payload = payload << bitOffset;
-
-            // Transform to count leading zeros
-            payload = bflag ? ~payload : payload;
-
-            // Return not found
-            if (payload == mask)
-                return -1;
-
-            return (long)Lzcnt.X64.LeadingZeroCount((ulong)payload);
-        }
-
-        /// <summary>
-        /// Main driver for bit position command.
-        /// </summary>
-        /// <param name="setVal"></param>
+        /// <param name="input"></param>
+        /// <param name="inputLen"></param>
         /// <param name="startOffset"></param>
         /// <param name="endOffset"></param>
+        /// <param name="searchFor"></param>
         /// <param name="offsetType"></param>
-        /// <param name="value">Pointer to start of bitmap.</param>
-        /// <param name="valLen">Length of bitmap.</param>
         /// <returns></returns>
-        public static long BitPosDriver(byte setVal, long startOffset, long endOffset, byte offsetType, byte* value, int valLen)
+        public static long BitPosDriver(byte* input, int inputLen, long startOffset, long endOffset, byte searchFor, byte offsetType)
         {
             if (offsetType == 0x0)
             {
-                startOffset = startOffset < 0 ? ProcessNegativeOffset(startOffset, valLen) : startOffset;
-                endOffset = endOffset < 0 ? ProcessNegativeOffset(endOffset, valLen) : endOffset;
+                startOffset = startOffset < 0 ? ProcessNegativeOffset(startOffset, inputLen) : startOffset;
+                endOffset = endOffset < 0 ? ProcessNegativeOffset(endOffset, inputLen) : endOffset;
 
-                if (startOffset >= valLen) // If startOffset greater that valLen always bitpos -1
+                if (startOffset >= inputLen) // If startOffset greater that valLen always bitpos -1
                     return -1;
 
                 if (startOffset > endOffset) // If start offset beyond endOffset return 0
                     return -1;
 
-                endOffset = endOffset >= valLen ? valLen : endOffset;
-                long pos = BitPosByte(value, setVal, startOffset, endOffset);
-                // check if position is exceeding the last byte in acceptable range
-                return pos >= ((endOffset + 1) * 8) ? -1 : pos;
+                endOffset = endOffset >= inputLen ? inputLen : endOffset;
+                // BYTE search
+                return BitPosByteSearch(input, inputLen, startOffset, endOffset, searchFor);
             }
-
-            startOffset = startOffset < 0 ? ProcessNegativeOffset(startOffset, valLen * 8) : startOffset;
-            endOffset = endOffset < 0 ? ProcessNegativeOffset(endOffset, valLen * 8) : endOffset;
-
-            var startByte = (startOffset / 8);
-            var endByte = (endOffset / 8);
-            if (startByte == endByte)
+            else
             {
-                // Search only inside single byte for pos
-                var leftBitIndex = (int)(startOffset & 7);
-                var rightBitIndex = (int)((endOffset + 1) & 7);
-                var _ipos = BitPosIndexBitSingleByteSearch(value[startByte], setVal, leftBitIndex, rightBitIndex);
-                return _ipos == -1 ? _ipos : startOffset + _ipos;
-            }
+                startOffset = startOffset < 0 ? ProcessNegativeOffset(startOffset, inputLen * 8) : startOffset;
+                endOffset = endOffset < 0 ? ProcessNegativeOffset(endOffset, inputLen * 8) : endOffset;
+
+                var startByteIndex = startOffset >> 3;
+                var endByteIndex = endOffset >> 3;
 
-            // Search prefix and terminate if found position of bit
-            var _ppos = BitPosIndexBitSearch(value, setVal, startOffset);
-            if (_ppos != -1) return startOffset + _ppos;
+                if (startByteIndex >= inputLen) // If startOffset greater that valLen always bitpos -1
+                    return -1;
 
-            // Adjust offsets to skip first and last byte
-            var _startOffset = (startOffset / 8) + 1;
-            var _endOffset = (endOffset / 8) - 1;
-            var _bpos = BitPosByte(value, setVal, _startOffset, _endOffset);
+                if (startByteIndex > endByteIndex) // If start offset beyond endOffset return 0
+                    return -1;
 
-            if (_bpos != -1 && _bpos < (_endOffset + 1) * 8) return _bpos;
+                endOffset = endByteIndex >= inputLen ? inputLen << 3 : endOffset;
 
-            // Search suffix
-            var _spos = BitPosIndexBitSearch(value, setVal, endOffset);
-            return _spos;
+                // BIT search
+                return BitPosBitSearch(input, inputLen, startOffset, endOffset, searchFor);
+            }
         }
 
         /// <summary>
-        /// Find pos of set/clear bit in a sequence of bytes.
+        /// Search for position of bit set in byte array using bit offset for start and end range
         /// </summary>
-        /// <param name="value">Pointer to start of bitmap.</param>
-        /// <param name="bSetVal">The bit value to search for (0 for cleared bit or 1 for set bit).</param>
-        /// <param name="startOffset">Starting offset into bitmap.</param>
-        /// <param name="endOffset">End offset into bitmap.</param>
+        /// <param name="input"></param>
+        /// <param name="inputLen"></param>
+        /// <param name="startBitOffset"></param>
+        /// <param name="endBitOffset"></param>
+        /// <param name="searchFor"></param>
         /// <returns></returns>
-        private static long BitPosByte(byte* value, byte bSetVal, long startOffset, long endOffset)
+        private static long BitPosBitSearch(byte* input, long inputLen, long startBitOffset, long endBitOffset, byte searchFor)
         {
-            // Mask set to look for 0 or 1 depending on clear/set flag
-            bool bflag = (bSetVal == 0);
-            long mask = bflag ? -1 : 0;
-            long len = (endOffset - startOffset) + 1;
-            long remainder = len & 7;
-            byte* curr = value + startOffset;
-            byte* end = curr + (len - remainder);
-
-            // Search for first word not matching mask.
-            while (curr < end)
+            var searchBit = searchFor == 1;
+            var invalidPayload = (byte)(searchBit ? 0x00 : 0xff);
+            var currentBitOffset = (int)startBitOffset;
+            while (currentBitOffset <= endBitOffset)
             {
-                long v = *(long*)(curr);
-                if (v != mask) break;
-                curr += 8;
-            }
+                var byteIndex = currentBitOffset >> 3;
+                var leftBitOffset = currentBitOffset & 7;
+                var boundary = 8 - leftBitOffset;
+                var rightBitOffset = currentBitOffset + boundary <= endBitOffset ? leftBitOffset + boundary : (int)(endBitOffset & 7) + 1;
+
+                // Trim byte to start and end bit index
+                var mask = (0xff >> leftBitOffset) ^ (0xff >> rightBitOffset);
+                var payload = (long)(input[byteIndex] & mask);
 
-            long pos = (((long)(curr - value)) << 3);
+                // Invalid only if equals the masked payload
+                var invalidMask = invalidPayload & mask;
 
-            long payload = 0;
-            // Adjust end so we can retrieve word
-            end = end + remainder;
+                // If transformed payload is invalid skip to next byte
+                if (payload != invalidMask)
+                {
+                    payload <<= (56 + leftBitOffset);
+                    payload = searchBit ? payload : ~payload;
 
-            // Build payload at least one byte to examine
-            if (curr < end) payload |= (long)curr[0] << 56;
-            if (curr + 1 < end) payload |= (long)curr[1] << 48;
-            if (curr + 2 < end) payload |= (long)curr[2] << 40;
-            if (curr + 3 < end) payload |= (long)curr[3] << 32;
-            if (curr + 4 < end) payload |= (long)curr[4] << 24;
-            if (curr + 5 < end) payload |= (long)curr[5] << 16;
-            if (curr + 6 < end) payload |= (long)curr[6] << 8;
-            if (curr + 7 < end) payload |= (long)curr[7];
+                    var lzcnt = (long)BitOperations.LeadingZeroCount((ulong)payload);
+                    return currentBitOffset + lzcnt;
+                }
 
-            // Transform to count leading zeros
-            payload = (bSetVal == 0) ? ~payload : payload;
+                currentBitOffset += boundary;
+            }
 
-            if (payload == mask)
-                return pos + 0;
+            return -1;
+        }
 
-            pos += (long)Lzcnt.X64.LeadingZeroCount((ulong)payload);
+        /// <summary>
+        /// Search for position of bit set in byte array using byte offset for start and end range
+        /// </summary>
+        /// <param name="input"></param>
+        /// <param name="inputLen"></param>
+        /// <param name="startOffset"></param>
+        /// <param name="endOffset"></param>
+        /// <param name="searchFor"></param>
+        /// <returns></returns>
+        private static long BitPosByteSearch(byte* input, long inputLen, long startOffset, long endOffset, byte searchFor)
+        {
+            // Initialize variables
+            var searchBit = searchFor == 1;
+            var invalidMask8 = searchBit ? 0x00 : 0xff;
+            var invalidMask32 = searchBit ? 0 : -1;
+            var invalidMask64 = searchBit ? 0L : -1L;
+            var currentStartOffset = startOffset;
+
+            while (currentStartOffset <= endOffset)
+            {
+                var remainder = endOffset - currentStartOffset + 1;
+                if (remainder >= 8)
+                {
+                    var payload = *(long*)(input + currentStartOffset);
+                    payload = BinaryPrimitives.ReverseEndianness(payload);
+
+                    // Process only if payload is valid (i.e. not all bits are set or clear based on searchFor parameter)
+                    if (payload != invalidMask64)
+                    {
+                        // Transform to count leading zeros
+                        payload = searchBit ? payload : ~payload;
+                        var lzcnt = (long)BitOperations.LeadingZeroCount((ulong)payload);
+                        return (currentStartOffset << 3) + lzcnt;
+                    }
+                    currentStartOffset += 8;
+                }
+                else if (remainder >= 4)
+                {
+                    var payload = *(int*)(input + currentStartOffset);
+                    payload = BinaryPrimitives.ReverseEndianness(payload);
+
+                    // Process only if payload is valid (i.e. not all bits are set or clear based on searchFor parameter)
+                    if (payload != invalidMask32)
+                    {
+                        // Transform to count leading zeros
+                        payload = searchBit ? payload : ~payload;
+                        var lzcnt = (long)BitOperations.LeadingZeroCount((uint)payload);
+                        return (currentStartOffset << 3) + lzcnt;
+                    }
+                    currentStartOffset += 4;
+                }
+                else
+                {
+                    // Process only if payload is valid (i.e. not all bits are set or clear based on searchFor parameter)
+                    if (input[currentStartOffset] != invalidMask8)
+                    {
+                        // Create a payload with the current byte shifted to the most significant byte position
+                        var payload = (long)input[currentStartOffset] << 56;
+                        // Transform to count leading zeros
+                        payload = searchBit ? payload : ~payload;
+                        var lzcnt = (long)BitOperations.LeadingZeroCount((ulong)payload);
+                        return (currentStartOffset << 3) + lzcnt;
+                    }
+                    currentStartOffset++;
+                }
+            }
 
-            return pos;
+            // Return -1 if no matching bit is found
+            return -1;
         }
     }
 }

libs/server/Storage/Functions/MainStore/PrivateMethods.cs

@@ -184,8 +184,14 @@ void CopyRespToWithInput(ref RawStringInput input, ref SpanByte value, ref SpanB
                         }
                     }
 
-                    var pos = BitmapManager.BitPosDriver(bpSetVal, bpStartOffset, bpEndOffset, bpOffsetType,
-                        value.ToPointer() + functionsState.etagState.etagSkippedStart, value.Length - functionsState.etagState.etagSkippedStart);
+                    var pos = BitmapManager.BitPosDriver(
+                        input: value.ToPointer() + functionsState.etagState.etagSkippedStart,
+                        inputLen: value.Length - functionsState.etagState.etagSkippedStart,
+                        startOffset: bpStartOffset,
+                        endOffset: bpEndOffset,
+                        searchFor: bpSetVal,
+                        offsetType: bpOffsetType
+                        );
                     *(long*)dst.SpanByte.ToPointer() = pos;
                     CopyRespNumber(pos, ref dst);
                     break;