bitlist.go 2.5 KB

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