github.com/outcaste-io/sroar@v0.0.0-20221229172112-1fb64f14314c/bitmap.go (about) 1 /* 2 * Copyright 2021 Dgraph Labs, Inc. and Contributors 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 17 package sroar 18 19 import ( 20 "fmt" 21 "math" 22 "sort" 23 "strings" 24 "sync" 25 26 "github.com/pkg/errors" 27 ) 28 29 var empty = make([]uint16, 16<<20) 30 31 const mask = uint64(0xFFFFFFFFFFFF0000) 32 33 type Bitmap struct { 34 data []uint16 35 keys node 36 37 // This _ptr is only used when we start with a []byte instead of a 38 // []uint16. Because we do an unsafe conversion to []uint16 data, and hence, 39 // do NOT own a valid pointer to the underlying array. 40 _ptr []byte 41 42 // memMoved keeps track of how many uint16 moves we had to do. The smaller 43 // this number, the more efficient we have been. 44 memMoved int 45 } 46 47 // FromBuffer returns a pointer to bitmap corresponding to the given buffer. This bitmap shouldn't 48 // be modified because it might corrupt the given buffer. 49 func FromBuffer(data []byte) *Bitmap { 50 assert(len(data)%2 == 0) 51 if len(data) < 8 { 52 return NewBitmap() 53 } 54 du := toUint16Slice(data) 55 x := toUint64Slice(du[:4])[indexNodeSize] 56 return &Bitmap{ 57 data: du, 58 _ptr: data, // Keep a hold of data, otherwise GC would do its thing. 59 keys: toUint64Slice(du[:x]), 60 } 61 } 62 63 // FromBufferWithCopy creates a copy of the given buffer and returns a bitmap based on the copied 64 // buffer. This bitmap is safe for both read and write operations. 65 func FromBufferWithCopy(src []byte) *Bitmap { 66 assert(len(src)%2 == 0) 67 if len(src) < 8 { 68 return NewBitmap() 69 } 70 src16 := toUint16Slice(src) 71 dst16 := make([]uint16, len(src16)) 72 copy(dst16, src16) 73 x := toUint64Slice(dst16[:4])[indexNodeSize] 74 75 return &Bitmap{ 76 data: dst16, 77 keys: toUint64Slice(dst16[:x]), 78 } 79 } 80 81 func (ra *Bitmap) ToBuffer() []byte { 82 if ra.IsEmpty() { 83 return nil 84 } 85 return toByteSlice(ra.data) 86 } 87 88 func (ra *Bitmap) ToBufferWithCopy() []byte { 89 if ra.IsEmpty() { 90 return nil 91 } 92 buf := make([]uint16, len(ra.data)) 93 copy(buf, ra.data) 94 return toByteSlice(buf) 95 } 96 97 func NewBitmap() *Bitmap { 98 return NewBitmapWith(2) 99 } 100 101 func NewBitmapWith(numKeys int) *Bitmap { 102 if numKeys < 2 { 103 panic("Must contain at least two keys.") 104 } 105 ra := &Bitmap{ 106 // Each key must also keep an offset. So, we need to double the number 107 // of uint64s allocated. Plus, we need to make space for the first 2 108 // uint64s to store the number of keys and node size. 109 data: make([]uint16, 4*(2*numKeys+2)), 110 } 111 ra.keys = toUint64Slice(ra.data) 112 ra.keys.setNodeSize(len(ra.data)) 113 114 // Always generate a container for key = 0x00. Otherwise, node gets confused 115 // about whether a zero key is a new key or not. 116 offset := ra.newContainer(minContainerSize) 117 // First two are for num keys. index=2 -> 0 key. index=3 -> offset. 118 ra.keys.setAt(indexNodeStart+1, offset) 119 ra.keys.setNumKeys(1) 120 121 return ra 122 } 123 124 func (ra *Bitmap) initSpaceForKeys(N int) { 125 if N == 0 { 126 return 127 } 128 curSize := uint64(len(ra.keys) * 4) // U64 -> U16 129 bySize := uint64(N * 8) // 2xU64 (key, value) -> 2x4xU16 130 131 // The following code is borrowed from setKey. 132 ra.scootRight(curSize, bySize) 133 ra.keys = toUint64Slice(ra.data[:curSize+bySize]) 134 ra.keys.setNodeSize(int(curSize + bySize)) 135 assert(1 == ra.keys.numKeys()) // This initialization assumes that the number of keys are 1. 136 137 // The containers have moved to the right bySize. So, update their offsets. 138 // Currently, there's only one container. 139 val := ra.keys.val(0) 140 ra.keys.setAt(valOffset(0), val+uint64(bySize)) 141 } 142 143 // setKey sets a key and container offset. 144 func (ra *Bitmap) setKey(k uint64, offset uint64) uint64 { 145 if added := ra.keys.set(k, offset); !added { 146 // No new key was added. So, we can just return. 147 return offset 148 } 149 // A new key was added. Let's ensure that ra.keys is not full. 150 if !ra.keys.isFull() { 151 return offset 152 } 153 154 // ra.keys is full. We should expand its size. 155 curSize := uint64(len(ra.keys) * 4) // Multiply by 4 for U64 -> U16. 156 bySize := curSize 157 if bySize > math.MaxUint16 { 158 bySize = math.MaxUint16 159 } 160 161 ra.scootRight(curSize, bySize) 162 ra.keys = toUint64Slice(ra.data[:curSize+bySize]) 163 ra.keys.setNodeSize(int(curSize + bySize)) 164 165 // All containers have moved to the right by bySize bytes. 166 // Update their offsets. 167 n := ra.keys 168 for i := 0; i < n.maxKeys(); i++ { 169 val := n.val(i) 170 if val > 0 { 171 n.setAt(valOffset(i), val+uint64(bySize)) 172 } 173 } 174 return offset + bySize 175 } 176 177 func (ra *Bitmap) fastExpand(bySize uint64) { 178 prev := len(ra.keys) * 4 // Multiply by 4 to convert from u16 to u64. 179 180 // This following statement also works. But, given how much fastExpand gets 181 // called (a lot), probably better to control allocation. 182 // ra.data = append(ra.data, empty[:bySize]...) 183 184 toSize := len(ra.data) + int(bySize) 185 if toSize <= cap(ra.data) { 186 ra.data = ra.data[:toSize] 187 return 188 } 189 growBy := cap(ra.data) 190 if growBy < int(bySize) { 191 growBy = int(bySize) 192 } 193 out := make([]uint16, cap(ra.data)+growBy) 194 copy(out, ra.data) 195 ra.data = out[:toSize] 196 ra._ptr = nil // Allow Go to GC whatever this was pointing to. 197 // Re-reference ra.keys correctly because underlying array has changed. 198 ra.keys = toUint64Slice(ra.data[:prev]) 199 } 200 201 // scootRight isn't aware of containers. It's going to create empty space of 202 // bySize at the given offset in ra.data. The offset doesn't need to line up 203 // with a container. 204 func (ra *Bitmap) scootRight(offset uint64, bySize uint64) { 205 left := ra.data[offset:] 206 207 ra.fastExpand(bySize) // Expand the buffer. 208 right := ra.data[len(ra.data)-len(left):] 209 n := copy(right, left) // Move data right. 210 ra.memMoved += n 211 212 Memclr(ra.data[offset : offset+uint64(bySize)]) // Zero out the space in the middle. 213 } 214 215 // scootLeft removes size number of uint16s starting from the given offset. 216 func (ra *Bitmap) scootLeft(offset uint64, size uint64) { 217 n := uint64(len(ra.data)) 218 right := ra.data[offset+size:] 219 ra.memMoved += copy(ra.data[offset:], right) 220 ra.data = ra.data[:n-size] 221 } 222 223 func (ra *Bitmap) newContainer(sz uint16) uint64 { 224 offset := uint64(len(ra.data)) 225 ra.fastExpand(uint64(sz)) 226 Memclr(ra.data[offset : offset+uint64(sz)]) 227 ra.data[offset] = sz 228 return offset 229 } 230 231 // expandContainer would expand a container at the given offset. It would typically double the size 232 // of the container, until it reaches a threshold, where the size of the container would reach 2^16. 233 // Expressed in uint16s, that'd be (2^16)/(2^4) = 2^12 = 4096. So, if the container size >= 2048, 234 // then doubling that would put it above 4096. That's why in the code below, you see the checks for 235 // size 2048. 236 func (ra *Bitmap) expandContainer(offset uint64) { 237 sz := ra.data[offset] 238 if sz == 0 { 239 panic("Container size should NOT be zero") 240 } 241 bySize := uint16(sz) 242 if sz >= 2048 { 243 // Size is in uint16. Half of max allowed size. If we're expanding the container by more 244 // than 2048, we should just cap it to max size of 4096. 245 assert(sz < maxContainerSize) 246 bySize = maxContainerSize - sz 247 } 248 249 // Select the portion to the right of the container, beyond its right boundary. 250 ra.scootRight(offset+uint64(sz), uint64(bySize)) 251 ra.keys.updateOffsets(offset, uint64(bySize), true) 252 253 if sz < 2048 { 254 ra.data[offset] = sz + bySize 255 256 } else { 257 // Convert to bitmap container. 258 src := array(ra.getContainer(offset)) 259 buf := src.toBitmapContainer(nil) 260 assert(copy(ra.data[offset:], buf) == maxContainerSize) 261 } 262 } 263 264 // stepSize is used for container expansion. For a container of given size n, 265 // stepSize would return the target size. This function is used to reduce the 266 // number of times expansion needs to happen for each container. 267 func stepSize(n uint16) uint16 { 268 // <=64 -> 128 269 // <=128 -> 256 270 // <=256 -> 512 271 // <=512 -> 1024 272 // <=1024 -> 2048 273 // >1024 -> maxSize (convert to bitmap) 274 for i := uint16(64); i <= 1024; i *= 2 { 275 if n <= i { 276 return i * 2 277 } 278 } 279 return maxContainerSize 280 } 281 282 // copyAt would copy over a given container via src, into the container at 283 // offset. If src is a bitmap, it would copy it over directly. If src is an 284 // array container, then it would follow these paths: 285 // - If src is smaller than dst, copy it over. 286 // - If not, look for target size for dst using the stepSize function. 287 // - If target size is maxSize, then convert src to a bitmap container, and 288 // copy to dst. 289 // - If target size is not max size, then expand dst container and copy src. 290 func (ra *Bitmap) copyAt(offset uint64, src []uint16) { 291 dstSize := ra.data[offset] 292 if dstSize == 0 { 293 panic("Container size should NOT be zero") 294 } 295 296 // The src is a bitmapContainer. Just copy it over. 297 if src[indexType] == typeBitmap { 298 assert(src[indexSize] == maxContainerSize) 299 bySize := uint16(maxContainerSize) - dstSize 300 // Select the portion to the right of the container, beyond its right boundary. 301 ra.scootRight(offset+uint64(dstSize), uint64(bySize)) 302 ra.keys.updateOffsets(offset, uint64(bySize), true) 303 assert(copy(ra.data[offset:], src) == len(src)) 304 return 305 } 306 307 // src is an array container. Check if dstSize >= src. If so, just copy. 308 // But, do keep dstSize intact, otherwise we'd lose portion of our container. 309 if dstSize >= src[indexSize] { 310 assert(copy(ra.data[offset:], src) == len(src)) 311 ra.data[offset] = dstSize 312 return 313 } 314 315 // dstSize < src. Determine the target size of the container. 316 targetSz := stepSize(dstSize) 317 for targetSz < src[indexSize] { 318 targetSz = stepSize(targetSz) 319 } 320 321 if targetSz == maxContainerSize { 322 // Looks like the targetSize is now maxSize. So, convert src to bitmap container. 323 s := array(src) 324 325 bySize := uint16(maxContainerSize) - dstSize 326 // Select the portion to the right of the container, beyond its right boundary. 327 ra.scootRight(offset+uint64(dstSize), uint64(bySize)) 328 ra.keys.updateOffsets(offset, uint64(bySize), true) 329 330 // Update the space of the container, so getContainer would work correctly. 331 ra.data[offset] = maxContainerSize 332 333 // Convert the src array to bitmap and write it directly over to the container. 334 out := ra.getContainer(offset) 335 Memclr(out) 336 s.toBitmapContainer(out) 337 return 338 } 339 340 // targetSize is not maxSize. Let's expand to targetSize and copy array. 341 bySize := targetSz - dstSize 342 ra.scootRight(offset+uint64(dstSize), uint64(bySize)) 343 ra.keys.updateOffsets(offset, uint64(bySize), true) 344 assert(copy(ra.data[offset:], src) == len(src)) 345 ra.data[offset] = targetSz 346 } 347 348 func (ra Bitmap) getContainer(offset uint64) []uint16 { 349 data := ra.data[offset:] 350 if len(data) == 0 { 351 panic(fmt.Sprintf("No container found at offset: %d\n", offset)) 352 } 353 sz := data[0] 354 return data[:sz] 355 } 356 357 func (ra *Bitmap) Clone() *Bitmap { 358 abuf := ra.ToBuffer() 359 bbuf := make([]byte, len(abuf)) 360 copy(bbuf, abuf) 361 return FromBuffer(bbuf) 362 } 363 364 func (ra *Bitmap) IsEmpty() bool { 365 if ra == nil { 366 return true 367 } 368 N := ra.keys.numKeys() 369 for i := 0; i < N; i++ { 370 offset := ra.keys.val(i) 371 cont := ra.getContainer(offset) 372 if c := getCardinality(cont); c > 0 { 373 return false 374 } 375 } 376 return true 377 } 378 379 func (ra *Bitmap) Set(x uint64) bool { 380 key := x & mask 381 offset, has := ra.keys.getValue(key) 382 if !has { 383 // We need to add a container. 384 o := ra.newContainer(minContainerSize) 385 // offset might have been updated by setKey. 386 offset = ra.setKey(key, o) 387 } 388 389 c := ra.getContainer(offset) 390 if c[indexType] == typeArray { 391 p := array(c) 392 if p.isFull() { 393 ra.expandContainer(offset) 394 // offsets might have changed. Safer to re-run Set. 395 return ra.Set(x) 396 } 397 } 398 399 switch c[indexType] { 400 case typeArray: 401 p := array(c) 402 return p.add(uint16(x)) 403 case typeBitmap: 404 b := bitmap(c) 405 return b.add(uint16(x)) 406 } 407 panic("we shouldn't reach here") 408 } 409 410 func FromSortedList(vals []uint64) *Bitmap { 411 var arr []uint16 412 var hi, lastHi, off uint64 413 414 ra := NewBitmap() 415 416 if len(vals) == 0 { 417 return ra 418 } 419 420 // Set the keys beforehand so that we don't need to move a lot of memory because of adding keys. 421 var numKeys int 422 for _, x := range vals { 423 hi = x & mask 424 if hi != 0 && hi != lastHi { 425 numKeys++ 426 } 427 lastHi = hi 428 } 429 ra.initSpaceForKeys(numKeys) 430 431 finalize := func(l []uint16, key uint64) { 432 if len(l) == 0 { 433 return 434 } 435 if len(l) <= 2048 { 436 // 4 uint16s for the header, and extra 4 uint16s so that adding more elements using 437 // Set operation doesn't fail. 438 sz := uint16(8 + len(l)) 439 off = ra.newContainer(sz) 440 c := ra.getContainer(off) 441 c[indexSize] = sz 442 c[indexType] = typeArray 443 setCardinality(c, len(l)) 444 for i := 0; i < len(l); i++ { 445 c[int(startIdx)+i] = l[i] 446 } 447 448 } else { 449 off = ra.newContainer(maxContainerSize) 450 c := ra.getContainer(off) 451 c[indexSize] = maxContainerSize 452 c[indexType] = typeBitmap 453 for _, v := range l { 454 bitmap(c).add(v) 455 } 456 } 457 ra.setKey(key, off) 458 return 459 } 460 461 lastHi = 0 462 for _, x := range vals { 463 hi = x & mask 464 // Finalize the last container before proceeding ahead 465 if hi != 0 && hi != lastHi { 466 finalize(arr, lastHi) 467 arr = arr[:0] 468 } 469 arr = append(arr, uint16(x)) 470 lastHi = hi 471 } 472 finalize(arr, lastHi) 473 return ra 474 } 475 476 // TODO: Potentially this can be optimized. 477 func (ra *Bitmap) SetMany(vals []uint64) { 478 for _, k := range vals { 479 ra.Set(k) 480 } 481 } 482 483 // Select returns the element at the xth index. (0-indexed) 484 func (ra *Bitmap) Select(x uint64) (uint64, error) { 485 if x >= uint64(ra.GetCardinality()) { 486 return 0, errors.Errorf("index %d is not less than the cardinality: %d", 487 x, ra.GetCardinality()) 488 } 489 n := ra.keys.numKeys() 490 for i := 0; i < n; i++ { 491 off := ra.keys.val(i) 492 con := ra.getContainer(off) 493 c := uint64(getCardinality(con)) 494 assert(c != uint64(invalidCardinality)) 495 if x < c { 496 key := ra.keys.key(i) 497 switch con[indexType] { 498 case typeArray: 499 return key | uint64(array(con).all()[x]), nil 500 case typeBitmap: 501 return key | uint64(bitmap(con).selectAt(int(x))), nil 502 } 503 } 504 x -= c 505 } 506 panic("should not reach here") 507 } 508 509 func (ra *Bitmap) Contains(x uint64) bool { 510 if ra == nil { 511 return false 512 } 513 key := x & mask 514 offset, has := ra.keys.getValue(key) 515 if !has { 516 return false 517 } 518 y := uint16(x) 519 520 c := ra.getContainer(offset) 521 switch c[indexType] { 522 case typeArray: 523 p := array(c) 524 return p.has(y) 525 case typeBitmap: 526 b := bitmap(c) 527 return b.has(y) 528 } 529 return false 530 } 531 532 func (ra *Bitmap) Remove(x uint64) bool { 533 if ra == nil { 534 return false 535 } 536 key := x & mask 537 offset, has := ra.keys.getValue(key) 538 if !has { 539 return false 540 } 541 c := ra.getContainer(offset) 542 switch c[indexType] { 543 case typeArray: 544 p := array(c) 545 return p.remove(uint16(x)) 546 case typeBitmap: 547 b := bitmap(c) 548 return b.remove(uint16(x)) 549 } 550 return true 551 } 552 553 // Remove range removes [lo, hi) from the bitmap. 554 func (ra *Bitmap) RemoveRange(lo, hi uint64) { 555 if lo > hi { 556 panic("lo should not be more than hi") 557 } 558 if lo == hi { 559 return 560 } 561 562 k1 := lo & mask 563 k2 := hi & mask 564 565 defer ra.Cleanup() 566 567 // Complete range lie in a single container 568 if k1 == k2 { 569 if off, has := ra.keys.getValue(k1); has { 570 c := ra.getContainer(off) 571 removeRangeContainer(c, uint16(lo), uint16(hi)-1) 572 } 573 return 574 } 575 576 // Remove all the containers in range [k1+1, k2-1]. 577 n := ra.keys.numKeys() 578 st := ra.keys.search(k1) 579 key := ra.keys.key(st) 580 if key == k1 { 581 st++ 582 } 583 584 for i := st; i < n; i++ { 585 key := ra.keys.key(i) 586 if key >= k2 { 587 break 588 } 589 if off, has := ra.keys.getValue(key); has { 590 zeroOutContainer(ra.getContainer(off)) 591 } 592 } 593 594 // Remove elements >= lo in k1's container 595 if off, has := ra.keys.getValue(k1); has { 596 c := ra.getContainer(off) 597 if uint16(lo) == 0 { 598 zeroOutContainer(c) 599 } else { 600 removeRangeContainer(c, uint16(lo), math.MaxUint16) 601 } 602 } 603 604 if uint16(hi) == 0 { 605 return 606 } 607 608 // Remove all elements < hi in k2's container 609 if off, has := ra.keys.getValue(k2); has { 610 c := ra.getContainer(off) 611 removeRangeContainer(c, 0, uint16(hi)-1) 612 } 613 } 614 615 func (ra *Bitmap) Reset() { 616 // reset ra.data to size enough for one container and corresponding key. 617 // 2 u64 is needed for header and another 2 u16 for the key 0. 618 ra.data = ra.data[:16+minContainerSize] 619 ra.keys = toUint64Slice(ra.data) 620 621 offset := ra.newContainer(minContainerSize) 622 ra.keys.setAt(indexNodeStart+1, offset) 623 ra.keys.setNumKeys(1) 624 } 625 626 func (ra *Bitmap) GetCardinality() int { 627 if ra == nil { 628 return 0 629 } 630 N := ra.keys.numKeys() 631 var sz int 632 for i := 0; i < N; i++ { 633 offset := ra.keys.val(i) 634 c := ra.getContainer(offset) 635 sz += getCardinality(c) 636 } 637 return sz 638 } 639 640 func (ra *Bitmap) ToArray() []uint64 { 641 if ra == nil { 642 return nil 643 } 644 res := make([]uint64, 0, ra.GetCardinality()) 645 N := ra.keys.numKeys() 646 for i := 0; i < N; i++ { 647 key := ra.keys.key(i) 648 off := ra.keys.val(i) 649 c := ra.getContainer(off) 650 651 switch c[indexType] { 652 case typeArray: 653 a := array(c) 654 for _, lo := range a.all() { 655 res = append(res, key|uint64(lo)) 656 } 657 case typeBitmap: 658 b := bitmap(c) 659 out := b.all() 660 for _, x := range out { 661 res = append(res, key|uint64(x)) 662 } 663 } 664 } 665 return res 666 } 667 668 func (ra *Bitmap) String() string { 669 var b strings.Builder 670 b.WriteRune('\n') 671 672 var usedSize, card int 673 usedSize += 4 * (ra.keys.numKeys()) 674 for i := 0; i < ra.keys.numKeys(); i++ { 675 k := ra.keys.key(i) 676 v := ra.keys.val(i) 677 c := ra.getContainer(v) 678 679 sz := c[indexSize] 680 usedSize += int(sz) 681 card += getCardinality(c) 682 683 b.WriteString(fmt.Sprintf( 684 "[%03d] Key: %#8x. Offset: %7d. Size: %4d. Type: %d. Card: %6d. Uint16/Uid: %.2f\n", 685 i, k, v, sz, c[indexType], getCardinality(c), float64(sz)/float64(getCardinality(c)))) 686 } 687 b.WriteString(fmt.Sprintf("Number of containers: %d. Cardinality: %d\n", 688 ra.keys.numKeys(), card)) 689 690 amp := float64(len(ra.data)-usedSize) / float64(usedSize) 691 b.WriteString(fmt.Sprintf( 692 "Size in Uint16s. Used: %d. Total: %d. Space Amplification: %.2f%%. Moved: %.2fx\n", 693 usedSize, len(ra.data), amp*100.0, float64(ra.memMoved)/float64(usedSize))) 694 695 b.WriteString(fmt.Sprintf("Used Uint16/Uid: %.2f. Total Uint16/Uid: %.2f", 696 float64(usedSize)/float64(card), float64(len(ra.data))/float64(card))) 697 698 return b.String() 699 } 700 701 const fwd int = 0x01 702 const rev int = 0x02 703 704 func (ra *Bitmap) Minimum() uint64 { return ra.extreme(fwd) } 705 func (ra *Bitmap) Maximum() uint64 { return ra.extreme(rev) } 706 707 func (ra *Bitmap) Debug(x uint64) string { 708 var b strings.Builder 709 hi := x & mask 710 off, found := ra.keys.getValue(hi) 711 if !found { 712 b.WriteString(fmt.Sprintf("Unable to find the container for x: %#x\n", hi)) 713 b.WriteString(ra.String()) 714 } 715 c := ra.getContainer(off) 716 lo := uint16(x) 717 718 b.WriteString(fmt.Sprintf("x: %#x lo: %#x. offset: %d\n", x, lo, off)) 719 720 switch c[indexType] { 721 case typeArray: 722 case typeBitmap: 723 idx := lo / 16 724 pos := lo % 16 725 b.WriteString(fmt.Sprintf("At idx: %d. Pos: %d val: %#b\n", idx, pos, c[startIdx+idx])) 726 } 727 return b.String() 728 } 729 730 func (ra *Bitmap) extreme(dir int) uint64 { 731 N := ra.keys.numKeys() 732 if N == 0 { 733 return 0 734 } 735 736 var k uint64 737 var c []uint16 738 739 if dir == fwd { 740 for i := 0; i < N; i++ { 741 offset := ra.keys.val(i) 742 c = ra.getContainer(offset) 743 if getCardinality(c) > 0 { 744 k = ra.keys.key(i) 745 break 746 } 747 } 748 } else { 749 for i := N - 1; i >= 0; i-- { 750 offset := ra.keys.val(i) 751 c = ra.getContainer(offset) 752 if getCardinality(c) > 0 { 753 k = ra.keys.key(i) 754 break 755 } 756 } 757 } 758 759 switch c[indexType] { 760 case typeArray: 761 a := array(c) 762 if dir == fwd { 763 return k | uint64(a.minimum()) 764 } 765 return k | uint64(a.maximum()) 766 case typeBitmap: 767 b := bitmap(c) 768 if dir == fwd { 769 return k | uint64(b.minimum()) 770 } 771 return k | uint64(b.maximum()) 772 default: 773 panic("We don't support this type of container") 774 } 775 } 776 777 func (ra *Bitmap) And(bm *Bitmap) { 778 if bm == nil { 779 ra.Reset() 780 return 781 } 782 783 a, b := ra, bm 784 ai, an := 0, a.keys.numKeys() 785 bi, bn := 0, b.keys.numKeys() 786 787 for ai < an && bi < bn { 788 ak := a.keys.key(ai) 789 bk := b.keys.key(bi) 790 if ak == bk { 791 off := a.keys.val(ai) 792 ac := a.getContainer(off) 793 794 off = b.keys.val(bi) 795 bc := b.getContainer(off) 796 797 // do the intersection 798 // TODO: See if we can do containerAnd operation in-place. 799 c := containerAnd(ac, bc) 800 801 // create a new container and update the key offset to this container. 802 offset := a.newContainer(uint16(len(c))) 803 copy(a.data[offset:], c) 804 a.setKey(ak, offset) 805 ai++ 806 bi++ 807 } else if ak < bk { 808 off := a.keys.val(ai) 809 zeroOutContainer(a.getContainer(off)) 810 ai++ 811 } else { 812 bi++ 813 } 814 } 815 for ai < an { 816 off := a.keys.val(ai) 817 zeroOutContainer(a.getContainer(off)) 818 ai++ 819 } 820 } 821 822 func And(a, b *Bitmap) *Bitmap { 823 ai, an := 0, a.keys.numKeys() 824 bi, bn := 0, b.keys.numKeys() 825 826 res := NewBitmap() 827 for ai < an && bi < bn { 828 ak := a.keys.key(ai) 829 bk := a.keys.key(bi) 830 if ak == bk { 831 // Do the intersection. 832 off := a.keys.val(ai) 833 ac := a.getContainer(off) 834 835 off = b.keys.val(bi) 836 bc := b.getContainer(off) 837 838 outc := containerAnd(ac, bc) 839 if getCardinality(outc) > 0 { 840 offset := res.newContainer(uint16(len(outc))) 841 copy(res.data[offset:], outc) 842 res.setKey(ak, offset) 843 } 844 ai++ 845 bi++ 846 } else if ak < bk { 847 ai++ 848 } else { 849 bi++ 850 } 851 } 852 return res 853 } 854 855 func (ra *Bitmap) AndNot(bm *Bitmap) { 856 if bm == nil { 857 return 858 } 859 a, b := ra, bm 860 var ai, bi int 861 862 buf := make([]uint16, maxContainerSize) 863 for ai < a.keys.numKeys() && bi < b.keys.numKeys() { 864 ak := a.keys.key(ai) 865 bk := b.keys.key(bi) 866 if ak == bk { 867 off := a.keys.val(ai) 868 ac := a.getContainer(off) 869 870 off = b.keys.val(bi) 871 bc := b.getContainer(off) 872 873 // TODO: See if we can do containerAndNot operation in-place. 874 c := containerAndNot(ac, bc, buf) 875 // create a new container and update the key offset to this container. 876 offset := a.newContainer(uint16(len(c))) 877 copy(a.data[offset:], c) 878 a.setKey(ak, offset) 879 880 ai++ 881 bi++ 882 continue 883 } 884 if ak < bk { 885 ai++ 886 } else { 887 bi++ 888 } 889 } 890 } 891 892 // TODO: Check if we want to use lazyMode 893 func (dst *Bitmap) Or(src *Bitmap) { 894 if src == nil { 895 return 896 } 897 dst.or(src, runInline) 898 } 899 900 func (dst *Bitmap) or(src *Bitmap, runMode int) { 901 srcIdx, numKeys := 0, src.keys.numKeys() 902 903 buf := make([]uint16, maxContainerSize) 904 for ; srcIdx < numKeys; srcIdx++ { 905 srcCont := src.getContainer(src.keys.val(srcIdx)) 906 if getCardinality(srcCont) == 0 { 907 continue 908 } 909 910 key := src.keys.key(srcIdx) 911 912 dstIdx := dst.keys.search(key) 913 if dstIdx >= dst.keys.numKeys() || dst.keys.key(dstIdx) != key { 914 // srcCont doesn't exist in dst. So, copy it over. 915 offset := dst.newContainer(uint16(len(srcCont))) 916 copy(dst.getContainer(offset), srcCont) 917 dst.setKey(key, offset) 918 } else { 919 // Container exists in dst as well. Do an inline containerOr. 920 offset := dst.keys.val(dstIdx) 921 dstCont := dst.getContainer(offset) 922 if c := containerOr(dstCont, srcCont, buf, runMode|runInline); len(c) > 0 { 923 dst.copyAt(offset, c) 924 dst.setKey(key, offset) 925 } 926 } 927 } 928 } 929 930 func Or(a, b *Bitmap) *Bitmap { 931 ai, an := 0, a.keys.numKeys() 932 bi, bn := 0, b.keys.numKeys() 933 934 buf := make([]uint16, maxContainerSize) 935 res := NewBitmap() 936 for ai < an && bi < bn { 937 ak := a.keys.key(ai) 938 ac := a.getContainer(a.keys.val(ai)) 939 940 bk := b.keys.key(bi) 941 bc := b.getContainer(b.keys.val(bi)) 942 943 if ak == bk { 944 // Do the union. 945 outc := containerOr(ac, bc, buf, 0) 946 offset := res.newContainer(uint16(len(outc))) 947 copy(res.data[offset:], outc) 948 res.setKey(ak, offset) 949 ai++ 950 bi++ 951 } else if ak < bk { 952 off := res.newContainer(uint16(len(ac))) 953 copy(res.getContainer(off), ac) 954 res.setKey(ak, off) 955 ai++ 956 } else { 957 off := res.newContainer(uint16(len(bc))) 958 copy(res.getContainer(off), bc) 959 res.setKey(bk, off) 960 bi++ 961 } 962 } 963 for ai < an { 964 ak := a.keys.key(ai) 965 ac := a.getContainer(a.keys.val(ai)) 966 off := res.newContainer(uint16(len(ac))) 967 968 copy(res.getContainer(off), ac) 969 res.setKey(ak, off) 970 ai++ 971 } 972 for bi < bn { 973 bk := b.keys.key(bi) 974 bc := b.getContainer(b.keys.val(bi)) 975 off := res.newContainer(uint16(len(bc))) 976 977 copy(res.getContainer(off), bc) 978 res.setKey(bk, off) 979 bi++ 980 } 981 return res 982 } 983 984 func (ra *Bitmap) Rank(x uint64) int { 985 key := x & mask 986 offset, has := ra.keys.getValue(key) 987 if !has { 988 return -1 989 } 990 c := ra.getContainer(offset) 991 y := uint16(x) 992 993 // Find the rank within the container 994 var rank int 995 switch c[indexType] { 996 case typeArray: 997 rank = array(c).rank(y) 998 case typeBitmap: 999 rank = bitmap(c).rank(y) 1000 } 1001 if rank < 0 { 1002 return -1 1003 } 1004 1005 // Add up cardinalities of all the containers on the left of container containing x. 1006 n := ra.keys.numKeys() 1007 for i := 0; i < n; i++ { 1008 if ra.keys.key(i) == key { 1009 break 1010 } 1011 cont := ra.getContainer(ra.keys.val(i)) 1012 rank += getCardinality(cont) 1013 } 1014 return rank 1015 } 1016 1017 func (ra *Bitmap) Cleanup() { 1018 type interval struct { 1019 start uint64 1020 end uint64 1021 } 1022 1023 // Find the ranges that needs to be removed in the key space and the container space. Also, 1024 // start the iteration from idx = 1 because we never remove the 0 key. 1025 var keyIntervals, contIntervals []interval 1026 for idx := 1; idx < ra.keys.numKeys(); idx++ { 1027 off := ra.keys.val(idx) 1028 cont := ra.getContainer(off) 1029 if getCardinality(cont) == 0 { 1030 ko := uint64(keyOffset(idx)) 1031 contIntervals = append(contIntervals, interval{off, off + uint64(cont[indexSize])}) 1032 keyIntervals = append(keyIntervals, interval{4 * ko, 4 * (ko + 2)}) 1033 } 1034 } 1035 if len(contIntervals) == 0 { 1036 return 1037 } 1038 1039 merge := func(intervals []interval) []interval { 1040 assert(len(intervals) > 0) 1041 1042 // Merge the ranges in order to reduce scootLeft 1043 merged := []interval{intervals[0]} 1044 for _, ir := range intervals[1:] { 1045 last := merged[len(merged)-1] 1046 if ir.start == last.end { 1047 last.end = ir.end 1048 merged[len(merged)-1] = last 1049 continue 1050 } 1051 merged = append(merged, ir) 1052 } 1053 return merged 1054 } 1055 1056 // Key intervals are already sorted, but container intervals needs to be sorted because 1057 // they are always added in the end of the ra.data. 1058 sort.Slice(contIntervals, func(i, j int) bool { 1059 return contIntervals[i].start < contIntervals[j].start 1060 }) 1061 1062 contIntervals = merge(contIntervals) 1063 keyIntervals = merge(keyIntervals) 1064 1065 // Cleanup the containers. 1066 moved := uint64(0) 1067 for _, ir := range contIntervals { 1068 assert(ir.start >= moved) 1069 sz := ir.end - ir.start 1070 ra.scootLeft(ir.start-moved, sz) 1071 ra.keys.updateOffsets(ir.end-moved-1, sz, false) 1072 moved += sz 1073 } 1074 1075 // Cleanup the key space. 1076 moved = uint64(0) 1077 for _, ir := range keyIntervals { 1078 assert(ir.start >= moved) 1079 sz := ir.end - ir.start 1080 ra.scootLeft(ir.start-moved, sz) 1081 1082 // sz is in number of u16s, hence number of key-value removed is sz/8. 1083 ra.keys.setNumKeys(ra.keys.numKeys() - int(sz/8)) 1084 ra.keys.setNodeSize(ra.keys.size() - int(sz)) 1085 ra.keys = ra.keys[:len(ra.keys)-int(sz/4)] 1086 ra.keys.updateOffsets(ir.end-moved-1, sz, false) 1087 moved += sz 1088 } 1089 } 1090 1091 func FastAnd(bitmaps ...*Bitmap) *Bitmap { 1092 if len(bitmaps) == 0 { 1093 return NewBitmap() 1094 } 1095 b := bitmaps[0] 1096 for _, bm := range bitmaps[1:] { 1097 b.And(bm) 1098 } 1099 b.Cleanup() 1100 return b 1101 } 1102 1103 // FastParOr would group up bitmaps and call FastOr on them concurrently. It 1104 // would then merge the groups into final Bitmap. This approach is simpler and 1105 // faster than operating at a container level, because we can't operate on array 1106 // containers belonging to the same Bitmap concurrently because array containers 1107 // can expand, leaving no clear boundaries. 1108 // 1109 // If FastParOr is called with numGo=1, it just calls FastOr. 1110 // 1111 // Experiments with numGo=4 shows that FastParOr would be 2x the speed of 1112 // FastOr, but 4x the memory usage, even under 50% CPU usage. So, use wisely. 1113 func FastParOr(numGo int, bitmaps ...*Bitmap) *Bitmap { 1114 if numGo == 1 { 1115 return FastOr(bitmaps...) 1116 } 1117 width := max(len(bitmaps)/numGo, 3) 1118 1119 var wg sync.WaitGroup 1120 var res []*Bitmap 1121 for start := 0; start < len(bitmaps); start += width { 1122 end := min(start+width, len(bitmaps)) 1123 res = append(res, nil) // Make space for result. 1124 wg.Add(1) 1125 1126 go func(start, end int) { 1127 idx := start / width 1128 res[idx] = FastOr(bitmaps[start:end]...) 1129 wg.Done() 1130 }(start, end) 1131 } 1132 wg.Wait() 1133 return FastOr(res...) 1134 } 1135 1136 // FastOr would merge given Bitmaps into one Bitmap. This is faster than 1137 // doing an OR over the bitmaps iteratively. 1138 func FastOr(bitmaps ...*Bitmap) *Bitmap { 1139 if len(bitmaps) == 0 { 1140 return NewBitmap() 1141 } 1142 if len(bitmaps) == 1 { 1143 return bitmaps[0] 1144 } 1145 1146 // We first figure out the container distribution across the bitmaps. We do 1147 // that by looking at the key of the container, and the cardinality. We 1148 // assume the worst-case scenario where the union would result in a 1149 // cardinality (per container) of the sum of cardinalities of each of the 1150 // corresponding containers in other bitmaps. 1151 containers := make(map[uint64]int) 1152 for _, b := range bitmaps { 1153 for i := 0; i < b.keys.numKeys(); i++ { 1154 offset := b.keys.val(i) 1155 cont := b.getContainer(offset) 1156 card := getCardinality(cont) 1157 containers[b.keys.key(i)] += card 1158 } 1159 } 1160 1161 // We use the above information to pre-generate the destination Bitmap and 1162 // allocate container sizes based on the calculated cardinalities. 1163 // var sz int 1164 dst := NewBitmap() 1165 // First create the keys. We do this as a separate step, because keys are 1166 // the left most portion of the data array. Adding space there requires 1167 // moving a lot of pieces. 1168 for key, card := range containers { 1169 if card > 0 { 1170 dst.setKey(key, 0) 1171 } 1172 } 1173 1174 // Then create the bitmap containers. 1175 for key, card := range containers { 1176 if card >= 4096 { 1177 offset := dst.newContainer(maxContainerSize) 1178 c := dst.getContainer(offset) 1179 c[indexSize] = maxContainerSize 1180 c[indexType] = typeBitmap 1181 dst.setKey(key, offset) 1182 } 1183 } 1184 1185 // Create the array containers at the end. This allows them to expand 1186 // without having to move a lot of memory. 1187 for key, card := range containers { 1188 // Ensure this condition exactly maps up with above. 1189 if card < 4096 && card > 0 { 1190 if card < minContainerSize { 1191 card = minContainerSize 1192 } 1193 offset := dst.newContainer(uint16(card)) 1194 c := dst.getContainer(offset) 1195 c[indexSize] = uint16(card) 1196 c[indexType] = typeArray 1197 dst.setKey(key, offset) 1198 } 1199 } 1200 1201 // dst Bitmap is ready to be ORed with the given Bitmaps. 1202 for _, b := range bitmaps { 1203 dst.or(b, runLazy) 1204 } 1205 1206 for i := 0; i < dst.keys.numKeys(); i++ { 1207 offset := dst.keys.val(i) 1208 c := dst.getContainer(offset) 1209 if getCardinality(c) == invalidCardinality { 1210 calculateAndSetCardinality(c) 1211 } 1212 } 1213 1214 return dst 1215 } 1216 1217 // Split splits the bitmap based on maxSz and the externalSize function. It splits the bitmap 1218 // such that size of each split bitmap + external size corresponding to its elements approximately 1219 // equal to maxSz (it can be greater than maxSz sometimes). The splits are returned in sorted order. 1220 // externalSize is a function that should return the external size corresponding to elements in 1221 // range [start, end]. External size is used to calculate the split boundaries. 1222 func (bm *Bitmap) Split(externalSize func(start, end uint64) uint64, maxSz uint64) []*Bitmap { 1223 splitFurther := func(b *Bitmap) []*Bitmap { 1224 itr := b.NewIterator() 1225 newBm := NewBitmap() 1226 var sz uint64 1227 var bms []*Bitmap 1228 for id := itr.Next(); id != 0; id = itr.Next() { 1229 sz += externalSize(id, id) 1230 newBm.Set(id) 1231 if sz >= maxSz { 1232 bms = append(bms, newBm) 1233 newBm = NewBitmap() 1234 sz = 0 1235 } 1236 } 1237 1238 if !newBm.IsEmpty() { 1239 bms = append(bms, newBm) 1240 } 1241 return bms 1242 } 1243 1244 create := func(keyToOffset map[uint64]uint64, totalSz uint64) []*Bitmap { 1245 var keys []uint64 1246 for key := range keyToOffset { 1247 keys = append(keys, key) 1248 } 1249 sort.Slice(keys, func(i, j int) bool { 1250 return keys[i] < keys[j] 1251 }) 1252 1253 newBm := NewBitmap() 1254 1255 // First set all the keys. 1256 var containerSz uint64 1257 for _, key := range keys { 1258 newBm.setKey(key, 0) 1259 1260 // Calculate the size of the containers. 1261 cont := bm.getContainer(keyToOffset[key]) 1262 containerSz += uint64(len(cont)) 1263 } 1264 // Allocate enough space to hold all the containers. 1265 beforeSize := len(newBm.data) 1266 newBm.fastExpand(containerSz) 1267 newBm.data = newBm.data[:beforeSize] 1268 1269 // Now, we can populate the containers. For that, we first expand the 1270 // bitmap. Calculate the total size we need to allocate all these containers. 1271 for _, key := range keys { 1272 cont := bm.getContainer(keyToOffset[key]) 1273 off := newBm.newContainer(uint16(len(cont))) 1274 copy(newBm.data[off:], cont) 1275 1276 newBm.setKey(key, off) 1277 } 1278 1279 if newBm.GetCardinality() == 0 { 1280 return nil 1281 } 1282 1283 if totalSz > maxSz { 1284 return splitFurther(newBm) 1285 } 1286 1287 return []*Bitmap{newBm} 1288 } 1289 1290 var splits []*Bitmap 1291 1292 containerMap := make(map[uint64]uint64) 1293 var totalSz uint64 // size of containers plus the external size of the container 1294 1295 for i := 0; i < bm.keys.numKeys(); i++ { 1296 key := bm.keys.key(i) 1297 off := bm.keys.val(i) 1298 cont := bm.getContainer(off) 1299 1300 start, end := key, addUint64(key, 1<<16-1) 1301 sz := externalSize(start, end) + 2*uint64(cont[indexSize]) // Converting to bytes. 1302 1303 // We can probably append more containers in the same bucket. 1304 if totalSz+sz < maxSz || len(containerMap) == 0 { 1305 // Include this container in the container map. 1306 containerMap[key] = off 1307 totalSz += sz 1308 continue 1309 } 1310 1311 // We have reached the maxSz limit. Hence, create a split. 1312 splits = append(splits, create(containerMap, totalSz)...) 1313 1314 containerMap = make(map[uint64]uint64) 1315 containerMap[key] = off 1316 totalSz = sz 1317 } 1318 if len(containerMap) > 0 { 1319 splits = append(splits, create(containerMap, totalSz)...) 1320 } 1321 1322 return splits 1323 }