github.com/hdt3213/godis@v1.2.9/datastruct/bitmap/bitmap.go (about)

     1  package bitmap
     2  
     3  type BitMap []byte
     4  
     5  func New() *BitMap {
     6  	b := BitMap(make([]byte, 0))
     7  	return &b
     8  }
     9  
    10  func toByteSize(bitSize int64) int64 {
    11  	if bitSize%8 == 0 {
    12  		return bitSize / 8
    13  	}
    14  	return bitSize/8 + 1
    15  }
    16  
    17  func (b *BitMap) grow(bitSize int64) {
    18  	byteSize := toByteSize(bitSize)
    19  	gap := byteSize - int64(len(*b))
    20  	if gap <= 0 {
    21  		return
    22  	}
    23  	*b = append(*b, make([]byte, gap)...)
    24  }
    25  
    26  func (b *BitMap) BitSize() int {
    27  	return len(*b) * 8
    28  }
    29  
    30  func FromBytes(bytes []byte) *BitMap {
    31  	bm := BitMap(bytes)
    32  	return &bm
    33  }
    34  
    35  func (b *BitMap) ToBytes() []byte {
    36  	return *b
    37  }
    38  
    39  func (b *BitMap) SetBit(offset int64, val byte) {
    40  	byteIndex := offset / 8
    41  	bitOffset := offset % 8
    42  	mask := byte(1 << bitOffset)
    43  	b.grow(offset + 1)
    44  	if val > 0 {
    45  		// set bit
    46  		(*b)[byteIndex] |= mask
    47  	} else {
    48  		// clear bit
    49  		(*b)[byteIndex] &^= mask
    50  	}
    51  }
    52  
    53  func (b *BitMap) GetBit(offset int64) byte {
    54  	byteIndex := offset / 8
    55  	bitOffset := offset % 8
    56  	if byteIndex >= int64(len(*b)) {
    57  		return 0
    58  	}
    59  	return ((*b)[byteIndex] >> bitOffset) & 0x01
    60  }
    61  
    62  type Callback func(offset int64, val byte) bool
    63  
    64  func (b *BitMap) ForEachBit(begin int64, end int64, cb Callback) {
    65  	offset := begin
    66  	byteIndex := offset / 8
    67  	bitOffset := offset % 8
    68  	for byteIndex < int64(len(*b)) {
    69  		b := (*b)[byteIndex]
    70  		for bitOffset < 8 {
    71  			bit := byte(b >> bitOffset & 0x01)
    72  			if !cb(offset, bit) {
    73  				return
    74  			}
    75  			bitOffset++
    76  			offset++
    77  			if offset >= end && end != 0 {
    78  				break
    79  			}
    80  		}
    81  		byteIndex++
    82  		bitOffset = 0
    83  		if end > 0 && offset >= end {
    84  			break
    85  		}
    86  	}
    87  }
    88  
    89  func (b *BitMap) ForEachByte(begin int, end int, cb Callback) {
    90  	if end == 0 {
    91  		end = len(*b)
    92  	} else if end > len(*b) {
    93  		end = len(*b)
    94  	}
    95  	for i := begin; i < end; i++ {
    96  		if !cb(int64(i), (*b)[i]) {
    97  			return
    98  		}
    99  	}
   100  }