gorgonia.org/gorgonia@v0.9.17/bitmap.go (about)

     1  package gorgonia
     2  
     3  import "math/bits"
     4  
     5  const (
     6  	bitmapBits = 64
     7  )
     8  
     9  // bitmap is a very simple bitmap. It only supports Set, IsSet and Clear methods. It's mostly used for tracking which element has been set
    10  type bitmap struct {
    11  	n   []uint64
    12  	max int
    13  }
    14  
    15  // newBitmap creates a new bitmap.
    16  func newBitmap(size int) *bitmap {
    17  	q, r := divmod(size, bitmapBits)
    18  
    19  	if r > 0 {
    20  		q++
    21  	}
    22  
    23  	return &bitmap{
    24  		n:   make([]uint64, q),
    25  		max: size,
    26  	}
    27  }
    28  
    29  // Set sets the ith bit of the bit map to 1. It panics if i is greater or equal to the defined max
    30  func (bm *bitmap) Set(i int) {
    31  	if i >= bm.max || i < 0 {
    32  		panic("Index out of range")
    33  	}
    34  
    35  	block, pos := divmod(i, bitmapBits)
    36  	bm.n[block] |= uint64(1) << uint64(pos)
    37  }
    38  
    39  // IsSet returns true if the ith bit is set. It panics if the i is greater or equal to the defined max
    40  func (bm *bitmap) IsSet(i int) bool {
    41  	if i >= bm.max || i < 0 {
    42  		panic("Index out of range")
    43  	}
    44  
    45  	block, pos := divmod(i, bitmapBits)
    46  	return bm.n[block]>>uint64(pos)&uint64(1) == uint64(1)
    47  }
    48  
    49  // Clear clears the ith bit. It panics if i is greater or equal to the defined max
    50  func (bm *bitmap) Clear(i int) {
    51  	if i >= bm.max || i < 0 {
    52  		panic("Index out of range")
    53  	}
    54  
    55  	block, pos := divmod(i, bitmapBits)
    56  	bm.n[block] &= ^(uint64(1) << uint64(pos))
    57  }
    58  
    59  // BlocksWithZero finds the first block with zeroes in the bit. atleast specifies how many consecutive zeroes need be found
    60  func (bm *bitmap) BlocksWithZero(atleast int) int {
    61  	retVal := -1
    62  	for i, b := range bm.n {
    63  		if bits.OnesCount64(b) != bitmapBits {
    64  			// shortcut:
    65  			if bits.LeadingZeros64(b) > atleast {
    66  				return i
    67  			}
    68  
    69  			var consecutive int
    70  			for j := 0; j < bitmapBits; j++ {
    71  				if b>>uint64(j)&uint64(1) == 0 {
    72  					consecutive++
    73  				} else {
    74  					consecutive = 0
    75  				}
    76  				if consecutive > atleast {
    77  					return i
    78  				}
    79  			}
    80  		}
    81  	}
    82  	return retVal
    83  }