bitlist.go raw

   1  package utils
   2  
   3  // BitList is a list that contains bits
   4  type BitList struct {
   5  	count int
   6  	data  []int32
   7  }
   8  
   9  // NewBitList returns a new BitList with the given length
  10  // all bits are initialize with false
  11  func NewBitList(capacity int) *BitList {
  12  	bl := new(BitList)
  13  	bl.count = capacity
  14  	x := 0
  15  	if capacity%32 != 0 {
  16  		x = 1
  17  	}
  18  	bl.data = make([]int32, capacity/32+x)
  19  	return bl
  20  }
  21  
  22  // Len returns the number of contained bits
  23  func (bl *BitList) Len() int {
  24  	return bl.count
  25  }
  26  
  27  func (bl *BitList) grow() {
  28  	growBy := len(bl.data)
  29  	if growBy < 128 {
  30  		growBy = 128
  31  	} else if growBy >= 1024 {
  32  		growBy = 1024
  33  	}
  34  
  35  	nd := make([]int32, len(bl.data)+growBy)
  36  	copy(nd, bl.data)
  37  	bl.data = nd
  38  }
  39  
  40  // AddBit appends the given bits to the end of the list
  41  func (bl *BitList) AddBit(bits ...bool) {
  42  	for _, bit := range bits {
  43  		itmIndex := bl.count / 32
  44  		for itmIndex >= len(bl.data) {
  45  			bl.grow()
  46  		}
  47  		bl.SetBit(bl.count, bit)
  48  		bl.count++
  49  	}
  50  }
  51  
  52  // SetBit sets the bit at the given index to the given value
  53  func (bl *BitList) SetBit(index int, value bool) {
  54  	itmIndex := index / 32
  55  	itmBitShift := 31 - (index % 32)
  56  	if value {
  57  		bl.data[itmIndex] = bl.data[itmIndex] | 1<<uint(itmBitShift)
  58  	} else {
  59  		bl.data[itmIndex] = bl.data[itmIndex] & ^(1 << uint(itmBitShift))
  60  	}
  61  }
  62  
  63  // GetBit returns the bit at the given index
  64  func (bl *BitList) GetBit(index int) bool {
  65  	itmIndex := index / 32
  66  	itmBitShift := 31 - (index % 32)
  67  	return ((bl.data[itmIndex] >> uint(itmBitShift)) & 1) == 1
  68  }
  69  
  70  // AddByte appends all 8 bits of the given byte to the end of the list
  71  func (bl *BitList) AddByte(b byte) {
  72  	for i := 7; i >= 0; i-- {
  73  		bl.AddBit(((b >> uint(i)) & 1) == 1)
  74  	}
  75  }
  76  
  77  // AddBits appends the last (LSB) 'count' bits of 'b' the the end of the list
  78  func (bl *BitList) AddBits(b int, count byte) {
  79  	for i := int(count) - 1; i >= 0; i-- {
  80  		bl.AddBit(((b >> uint(i)) & 1) == 1)
  81  	}
  82  }
  83  
  84  // GetBytes returns all bits of the BitList as a []byte
  85  func (bl *BitList) GetBytes() []byte {
  86  	len := bl.count >> 3
  87  	if (bl.count % 8) != 0 {
  88  		len++
  89  	}
  90  	result := make([]byte, len)
  91  	for i := 0; i < len; i++ {
  92  		shift := (3 - (i % 4)) * 8
  93  		result[i] = (byte)((bl.data[i/4] >> uint(shift)) & 0xFF)
  94  	}
  95  	return result
  96  }
  97  
  98  // IterateBytes iterates through all bytes contained in the BitList
  99  func (bl *BitList) IterateBytes() <-chan byte {
 100  	res := make(chan byte)
 101  
 102  	go func() {
 103  		c := bl.count
 104  		shift := 24
 105  		i := 0
 106  		for c > 0 {
 107  			res <- byte((bl.data[i] >> uint(shift)) & 0xFF)
 108  			shift -= 8
 109  			if shift < 0 {
 110  				shift = 24
 111  				i++
 112  			}
 113  			c -= 8
 114  		}
 115  		close(res)
 116  	}()
 117  
 118  	return res
 119  }
 120