You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

444 lines
10 KiB

9 years ago
9 years ago
9 years ago
9 years ago
9 years ago
9 years ago
9 years ago
9 years ago
9 years ago
9 years ago
9 years ago
9 years ago
9 years ago
9 years ago
9 years ago
9 years ago
9 years ago
9 years ago
9 years ago
9 years ago
9 years ago
9 years ago
9 years ago
9 years ago
9 years ago
9 years ago
9 years ago
9 years ago
9 years ago
9 years ago
9 years ago
9 years ago
9 years ago
9 years ago
9 years ago
9 years ago
9 years ago
9 years ago
9 years ago
  1. package bits
  2. import (
  3. "encoding/binary"
  4. "fmt"
  5. "regexp"
  6. "strings"
  7. "sync"
  8. tmmath "github.com/tendermint/tendermint/libs/math"
  9. tmrand "github.com/tendermint/tendermint/libs/rand"
  10. tmprotobits "github.com/tendermint/tendermint/proto/tendermint/libs/bits"
  11. )
  12. // BitArray is a thread-safe implementation of a bit array.
  13. type BitArray struct {
  14. mtx sync.Mutex
  15. Bits int `json:"bits"` // NOTE: persisted via reflect, must be exported
  16. Elems []uint64 `json:"elems"` // NOTE: persisted via reflect, must be exported
  17. }
  18. // NewBitArray returns a new bit array.
  19. // It returns nil if the number of bits is zero.
  20. func NewBitArray(bits int) *BitArray {
  21. if bits <= 0 {
  22. return nil
  23. }
  24. return &BitArray{
  25. Bits: bits,
  26. Elems: make([]uint64, (bits+63)/64),
  27. }
  28. }
  29. // Size returns the number of bits in the bitarray
  30. func (bA *BitArray) Size() int {
  31. if bA == nil {
  32. return 0
  33. }
  34. return bA.Bits
  35. }
  36. // GetIndex returns the bit at index i within the bit array.
  37. // The behavior is undefined if i >= bA.Bits
  38. func (bA *BitArray) GetIndex(i int) bool {
  39. if bA == nil {
  40. return false
  41. }
  42. bA.mtx.Lock()
  43. defer bA.mtx.Unlock()
  44. return bA.getIndex(i)
  45. }
  46. func (bA *BitArray) getIndex(i int) bool {
  47. if i >= bA.Bits {
  48. return false
  49. }
  50. return bA.Elems[i/64]&(uint64(1)<<uint(i%64)) > 0
  51. }
  52. // SetIndex sets the bit at index i within the bit array.
  53. // The behavior is undefined if i >= bA.Bits
  54. func (bA *BitArray) SetIndex(i int, v bool) bool {
  55. if bA == nil {
  56. return false
  57. }
  58. bA.mtx.Lock()
  59. defer bA.mtx.Unlock()
  60. return bA.setIndex(i, v)
  61. }
  62. func (bA *BitArray) setIndex(i int, v bool) bool {
  63. if i >= bA.Bits {
  64. return false
  65. }
  66. if v {
  67. bA.Elems[i/64] |= (uint64(1) << uint(i%64))
  68. } else {
  69. bA.Elems[i/64] &= ^(uint64(1) << uint(i%64))
  70. }
  71. return true
  72. }
  73. // Copy returns a copy of the provided bit array.
  74. func (bA *BitArray) Copy() *BitArray {
  75. if bA == nil {
  76. return nil
  77. }
  78. bA.mtx.Lock()
  79. defer bA.mtx.Unlock()
  80. return bA.copy()
  81. }
  82. func (bA *BitArray) copy() *BitArray {
  83. c := make([]uint64, len(bA.Elems))
  84. copy(c, bA.Elems)
  85. return &BitArray{
  86. Bits: bA.Bits,
  87. Elems: c,
  88. }
  89. }
  90. func (bA *BitArray) copyBits(bits int) *BitArray {
  91. c := make([]uint64, (bits+63)/64)
  92. copy(c, bA.Elems)
  93. return &BitArray{
  94. Bits: bits,
  95. Elems: c,
  96. }
  97. }
  98. // Or returns a bit array resulting from a bitwise OR of the two bit arrays.
  99. // If the two bit-arrys have different lengths, Or right-pads the smaller of the two bit-arrays with zeroes.
  100. // Thus the size of the return value is the maximum of the two provided bit arrays.
  101. func (bA *BitArray) Or(o *BitArray) *BitArray {
  102. if bA == nil && o == nil {
  103. return nil
  104. }
  105. if bA == nil && o != nil {
  106. return o.Copy()
  107. }
  108. if o == nil {
  109. return bA.Copy()
  110. }
  111. bA.mtx.Lock()
  112. o.mtx.Lock()
  113. c := bA.copyBits(tmmath.MaxInt(bA.Bits, o.Bits))
  114. smaller := tmmath.MinInt(len(bA.Elems), len(o.Elems))
  115. for i := 0; i < smaller; i++ {
  116. c.Elems[i] |= o.Elems[i]
  117. }
  118. bA.mtx.Unlock()
  119. o.mtx.Unlock()
  120. return c
  121. }
  122. // And returns a bit array resulting from a bitwise AND of the two bit arrays.
  123. // If the two bit-arrys have different lengths, this truncates the larger of the two bit-arrays from the right.
  124. // Thus the size of the return value is the minimum of the two provided bit arrays.
  125. func (bA *BitArray) And(o *BitArray) *BitArray {
  126. if bA == nil || o == nil {
  127. return nil
  128. }
  129. bA.mtx.Lock()
  130. o.mtx.Lock()
  131. defer func() {
  132. bA.mtx.Unlock()
  133. o.mtx.Unlock()
  134. }()
  135. return bA.and(o)
  136. }
  137. func (bA *BitArray) and(o *BitArray) *BitArray {
  138. c := bA.copyBits(tmmath.MinInt(bA.Bits, o.Bits))
  139. for i := 0; i < len(c.Elems); i++ {
  140. c.Elems[i] &= o.Elems[i]
  141. }
  142. return c
  143. }
  144. // Not returns a bit array resulting from a bitwise Not of the provided bit array.
  145. func (bA *BitArray) Not() *BitArray {
  146. if bA == nil {
  147. return nil // Degenerate
  148. }
  149. bA.mtx.Lock()
  150. defer bA.mtx.Unlock()
  151. return bA.not()
  152. }
  153. func (bA *BitArray) not() *BitArray {
  154. c := bA.copy()
  155. for i := 0; i < len(c.Elems); i++ {
  156. c.Elems[i] = ^c.Elems[i]
  157. }
  158. return c
  159. }
  160. // Sub subtracts the two bit-arrays bitwise, without carrying the bits.
  161. // Note that carryless subtraction of a - b is (a and not b).
  162. // The output is the same as bA, regardless of o's size.
  163. // If bA is longer than o, o is right padded with zeroes
  164. func (bA *BitArray) Sub(o *BitArray) *BitArray {
  165. if bA == nil || o == nil {
  166. // TODO: Decide if we should do 1's complement here?
  167. return nil
  168. }
  169. bA.mtx.Lock()
  170. o.mtx.Lock()
  171. // output is the same size as bA
  172. c := bA.copyBits(bA.Bits)
  173. // Only iterate to the minimum size between the two.
  174. // If o is longer, those bits are ignored.
  175. // If bA is longer, then skipping those iterations is equivalent
  176. // to right padding with 0's
  177. smaller := tmmath.MinInt(len(bA.Elems), len(o.Elems))
  178. for i := 0; i < smaller; i++ {
  179. // &^ is and not in golang
  180. c.Elems[i] &^= o.Elems[i]
  181. }
  182. bA.mtx.Unlock()
  183. o.mtx.Unlock()
  184. return c
  185. }
  186. // IsEmpty returns true iff all bits in the bit array are 0
  187. func (bA *BitArray) IsEmpty() bool {
  188. if bA == nil {
  189. return true // should this be opposite?
  190. }
  191. bA.mtx.Lock()
  192. defer bA.mtx.Unlock()
  193. for _, e := range bA.Elems {
  194. if e > 0 {
  195. return false
  196. }
  197. }
  198. return true
  199. }
  200. // IsFull returns true iff all bits in the bit array are 1.
  201. func (bA *BitArray) IsFull() bool {
  202. if bA == nil {
  203. return true
  204. }
  205. bA.mtx.Lock()
  206. defer bA.mtx.Unlock()
  207. // Check all elements except the last
  208. for _, elem := range bA.Elems[:len(bA.Elems)-1] {
  209. if (^elem) != 0 {
  210. return false
  211. }
  212. }
  213. // Check that the last element has (lastElemBits) 1's
  214. lastElemBits := (bA.Bits+63)%64 + 1
  215. lastElem := bA.Elems[len(bA.Elems)-1]
  216. return (lastElem+1)&((uint64(1)<<uint(lastElemBits))-1) == 0
  217. }
  218. // PickRandom returns a random index for a set bit in the bit array.
  219. // If there is no such value, it returns 0, false.
  220. // It uses the global randomness in `random.go` to get this index.
  221. func (bA *BitArray) PickRandom() (int, bool) {
  222. if bA == nil {
  223. return 0, false
  224. }
  225. bA.mtx.Lock()
  226. trueIndices := bA.getTrueIndices()
  227. bA.mtx.Unlock()
  228. if len(trueIndices) == 0 { // no bits set to true
  229. return 0, false
  230. }
  231. return trueIndices[tmrand.Intn(len(trueIndices))], true
  232. }
  233. func (bA *BitArray) getTrueIndices() []int {
  234. trueIndices := make([]int, 0, bA.Bits)
  235. curBit := 0
  236. numElems := len(bA.Elems)
  237. // set all true indices
  238. for i := 0; i < numElems-1; i++ {
  239. elem := bA.Elems[i]
  240. if elem == 0 {
  241. curBit += 64
  242. continue
  243. }
  244. for j := 0; j < 64; j++ {
  245. if (elem & (uint64(1) << uint64(j))) > 0 {
  246. trueIndices = append(trueIndices, curBit)
  247. }
  248. curBit++
  249. }
  250. }
  251. // handle last element
  252. lastElem := bA.Elems[numElems-1]
  253. numFinalBits := bA.Bits - curBit
  254. for i := 0; i < numFinalBits; i++ {
  255. if (lastElem & (uint64(1) << uint64(i))) > 0 {
  256. trueIndices = append(trueIndices, curBit)
  257. }
  258. curBit++
  259. }
  260. return trueIndices
  261. }
  262. // String returns a string representation of BitArray: BA{<bit-string>},
  263. // where <bit-string> is a sequence of 'x' (1) and '_' (0).
  264. // The <bit-string> includes spaces and newlines to help people.
  265. // For a simple sequence of 'x' and '_' characters with no spaces or newlines,
  266. // see the MarshalJSON() method.
  267. // Example: "BA{_x_}" or "nil-BitArray" for nil.
  268. func (bA *BitArray) String() string {
  269. return bA.StringIndented("")
  270. }
  271. // StringIndented returns the same thing as String(), but applies the indent
  272. // at every 10th bit, and twice at every 50th bit.
  273. func (bA *BitArray) StringIndented(indent string) string {
  274. if bA == nil {
  275. return "nil-BitArray"
  276. }
  277. bA.mtx.Lock()
  278. defer bA.mtx.Unlock()
  279. return bA.stringIndented(indent)
  280. }
  281. func (bA *BitArray) stringIndented(indent string) string {
  282. lines := []string{}
  283. bits := ""
  284. for i := 0; i < bA.Bits; i++ {
  285. if bA.getIndex(i) {
  286. bits += "x"
  287. } else {
  288. bits += "_"
  289. }
  290. if i%100 == 99 {
  291. lines = append(lines, bits)
  292. bits = ""
  293. }
  294. if i%10 == 9 {
  295. bits += indent
  296. }
  297. if i%50 == 49 {
  298. bits += indent
  299. }
  300. }
  301. if len(bits) > 0 {
  302. lines = append(lines, bits)
  303. }
  304. return fmt.Sprintf("BA{%v:%v}", bA.Bits, strings.Join(lines, indent))
  305. }
  306. // Bytes returns the byte representation of the bits within the bitarray.
  307. func (bA *BitArray) Bytes() []byte {
  308. bA.mtx.Lock()
  309. defer bA.mtx.Unlock()
  310. numBytes := (bA.Bits + 7) / 8
  311. bytes := make([]byte, numBytes)
  312. for i := 0; i < len(bA.Elems); i++ {
  313. elemBytes := [8]byte{}
  314. binary.LittleEndian.PutUint64(elemBytes[:], bA.Elems[i])
  315. copy(bytes[i*8:], elemBytes[:])
  316. }
  317. return bytes
  318. }
  319. // Update sets the bA's bits to be that of the other bit array.
  320. // The copying begins from the begin of both bit arrays.
  321. func (bA *BitArray) Update(o *BitArray) {
  322. if bA == nil || o == nil {
  323. return
  324. }
  325. bA.mtx.Lock()
  326. o.mtx.Lock()
  327. copy(bA.Elems, o.Elems)
  328. o.mtx.Unlock()
  329. bA.mtx.Unlock()
  330. }
  331. // MarshalJSON implements json.Marshaler interface by marshaling bit array
  332. // using a custom format: a string of '-' or 'x' where 'x' denotes the 1 bit.
  333. func (bA *BitArray) MarshalJSON() ([]byte, error) {
  334. if bA == nil {
  335. return []byte("null"), nil
  336. }
  337. bA.mtx.Lock()
  338. defer bA.mtx.Unlock()
  339. bits := `"`
  340. for i := 0; i < bA.Bits; i++ {
  341. if bA.getIndex(i) {
  342. bits += `x`
  343. } else {
  344. bits += `_`
  345. }
  346. }
  347. bits += `"`
  348. return []byte(bits), nil
  349. }
  350. var bitArrayJSONRegexp = regexp.MustCompile(`\A"([_x]*)"\z`)
  351. // UnmarshalJSON implements json.Unmarshaler interface by unmarshaling a custom
  352. // JSON description.
  353. func (bA *BitArray) UnmarshalJSON(bz []byte) error {
  354. b := string(bz)
  355. if b == "null" {
  356. // This is required e.g. for encoding/json when decoding
  357. // into a pointer with pre-allocated BitArray.
  358. bA.Bits = 0
  359. bA.Elems = nil
  360. return nil
  361. }
  362. // Validate 'b'.
  363. match := bitArrayJSONRegexp.FindStringSubmatch(b)
  364. if match == nil {
  365. return fmt.Errorf("bitArray in JSON should be a string of format %q but got %s", bitArrayJSONRegexp.String(), b)
  366. }
  367. bits := match[1]
  368. // Construct new BitArray and copy over.
  369. numBits := len(bits)
  370. bA2 := NewBitArray(numBits)
  371. for i := 0; i < numBits; i++ {
  372. if bits[i] == 'x' {
  373. bA2.SetIndex(i, true)
  374. }
  375. }
  376. *bA = *bA2 //nolint:govet
  377. return nil
  378. }
  379. // ToProto converts BitArray to protobuf
  380. func (bA *BitArray) ToProto() *tmprotobits.BitArray {
  381. if bA == nil || len(bA.Elems) == 0 {
  382. return nil
  383. }
  384. return &tmprotobits.BitArray{
  385. Bits: int64(bA.Bits),
  386. Elems: bA.Elems,
  387. }
  388. }
  389. // FromProto sets a protobuf BitArray to the given pointer.
  390. func (bA *BitArray) FromProto(protoBitArray *tmprotobits.BitArray) {
  391. if protoBitArray == nil {
  392. bA = nil
  393. return
  394. }
  395. bA.Bits = int(protoBitArray.Bits)
  396. if len(protoBitArray.Elems) > 0 {
  397. bA.Elems = protoBitArray.Elems
  398. }
  399. }