Skip to content

Commit d7d2bc6

Browse files
vazoisvazois
authored
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
1 parent c65774a commit d7d2bc6

File tree

3 files changed

+338
-214
lines changed

3 files changed

+338
-214
lines changed

‎libs/server/Resp/Bitmap/BitmapManagerBitPos.cs

+128-139
Original file line numberDiff line numberDiff line change
@@ -1,185 +1,174 @@
11
// Copyright (c) Microsoft Corporation.
22
// Licensed under the MIT license.
33

4-
using System.Diagnostics;
5-
using System.Runtime.Intrinsics.X86;
4+
using System.Buffers.Binary;
5+
using System.Numerics;
66

77
namespace Garnet.server
88
{
99
public unsafe partial class BitmapManager
1010
{
1111
/// <summary>
12-
/// Find pos of bit set/clear for given bit offsets within a single byte.
12+
/// Main driver for BITPOS command
1313
/// </summary>
14-
/// <param name="value">Byte value to search within.</param>
15-
/// <param name="bSetVal">Bit value to search for (0|1).</param>
16-
/// <param name="startBitOffset">Start most significant bit offset in byte value.</param>
17-
/// <param name="endBitOffset">End most significant bit offset in bitmap.</param>
18-
/// <returns></returns>
19-
private static long BitPosIndexBitSingleByteSearch(byte value, byte bSetVal, int startBitOffset = 0, int endBitOffset = 8)
20-
{
21-
Debug.Assert(startBitOffset >= 0 && startBitOffset <= 8);
22-
Debug.Assert(endBitOffset >= 0 && endBitOffset <= 8);
23-
bool bflag = (bSetVal == 0);
24-
long mask = bflag ? -1 : 0;
25-
26-
int leftBitIndex = 1 << (8 - startBitOffset);
27-
int rightBitIndex = 1 << (8 - endBitOffset);
28-
29-
// Create extraction mask
30-
long extract = leftBitIndex - rightBitIndex;
31-
32-
long payload = (long)(value & extract) << 56;
33-
// Trim leading bits
34-
payload = payload << startBitOffset;
35-
36-
// Transform to count leading zeros
37-
payload = bflag ? ~payload : payload;
38-
39-
// Return not found
40-
if (payload == mask) return -1;
41-
42-
return (long)Lzcnt.X64.LeadingZeroCount((ulong)payload);
43-
}
44-
45-
/// <summary>
46-
/// Find pos of bit set/clear for given bit offset.
47-
/// </summary>
48-
/// <param name="value">Pointer to start of bitmap.</param>
49-
/// <param name="bSetVal">Bit value to search for (0|1).</param>
50-
/// <param name="offset">Bit offset in bitmap.</param>
51-
/// <returns></returns>
52-
private static long BitPosIndexBitSearch(byte* value, byte bSetVal, long offset = 0)
53-
{
54-
bool bflag = (bSetVal == 0);
55-
long mask = bflag ? -1 : 0;
56-
long startByteOffset = (offset / 8);
57-
int bitOffset = (int)(offset & 7);
58-
59-
long payload = (long)value[startByteOffset] << 56;
60-
// Trim leading bits
61-
payload = payload << bitOffset;
62-
63-
// Transform to count leading zeros
64-
payload = bflag ? ~payload : payload;
65-
66-
// Return not found
67-
if (payload == mask)
68-
return -1;
69-
70-
return (long)Lzcnt.X64.LeadingZeroCount((ulong)payload);
71-
}
72-
73-
/// <summary>
74-
/// Main driver for bit position command.
75-
/// </summary>
76-
/// <param name="setVal"></param>
14+
/// <param name="input"></param>
15+
/// <param name="inputLen"></param>
7716
/// <param name="startOffset"></param>
7817
/// <param name="endOffset"></param>
18+
/// <param name="searchFor"></param>
7919
/// <param name="offsetType"></param>
80-
/// <param name="value">Pointer to start of bitmap.</param>
81-
/// <param name="valLen">Length of bitmap.</param>
8220
/// <returns></returns>
83-
public static long BitPosDriver(byte setVal, long startOffset, long endOffset, byte offsetType, byte* value, int valLen)
21+
public static long BitPosDriver(byte* input, int inputLen, long startOffset, long endOffset, byte searchFor, byte offsetType)
8422
{
8523
if (offsetType == 0x0)
8624
{
87-
startOffset = startOffset < 0 ? ProcessNegativeOffset(startOffset, valLen) : startOffset;
88-
endOffset = endOffset < 0 ? ProcessNegativeOffset(endOffset, valLen) : endOffset;
25+
startOffset = startOffset < 0 ? ProcessNegativeOffset(startOffset, inputLen) : startOffset;
26+
endOffset = endOffset < 0 ? ProcessNegativeOffset(endOffset, inputLen) : endOffset;
8927

90-
if (startOffset >= valLen) // If startOffset greater that valLen always bitpos -1
28+
if (startOffset >= inputLen) // If startOffset greater that valLen always bitpos -1
9129
return -1;
9230

9331
if (startOffset > endOffset) // If start offset beyond endOffset return 0
9432
return -1;
9533

96-
endOffset = endOffset >= valLen ? valLen : endOffset;
97-
long pos = BitPosByte(value, setVal, startOffset, endOffset);
98-
// check if position is exceeding the last byte in acceptable range
99-
return pos >= ((endOffset + 1) * 8) ? -1 : pos;
34+
endOffset = endOffset >= inputLen ? inputLen : endOffset;
35+
// BYTE search
36+
return BitPosByteSearch(input, inputLen, startOffset, endOffset, searchFor);
10037
}
101-
102-
startOffset = startOffset < 0 ? ProcessNegativeOffset(startOffset, valLen * 8) : startOffset;
103-
endOffset = endOffset < 0 ? ProcessNegativeOffset(endOffset, valLen * 8) : endOffset;
104-
105-
var startByte = (startOffset / 8);
106-
var endByte = (endOffset / 8);
107-
if (startByte == endByte)
38+
else
10839
{
109-
// Search only inside single byte for pos
110-
var leftBitIndex = (int)(startOffset & 7);
111-
var rightBitIndex = (int)((endOffset + 1) & 7);
112-
var _ipos = BitPosIndexBitSingleByteSearch(value[startByte], setVal, leftBitIndex, rightBitIndex);
113-
return _ipos == -1 ? _ipos : startOffset + _ipos;
114-
}
40+
startOffset = startOffset < 0 ? ProcessNegativeOffset(startOffset, inputLen * 8) : startOffset;
41+
endOffset = endOffset < 0 ? ProcessNegativeOffset(endOffset, inputLen * 8) : endOffset;
42+
43+
var startByteIndex = startOffset >> 3;
44+
var endByteIndex = endOffset >> 3;
11545

116-
// Search prefix and terminate if found position of bit
117-
var _ppos = BitPosIndexBitSearch(value, setVal, startOffset);
118-
if (_ppos != -1) return startOffset + _ppos;
46+
if (startByteIndex >= inputLen) // If startOffset greater that valLen always bitpos -1
47+
return -1;
11948

120-
// Adjust offsets to skip first and last byte
121-
var _startOffset = (startOffset / 8) + 1;
122-
var _endOffset = (endOffset / 8) - 1;
123-
var _bpos = BitPosByte(value, setVal, _startOffset, _endOffset);
49+
if (startByteIndex > endByteIndex) // If start offset beyond endOffset return 0
50+
return -1;
12451

125-
if (_bpos != -1 && _bpos < (_endOffset + 1) * 8) return _bpos;
52+
endOffset = endByteIndex >= inputLen ? inputLen << 3 : endOffset;
12653

127-
// Search suffix
128-
var _spos = BitPosIndexBitSearch(value, setVal, endOffset);
129-
return _spos;
54+
// BIT search
55+
return BitPosBitSearch(input, inputLen, startOffset, endOffset, searchFor);
56+
}
13057
}
13158

13259
/// <summary>
133-
/// Find pos of set/clear bit in a sequence of bytes.
60+
/// Search for position of bit set in byte array using bit offset for start and end range
13461
/// </summary>
135-
/// <param name="value">Pointer to start of bitmap.</param>
136-
/// <param name="bSetVal">The bit value to search for (0 for cleared bit or 1 for set bit).</param>
137-
/// <param name="startOffset">Starting offset into bitmap.</param>
138-
/// <param name="endOffset">End offset into bitmap.</param>
62+
/// <param name="input"></param>
63+
/// <param name="inputLen"></param>
64+
/// <param name="startBitOffset"></param>
65+
/// <param name="endBitOffset"></param>
66+
/// <param name="searchFor"></param>
13967
/// <returns></returns>
140-
private static long BitPosByte(byte* value, byte bSetVal, long startOffset, long endOffset)
68+
private static long BitPosBitSearch(byte* input, long inputLen, long startBitOffset, long endBitOffset, byte searchFor)
14169
{
142-
// Mask set to look for 0 or 1 depending on clear/set flag
143-
bool bflag = (bSetVal == 0);
144-
long mask = bflag ? -1 : 0;
145-
long len = (endOffset - startOffset) + 1;
146-
long remainder = len & 7;
147-
byte* curr = value + startOffset;
148-
byte* end = curr + (len - remainder);
149-
150-
// Search for first word not matching mask.
151-
while (curr < end)
70+
var searchBit = searchFor == 1;
71+
var invalidPayload = (byte)(searchBit ? 0x00 : 0xff);
72+
var currentBitOffset = (int)startBitOffset;
73+
while (currentBitOffset <= endBitOffset)
15274
{
153-
long v = *(long*)(curr);
154-
if (v != mask) break;
155-
curr += 8;
156-
}
75+
var byteIndex = currentBitOffset >> 3;
76+
var leftBitOffset = currentBitOffset & 7;
77+
var boundary = 8 - leftBitOffset;
78+
var rightBitOffset = currentBitOffset + boundary <= endBitOffset ? leftBitOffset + boundary : (int)(endBitOffset & 7) + 1;
79+
80+
// Trim byte to start and end bit index
81+
var mask = (0xff >> leftBitOffset) ^ (0xff >> rightBitOffset);
82+
var payload = (long)(input[byteIndex] & mask);
15783

158-
long pos = (((long)(curr - value)) << 3);
84+
// Invalid only if equals the masked payload
85+
var invalidMask = invalidPayload & mask;
15986

160-
long payload = 0;
161-
// Adjust end so we can retrieve word
162-
end = end + remainder;
87+
// If transformed payload is invalid skip to next byte
88+
if (payload != invalidMask)
89+
{
90+
payload <<= (56 + leftBitOffset);
91+
payload = searchBit ? payload : ~payload;
16392

164-
// Build payload at least one byte to examine
165-
if (curr < end) payload |= (long)curr[0] << 56;
166-
if (curr + 1 < end) payload |= (long)curr[1] << 48;
167-
if (curr + 2 < end) payload |= (long)curr[2] << 40;
168-
if (curr + 3 < end) payload |= (long)curr[3] << 32;
169-
if (curr + 4 < end) payload |= (long)curr[4] << 24;
170-
if (curr + 5 < end) payload |= (long)curr[5] << 16;
171-
if (curr + 6 < end) payload |= (long)curr[6] << 8;
172-
if (curr + 7 < end) payload |= (long)curr[7];
93+
var lzcnt = (long)BitOperations.LeadingZeroCount((ulong)payload);
94+
return currentBitOffset + lzcnt;
95+
}
17396

174-
// Transform to count leading zeros
175-
payload = (bSetVal == 0) ? ~payload : payload;
97+
currentBitOffset += boundary;
98+
}
17699

177-
if (payload == mask)
178-
return pos + 0;
100+
return -1;
101+
}
179102

180-
pos += (long)Lzcnt.X64.LeadingZeroCount((ulong)payload);
103+
/// <summary>
104+
/// Search for position of bit set in byte array using byte offset for start and end range
105+
/// </summary>
106+
/// <param name="input"></param>
107+
/// <param name="inputLen"></param>
108+
/// <param name="startOffset"></param>
109+
/// <param name="endOffset"></param>
110+
/// <param name="searchFor"></param>
111+
/// <returns></returns>
112+
private static long BitPosByteSearch(byte* input, long inputLen, long startOffset, long endOffset, byte searchFor)
113+
{
114+
// Initialize variables
115+
var searchBit = searchFor == 1;
116+
var invalidMask8 = searchBit ? 0x00 : 0xff;
117+
var invalidMask32 = searchBit ? 0 : -1;
118+
var invalidMask64 = searchBit ? 0L : -1L;
119+
var currentStartOffset = startOffset;
120+
121+
while (currentStartOffset <= endOffset)
122+
{
123+
var remainder = endOffset - currentStartOffset + 1;
124+
if (remainder >= 8)
125+
{
126+
var payload = *(long*)(input + currentStartOffset);
127+
payload = BinaryPrimitives.ReverseEndianness(payload);
128+
129+
// Process only if payload is valid (i.e. not all bits are set or clear based on searchFor parameter)
130+
if (payload != invalidMask64)
131+
{
132+
// Transform to count leading zeros
133+
payload = searchBit ? payload : ~payload;
134+
var lzcnt = (long)BitOperations.LeadingZeroCount((ulong)payload);
135+
return (currentStartOffset << 3) + lzcnt;
136+
}
137+
currentStartOffset += 8;
138+
}
139+
else if (remainder >= 4)
140+
{
141+
var payload = *(int*)(input + currentStartOffset);
142+
payload = BinaryPrimitives.ReverseEndianness(payload);
143+
144+
// Process only if payload is valid (i.e. not all bits are set or clear based on searchFor parameter)
145+
if (payload != invalidMask32)
146+
{
147+
// Transform to count leading zeros
148+
payload = searchBit ? payload : ~payload;
149+
var lzcnt = (long)BitOperations.LeadingZeroCount((uint)payload);
150+
return (currentStartOffset << 3) + lzcnt;
151+
}
152+
currentStartOffset += 4;
153+
}
154+
else
155+
{
156+
// Process only if payload is valid (i.e. not all bits are set or clear based on searchFor parameter)
157+
if (input[currentStartOffset] != invalidMask8)
158+
{
159+
// Create a payload with the current byte shifted to the most significant byte position
160+
var payload = (long)input[currentStartOffset] << 56;
161+
// Transform to count leading zeros
162+
payload = searchBit ? payload : ~payload;
163+
var lzcnt = (long)BitOperations.LeadingZeroCount((ulong)payload);
164+
return (currentStartOffset << 3) + lzcnt;
165+
}
166+
currentStartOffset++;
167+
}
168+
}
181169

182-
return pos;
170+
// Return -1 if no matching bit is found
171+
return -1;
183172
}
184173
}
185174
}

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

+8-2
Original file line numberDiff line numberDiff line change
@@ -184,8 +184,14 @@ void CopyRespToWithInput(ref RawStringInput input, ref SpanByte value, ref SpanB
184184
}
185185
}
186186

187-
var pos = BitmapManager.BitPosDriver(bpSetVal, bpStartOffset, bpEndOffset, bpOffsetType,
188-
value.ToPointer() + functionsState.etagState.etagSkippedStart, value.Length - functionsState.etagState.etagSkippedStart);
187+
var pos = BitmapManager.BitPosDriver(
188+
input: value.ToPointer() + functionsState.etagState.etagSkippedStart,
189+
inputLen: value.Length - functionsState.etagState.etagSkippedStart,
190+
startOffset: bpStartOffset,
191+
endOffset: bpEndOffset,
192+
searchFor: bpSetVal,
193+
offsetType: bpOffsetType
194+
);
189195
*(long*)dst.SpanByte.ToPointer() = pos;
190196
CopyRespNumber(pos, ref dst);
191197
break;

0 commit comments

Comments
 (0)