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.

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