Line data Source code
1 : // Copyright 2024 The LevelDB-Go and Pebble Authors. All rights reserved. Use
2 : // of this source code is governed by a BSD-style license that can be found in
3 : // the LICENSE file.
4 :
5 : package colblk
6 :
7 : import (
8 : "fmt"
9 : "io"
10 : "math"
11 : "math/bits"
12 : "strings"
13 : "unsafe"
14 :
15 : "github.com/cockroachdb/errors"
16 : "github.com/cockroachdb/pebble/internal/binfmt"
17 : "github.com/cockroachdb/pebble/internal/invariants"
18 : "github.com/cockroachdb/pebble/internal/treeprinter"
19 : )
20 :
21 : // Bitmap is a bitmap structure built on a []uint64. A bitmap utilizes ~1
22 : // physical bit/logical bit (~0.125 bytes/row). The bitmap is encoded into an
23 : // 8-byte aligned array of 64-bit words which is (nRows+63)/64 words in length.
24 : //
25 : // A summary bitmap is stored after the primary bitmap in which each bit in the
26 : // summary bitmap corresponds to 1 word in the primary bitmap. A bit is set in
27 : // the summary bitmap if the corresponding word in the primary bitmap is
28 : // non-zero. The summary bitmap accelerates predecessor and successor
29 : // operations.
30 : type Bitmap struct {
31 : // data contains the bitmap data, according to defaultBitmapEncoding, or it
32 : // is nil if the bitmap is all zeros.
33 : data UnsafeRawSlice[uint64]
34 : bitCount int
35 : }
36 :
37 : // Assert that Bitmap implements Array[bool].
38 : var _ Array[bool] = Bitmap{}
39 :
40 : // DecodeBitmap decodes the structure of a Bitmap and returns a Bitmap that
41 : // reads from b supporting bitCount logical bits. No bounds checking is
42 : // performed, so the caller must guarantee the bitmap is appropriately sized and
43 : // the provided bitCount correctly identifies the number of bits in the bitmap.
44 1 : func DecodeBitmap(b []byte, off uint32, bitCount int) (bitmap Bitmap, endOffset uint32) {
45 1 : encoding := bitmapEncoding(b[off])
46 1 : off++
47 1 : if encoding == zeroBitmapEncoding {
48 1 : return Bitmap{bitCount: bitCount}, off
49 1 : }
50 1 : off = align(off, align64)
51 1 : sz := bitmapRequiredSize(bitCount)
52 1 : if len(b) < int(off)+sz {
53 0 : panic(errors.AssertionFailedf("bitmap of %d bits requires at least %d bytes; provided with %d-byte slice",
54 0 : bitCount, bitmapRequiredSize(bitCount), len(b[off:])))
55 : }
56 1 : return Bitmap{
57 1 : data: makeUnsafeRawSlice[uint64](unsafe.Pointer(&b[off])),
58 1 : bitCount: bitCount,
59 1 : }, off + uint32(sz)
60 : }
61 :
62 : // Assert that DecodeBitmap implements DecodeFunc.
63 : var _ DecodeFunc[Bitmap] = DecodeBitmap
64 :
65 : // At returns true if the bit at position i is set and false otherwise.
66 : //
67 : //gcassert:inline
68 1 : func (b Bitmap) At(i int) bool {
69 1 : if b.data.ptr == nil {
70 1 : // zero bitmap case.
71 1 : return false
72 1 : }
73 : // Inline b.data.At(i/64).
74 : // The offset of the correct word is i / 64 * 8 = (i >> 3) &^ 0b111
75 1 : const mask = ^uintptr(0b111)
76 1 : val := *(*uint64)(unsafe.Pointer(uintptr(b.data.ptr) + (uintptr(i)>>3)&mask))
77 1 : return val&(1<<(uint(i)&63)) != 0
78 : }
79 :
80 : // SeekSetBitGE returns the next bit greater than or equal to i set in the bitmap.
81 : // The i parameter must be ≥ 0. Returns the number of bits
82 : // represented by the bitmap if no next bit is set (or if i >= bitCount).
83 1 : func (b Bitmap) SeekSetBitGE(i int) int {
84 1 : if b.data.ptr == nil || i >= b.bitCount {
85 1 : // Zero bitmap case.
86 1 : return b.bitCount
87 1 : }
88 :
89 1 : wordIdx := i >> 6 // i/64
90 1 : // Fast path for common case of reasonably dense bitmaps; if the there's a
91 1 : // bit ≥ i set in the same word, return it.
92 1 : if next := nextBitInWord(b.data.At(wordIdx), uint(i)&63); next < 64 {
93 1 : return wordIdx<<6 + next
94 1 : }
95 :
96 : // Consult summary structure to find the next word with a set bit. The word
97 : // we just checked (wordIdx) is represented by the wordIdx%64'th bit in the
98 : // wordIdx/64'th summary word. We want to know if any of the other later
99 : // words that are summarized together have a set bit. We call [nextInWord]
100 : // on the summary word to get the index of which word has a set bit, if any.
101 1 : summaryTableOffset, summaryTableEnd := b.summaryTableBounds()
102 1 : summaryWordIdx := summaryTableOffset + wordIdx>>6
103 1 : summaryNext := nextBitInWord(b.data.At(summaryWordIdx), uint(wordIdx%64)+1)
104 1 : // If [summaryNext] == 64, then there are no set bits in any of the earlier
105 1 : // words represented by the summary word at [summaryWordIdx]. In that case,
106 1 : // we need to keep scanning the summary table forwards.
107 1 : if summaryNext == 64 {
108 1 : for summaryWordIdx++; ; summaryWordIdx++ {
109 1 : // When we fall off the end of the summary table, we've determined
110 1 : // there are no set bits after i across the entirety of the bitmap.
111 1 : if summaryWordIdx >= summaryTableEnd {
112 1 : return b.bitCount
113 1 : }
114 0 : if summaryWord := b.data.At(summaryWordIdx); summaryWord != 0 {
115 0 : summaryNext = bits.TrailingZeros64(summaryWord)
116 0 : break
117 : }
118 : }
119 : }
120 : // The summary word index and the summaryNext together tell us which word
121 : // has a set bit. The number of leading zeros in the word itself tell us
122 : // which bit is set.
123 1 : wordIdx = ((summaryWordIdx - summaryTableOffset) << 6) + summaryNext
124 1 : return (wordIdx << 6) + bits.TrailingZeros64(b.data.At(wordIdx))
125 : }
126 :
127 : // SeekSetBitLE returns the previous bit less than or equal to i set in the
128 : // bitmap. The i parameter must be in [0, bitCount). Returns -1 if no previous
129 : // bit is set.
130 0 : func (b Bitmap) SeekSetBitLE(i int) int {
131 0 : if b.data.ptr == nil {
132 0 : // Zero bitmap case.
133 0 : return -1
134 0 : }
135 0 : wordIdx := i >> 6 // i/64
136 0 : // Fast path for common case of reasonably dense bitmaps; if the there's a
137 0 : // bit ≤ i set in the same word, return it.
138 0 : if prev := prevBitInWord(b.data.At(wordIdx), uint(i)&63); prev >= 0 {
139 0 : return (wordIdx << 6) + prev
140 0 : }
141 :
142 : // Consult summary structure to find the next word with a set bit. The word
143 : // we just checked (wordIdx) is represented by the wordIdx%64'th bit in the
144 : // wordIdx/64'th summary word. We want to know if any of other earlier words
145 : // that are summarized together have a set bit. We call [prevInWord] on the
146 : // summary word to get the index of which word has a set bit, if any.
147 0 : summaryTableOffset, _ := b.summaryTableBounds()
148 0 : summaryWordIdx := summaryTableOffset + wordIdx>>6
149 0 : summaryPrev := prevBitInWord(b.data.At(summaryWordIdx), uint(wordIdx%64)-1)
150 0 : // If [summaryPrev] is negative, then there are no set bits in any of the
151 0 : // earlier words represented by the summary word at [summaryWordIdx]. In
152 0 : // that case, we need to keep scanning the summary table backwards.
153 0 : if summaryPrev < 0 {
154 0 : for summaryWordIdx--; ; summaryWordIdx-- {
155 0 : // When we fall below the beginning of the summary table, we've
156 0 : // determined there are no set bits before i across the entirety of
157 0 : // the bitmap.
158 0 : if summaryWordIdx < summaryTableOffset {
159 0 : return -1
160 0 : }
161 0 : if summaryWord := b.data.At(summaryWordIdx); summaryWord != 0 {
162 0 : summaryPrev = 63 - bits.LeadingZeros64(summaryWord)
163 0 : break
164 : }
165 : }
166 : }
167 : // The summary word index and the summary prev together tell us which word
168 : // has a set bit. The number of trailing zeros in the word itself tell us
169 : // which bit is set.
170 0 : wordIdx = ((summaryWordIdx - summaryTableOffset) << 6) + summaryPrev
171 0 : return (wordIdx << 6) + 63 - bits.LeadingZeros64(b.data.At(wordIdx))
172 : }
173 :
174 : // SeekUnsetBitGE returns the next bit greater than or equal to i that is unset
175 : // in the bitmap. The i parameter must be in [0, bitCount). Returns the number
176 : // of bits represented by the bitmap if no next bit is unset.
177 1 : func (b Bitmap) SeekUnsetBitGE(i int) int {
178 1 : if b.data.ptr == nil {
179 0 : // Zero bitmap case.
180 0 : return i
181 0 : }
182 :
183 1 : wordIdx := i >> 6 // i/64
184 1 : // If the there's a bit ≥ i unset in the same word, return it.
185 1 : if next := nextBitInWord(^b.data.At(wordIdx), uint(i)&63); next < 64 {
186 1 : return wordIdx<<6 + next
187 1 : }
188 1 : numWords := (b.bitCount + 63) >> 6
189 1 : var word uint64
190 1 : for wordIdx++; ; wordIdx++ {
191 1 : if wordIdx >= numWords {
192 0 : return b.bitCount
193 0 : }
194 1 : word = b.data.At(wordIdx)
195 1 : if word != math.MaxUint64 {
196 1 : break
197 : }
198 : }
199 1 : return wordIdx<<6 + bits.TrailingZeros64(^word)
200 : }
201 :
202 : // SeekUnsetBitLE returns the previous bit less than or equal to i set in the
203 : // bitmap. The i parameter must be in [0, bitCount). Returns -1 if no previous
204 : // bit is unset.
205 1 : func (b Bitmap) SeekUnsetBitLE(i int) int {
206 1 : if b.data.ptr == nil {
207 0 : // Zero bitmap case.
208 0 : return i
209 0 : }
210 :
211 1 : wordIdx := i >> 6 // i/64
212 1 : // If there's a bit ≤ i unset in the same word, return it.
213 1 : if prev := prevBitInWord(^b.data.At(wordIdx), uint(i)&63); prev >= 0 {
214 1 : return (wordIdx << 6) + prev
215 1 : }
216 1 : var word uint64
217 1 : for wordIdx--; ; wordIdx-- {
218 1 : if wordIdx < 0 {
219 1 : return -1
220 1 : }
221 1 : word = b.data.At(wordIdx)
222 1 : if word != math.MaxUint64 {
223 1 : break
224 : }
225 : }
226 1 : return (wordIdx << 6) + 63 - bits.LeadingZeros64(^word)
227 : }
228 :
229 : // summaryTableBounds returns the indexes of the bitmap words containing the
230 : // summary table. The summary table's words lie within [startOffset, endOffset).
231 1 : func (b Bitmap) summaryTableBounds() (startOffset, endOffset int) {
232 1 : startOffset = (b.bitCount + 63) >> 6
233 1 : endOffset = startOffset + (startOffset+63)>>6
234 1 : return startOffset, endOffset
235 1 : }
236 :
237 : // String returns a string representation of the entire bitmap.
238 0 : func (b Bitmap) String() string {
239 0 : var sb strings.Builder
240 0 : for w := 0; w < (b.bitCount+63)/64; w++ {
241 0 : fmt.Fprintf(&sb, "%064b", b.data.At(w))
242 0 : }
243 0 : return sb.String()
244 : }
245 :
246 : // BitmapBuilder constructs a Bitmap. Bits default to false.
247 : type BitmapBuilder struct {
248 : words []uint64
249 : // minNonZeroRowCount is the row count at which the bitmap should begin to
250 : // use the defaultBitmapEncoding (as opposed to the zeroBitmapEncoding).
251 : // It's updated on the first call to Set and defaults to zero.
252 : minNonZeroRowCount int
253 : }
254 :
255 : type bitmapEncoding uint8
256 :
257 : const (
258 : // defaultBitmapEncoding encodes the bitmap using ⌈n/64⌉ words followed by
259 : // ⌈⌈n/64⌉/64⌉ summary words.
260 : defaultBitmapEncoding bitmapEncoding = iota
261 : // zeroBitmapEncoding is used for the special case when the bitmap is empty.
262 : zeroBitmapEncoding
263 : )
264 :
265 : // Assert that BitmapBuilder implements ColumnWriter.
266 : var _ ColumnWriter = (*BitmapBuilder)(nil)
267 :
268 : // bitmapRequiredSize returns the size of an encoded bitmap in bytes, using the
269 : // defaultBitmapEncoding.
270 1 : func bitmapRequiredSize(total int) int {
271 1 : nWords := (total + 63) >> 6 // divide by 64
272 1 : nSummaryWords := (nWords + 63) >> 6 // divide by 64
273 1 : return (nWords + nSummaryWords) << 3 // multiply by 8
274 1 : }
275 :
276 : // Set sets the bit at position i to true.
277 1 : func (b *BitmapBuilder) Set(i int) {
278 1 : // Update minNonZeroRowCount if necessary. This is used to determine whether
279 1 : // the bitmap should be encoded using the all-zeros encoding.
280 1 : if b.isZero(i + 1) {
281 1 : b.minNonZeroRowCount = i + 1
282 1 : }
283 1 : w := i >> 6 // divide by 64
284 1 : for len(b.words) <= w {
285 1 : // We append zeros because if b.words has additional capacity, it has
286 1 : // not been zeroed.
287 1 : b.words = append(b.words, 0)
288 1 : }
289 1 : b.words[w] |= 1 << uint(i&63)
290 : }
291 :
292 : // isZero returns true if no bits are set and Invert was not called.
293 : //
294 : //gcassert:inline
295 1 : func (b *BitmapBuilder) isZero(rows int) bool {
296 1 : return b.minNonZeroRowCount == 0 || rows < b.minNonZeroRowCount
297 1 : }
298 :
299 : // Reset resets the bitmap to the empty state.
300 1 : func (b *BitmapBuilder) Reset() {
301 1 : if invariants.Sometimes(50) {
302 1 : // Sometimes trash the bitmap with all ones to catch bugs that assume
303 1 : // b.words is zeroed.
304 1 : for i := 0; i < len(b.words); i++ {
305 1 : b.words[i] = ^uint64(0)
306 1 : }
307 : }
308 :
309 : // NB: We don't zero the contents of b.words. When the BitmapBuilder reuses
310 : // b.words, it must ensure it zeroes the contents as necessary.
311 1 : b.words = b.words[:0]
312 1 : b.minNonZeroRowCount = 0
313 : }
314 :
315 : // NumColumns implements the ColumnWriter interface.
316 1 : func (b *BitmapBuilder) NumColumns() int { return 1 }
317 :
318 : // DataType implements the ColumnWriter interface.
319 1 : func (b *BitmapBuilder) DataType(int) DataType { return DataTypeBool }
320 :
321 : // Size implements the ColumnWriter interface.
322 1 : func (b *BitmapBuilder) Size(rows int, offset uint32) uint32 {
323 1 : // First byte will be the encoding type.
324 1 : offset++
325 1 : if b.isZero(rows) {
326 1 : return offset
327 1 : }
328 1 : offset = align(offset, align64)
329 1 : return offset + uint32(bitmapRequiredSize(rows))
330 : }
331 :
332 : // InvertedSize returns the size of the encoded bitmap, assuming Invert will be called.
333 1 : func (b *BitmapBuilder) InvertedSize(rows int, offset uint32) uint32 {
334 1 : // First byte will be the encoding type.
335 1 : offset++
336 1 : // An inverted bitmap will never use all-zeros encoding (even if it happens to
337 1 : // be all zero).
338 1 : offset = align(offset, align64)
339 1 : return offset + uint32(bitmapRequiredSize(rows))
340 1 : }
341 :
342 : // Invert inverts the bitmap, setting all bits that are not set and clearing all
343 : // bits that are set. If the bitmap's tail is sparse and is not large enough to
344 : // represent nRows rows, it's first materialized.
345 : //
346 : // Note that Invert can affect the Size of the bitmap. Use InvertedSize() if you
347 : // intend to invert the bitmap before finishing.
348 1 : func (b *BitmapBuilder) Invert(nRows int) {
349 1 : // Inverted bitmaps never use the all-zero encoding, so we set
350 1 : // rowCountIncludingFirstSetBit to 1 so that as long as the bitmap is
351 1 : // finished encoding any rows at all, it uses the default encoding.
352 1 : b.minNonZeroRowCount = 1
353 1 : // If the tail of b is sparse, fill in zeroes before inverting.
354 1 : nBitmapWords := (nRows + 63) >> 6
355 1 : for len(b.words) < nBitmapWords {
356 1 : // We append zeros because if b.words has additional capacity, it has
357 1 : // not been zeroed.
358 1 : b.words = append(b.words, 0)
359 1 : }
360 1 : b.words = b.words[:nBitmapWords]
361 1 : for i := range b.words {
362 1 : b.words[i] = ^b.words[i]
363 1 : }
364 : }
365 :
366 : // Finish finalizes the bitmap, computing the per-word summary bitmap and
367 : // writing the resulting data to buf at offset.
368 1 : func (b *BitmapBuilder) Finish(col, nRows int, offset uint32, buf []byte) uint32 {
369 1 : if b.isZero(nRows) {
370 1 : buf[offset] = byte(zeroBitmapEncoding)
371 1 : return offset + 1
372 1 : }
373 1 : buf[offset] = byte(defaultBitmapEncoding)
374 1 : offset++
375 1 : offset = alignWithZeroes(buf, offset, align64)
376 1 : dest := makeUnsafeRawSlice[uint64](unsafe.Pointer(&buf[offset]))
377 1 :
378 1 : nBitmapWords := (nRows + 63) >> 6
379 1 : // Truncate the bitmap to the number of words required to represent nRows.
380 1 : // The caller may have written more bits than nRows and no longer cares to
381 1 : // write them out.
382 1 : if len(b.words) > nBitmapWords {
383 0 : b.words = b.words[:nBitmapWords]
384 0 : }
385 : // Ensure the last word of the bitmap does not contain any set bits beyond
386 : // the last row. This is not just for determinism but also to ensure that
387 : // the summary bitmap is correct (which is necessary for Bitmap.SeekSetBitGE
388 : // correctness).
389 1 : if i := nRows % 64; len(b.words) >= nBitmapWords && i != 0 {
390 1 : b.words[nBitmapWords-1] &= (1 << i) - 1
391 1 : }
392 :
393 : // Copy all the words of the bitmap into the destination buffer.
394 1 : offset += uint32(copy(dest.Slice(len(b.words)), b.words)) << align64Shift
395 1 :
396 1 : // The caller may have written fewer than nRows rows if the tail is all
397 1 : // zeroes, relying on these bits being implicitly zero. If the tail of b is
398 1 : // sparse, fill in zeroes.
399 1 : for i := len(b.words); i < nBitmapWords; i++ {
400 1 : dest.set(i, 0)
401 1 : offset += align64
402 1 : }
403 :
404 : // Add the summary bitmap.
405 1 : nSummaryWords := (nBitmapWords + 63) >> 6
406 1 : for i := 0; i < nSummaryWords; i++ {
407 1 : wordsOff := (i << 6) // i*64
408 1 : nWords := min(64, len(b.words)-wordsOff)
409 1 : var summaryWord uint64
410 1 : for j := 0; j < nWords; j++ {
411 1 : if (b.words)[wordsOff+j] != 0 {
412 1 : summaryWord |= 1 << j
413 1 : }
414 : }
415 1 : dest.set(nBitmapWords+i, summaryWord)
416 : }
417 1 : return offset + uint32(nSummaryWords)<<align64Shift
418 : }
419 :
420 : // WriteDebug implements the ColumnWriter interface.
421 0 : func (b *BitmapBuilder) WriteDebug(w io.Writer, rows int) {
422 0 : // TODO(jackson): Add more detailed debugging information.
423 0 : fmt.Fprint(w, "bitmap")
424 0 : }
425 :
426 0 : func bitmapToBinFormatter(f *binfmt.Formatter, tp treeprinter.Node, rows int) {
427 0 : encoding := bitmapEncoding(f.PeekUint(1))
428 0 : if encoding == zeroBitmapEncoding {
429 0 : f.HexBytesln(1, "zero bitmap encoding")
430 0 : f.ToTreePrinter(tp)
431 0 : return
432 0 : }
433 0 : if encoding != defaultBitmapEncoding {
434 0 : panic(fmt.Sprintf("unknown bitmap encoding %d", encoding))
435 : }
436 0 : f.HexBytesln(1, "default bitmap encoding")
437 0 : if aligned := align(f.RelativeOffset(), 8); aligned-f.RelativeOffset() != 0 {
438 0 : f.HexBytesln(aligned-f.RelativeOffset(), "padding to align to 64-bit boundary")
439 0 : }
440 0 : bitmapWords := (rows + 63) / 64
441 0 : for i := 0; i < bitmapWords; i++ {
442 0 : f.Line(8).Append("b ").Binary(8).Done("bitmap word %d", i)
443 0 : }
444 0 : summaryWords := (bitmapWords + 63) / 64
445 0 : for i := 0; i < summaryWords; i++ {
446 0 : f.Line(8).Append("b ").Binary(8).Done("bitmap summary word %d-%d", i*64, i*64+63)
447 0 : }
448 0 : f.ToTreePrinter(tp)
449 : }
450 :
451 : // nextBitInWord returns the index of the smallest set bit with an index ≥ bit
452 : // within the provided word. The given index must be in the [0, 63] interval.
453 : // The returned index is an index local to the word. Returns 64 if no set bit is
454 : // found.
455 1 : func nextBitInWord(word uint64, bit uint) int {
456 1 : // We want to find the index of the next set bit. We can accomplish this
457 1 : // by clearing the trailing `bit` bits from the word and counting the
458 1 : // number of trailing zeros. For example, consider the word and bit=37:
459 1 : //
460 1 : // word: 1010101010111111111110000001110101010101011111111111000000111011
461 1 : //
462 1 : // 1<<bit: 0000000000000000000000000010000000000000000000000000000000000000
463 1 : // 1<<bit-1: 0000000000000000000000000001111111111111111111111111111111111111
464 1 : // ^1<<bit-1: 1111111111111111111111111110000000000000000000000000000000000000
465 1 : // word&^1<<bit-1: 1010101010111111111110000000000000000000000000000000000000000000
466 1 : //
467 1 : // Counting the trailing zeroes of this last value gives us 43. For
468 1 : // visualizing, 1<<43 is:
469 1 : //
470 1 : // 0000000000000000000010000000000000000000000000000000000000000000
471 1 : //
472 1 : return bits.TrailingZeros64(word &^ ((1 << bit) - 1))
473 1 : }
474 :
475 : // prevBitInWord returns the index of the largest set bit ≤ bit within the
476 : // provided word. The given bit index must be in the [0, 63] interval. The
477 : // returned bit index is an index local to the word. Returns -1 if no set bit is
478 : // found.
479 1 : func prevBitInWord(word uint64, bit uint) int {
480 1 : // We want to find the index of the previous set bit. We can accomplish
481 1 : // this by clearing the leading `bit` bits from the word and counting
482 1 : // the number of leading zeros. For example, consider the word and
483 1 : // bit=42:
484 1 : //
485 1 : // word: 1010101010111111111110000001110101010101011111111111000000111011
486 1 : //
487 1 : // 1<<(bit+1): 0000000000000000000010000000000000000000000000000000000000000000
488 1 : // 1<<(bit+1)-1: 0000000000000000000001111111111111111111111111111111111111111111
489 1 : // word&1<<(bit+1)-1: 0000000000000000000000000001110101010101011111111111000000111011
490 1 : //
491 1 : // Counting the leading zeroes of this last value gives us 27 leading
492 1 : // zeros. 63-27 gives index 36. For visualizing, 1<<36 is:
493 1 : //
494 1 : // 0000000000000000000000000001000000000000000000000000000000000000
495 1 : //
496 1 : return 63 - bits.LeadingZeros64(word&((1<<(bit+1))-1))
497 1 : }
|
