github.com/zeebo/mon@v0.0.0-20211012163247-13d39bdb54fa/internal/bitmap/bitmap.go (about)

     1  package bitmap
     2  
     3  import (
     4  	"math/bits"
     5  	"sync/atomic"
     6  )
     7  
     8  //
     9  // 64 bits
    10  //
    11  
    12  type B64 [1]uint64
    13  
    14  func (b *B64) Clone() B64 {
    15  	return B64{atomic.LoadUint64(&b[0])}
    16  }
    17  
    18  func (b *B64) Set(idx uint) {
    19  	atomic.AddUint64(&b[0], 1<<(idx&63))
    20  }
    21  
    22  func (b *B64) Has(idx uint) bool {
    23  	return atomic.LoadUint64(&b[0])&(1<<(idx&63)) > 0
    24  }
    25  
    26  func (b *B64) Next() (idx uint, ok bool) {
    27  	u := b[0]
    28  	c := u & (u - 1)
    29  	idx = uint(bits.Len64(u ^ c))
    30  	b[0] = c
    31  	return (idx - 1) % 64, u > 0
    32  }
    33  
    34  //
    35  // 128 bits
    36  //
    37  
    38  type B128 [2]uint64
    39  
    40  func (b *B128) Clone() B128 {
    41  	return B128{atomic.LoadUint64(&b[0]), atomic.LoadUint64(&b[1])}
    42  }
    43  
    44  func (b *B128) Set(idx uint) {
    45  	atomic.AddUint64(&b[(idx>>6)&1], 1<<(idx&63))
    46  }
    47  
    48  func (b *B128) Has(idx uint) bool {
    49  	return atomic.LoadUint64(&b[(idx>>6)&1])&(1<<(idx&63)) > 0
    50  }
    51  
    52  func (b *B128) Next() (idx uint, ok bool) {
    53  	u := b[0]
    54  	c := u & (u - 1)
    55  	idx = uint(bits.Len64(u ^ c))
    56  	b[0] = c
    57  
    58  	if u > 0 {
    59  		return (idx - 1) % 128, true
    60  	}
    61  
    62  	u = b[1]
    63  	c = u & (u - 1)
    64  	idx = 63 + uint(bits.Len64(u^c))
    65  	b[1] = c
    66  
    67  	return idx % 128, u > 0
    68  }