bitset.go 28 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083
  1. /*
  2. Package bitset implements bitsets, a mapping
  3. between non-negative integers and boolean values. It should be more
  4. efficient than map[uint] bool.
  5. It provides methods for setting, clearing, flipping, and testing
  6. individual integers.
  7. But it also provides set intersection, union, difference,
  8. complement, and symmetric operations, as well as tests to
  9. check whether any, all, or no bits are set, and querying a
  10. bitset's current length and number of positive bits.
  11. BitSets are expanded to the size of the largest set bit; the
  12. memory allocation is approximately Max bits, where Max is
  13. the largest set bit. BitSets are never shrunk. On creation,
  14. a hint can be given for the number of bits that will be used.
  15. Many of the methods, including Set,Clear, and Flip, return
  16. a BitSet pointer, which allows for chaining.
  17. Example use:
  18. import "bitset"
  19. var b BitSet
  20. b.Set(10).Set(11)
  21. if b.Test(1000) {
  22. b.Clear(1000)
  23. }
  24. if B.Intersection(bitset.New(100).Set(10)).Count() > 1 {
  25. fmt.Println("Intersection works.")
  26. }
  27. As an alternative to BitSets, one should check out the 'big' package,
  28. which provides a (less set-theoretical) view of bitsets.
  29. */
  30. package bitset
  31. import (
  32. "bytes"
  33. "encoding/base64"
  34. "encoding/binary"
  35. "encoding/json"
  36. "errors"
  37. "fmt"
  38. "io"
  39. "strconv"
  40. )
  41. // the wordSize of a bit set
  42. const wordSize = uint(64)
  43. // the wordSize of a bit set in bytes
  44. const wordBytes = wordSize / 8
  45. // log2WordSize is lg(wordSize)
  46. const log2WordSize = uint(6)
  47. // allBits has every bit set
  48. const allBits uint64 = 0xffffffffffffffff
  49. // default binary BigEndian
  50. var binaryOrder binary.ByteOrder = binary.BigEndian
  51. // default json encoding base64.URLEncoding
  52. var base64Encoding = base64.URLEncoding
  53. // Base64StdEncoding Marshal/Unmarshal BitSet with base64.StdEncoding(Default: base64.URLEncoding)
  54. func Base64StdEncoding() { base64Encoding = base64.StdEncoding }
  55. // LittleEndian Marshal/Unmarshal Binary as Little Endian(Default: binary.BigEndian)
  56. func LittleEndian() { binaryOrder = binary.LittleEndian }
  57. // A BitSet is a set of bits. The zero value of a BitSet is an empty set of length 0.
  58. type BitSet struct {
  59. length uint
  60. set []uint64
  61. }
  62. // Error is used to distinguish errors (panics) generated in this package.
  63. type Error string
  64. // safeSet will fixup b.set to be non-nil and return the field value
  65. func (b *BitSet) safeSet() []uint64 {
  66. if b.set == nil {
  67. b.set = make([]uint64, wordsNeeded(0))
  68. }
  69. return b.set
  70. }
  71. // SetBitsetFrom fills the bitset with an array of integers without creating a new BitSet instance
  72. func (b *BitSet) SetBitsetFrom(buf []uint64) {
  73. b.length = uint(len(buf)) * 64
  74. b.set = buf
  75. }
  76. // From is a constructor used to create a BitSet from an array of integers
  77. func From(buf []uint64) *BitSet {
  78. return FromWithLength(uint(len(buf))*64, buf)
  79. }
  80. // FromWithLength constructs from an array of integers and length.
  81. func FromWithLength(len uint, set []uint64) *BitSet {
  82. return &BitSet{len, set}
  83. }
  84. // Bytes returns the bitset as array of integers
  85. func (b *BitSet) Bytes() []uint64 {
  86. return b.set
  87. }
  88. // wordsNeeded calculates the number of words needed for i bits
  89. func wordsNeeded(i uint) int {
  90. if i > (Cap() - wordSize + 1) {
  91. return int(Cap() >> log2WordSize)
  92. }
  93. return int((i + (wordSize - 1)) >> log2WordSize)
  94. }
  95. // wordsNeededUnbound calculates the number of words needed for i bits, possibly exceeding the capacity.
  96. // This function is useful if you know that the capacity cannot be exceeded (e.g., you have an existing bitmap).
  97. func wordsNeededUnbound(i uint) int {
  98. return int((i + (wordSize - 1)) >> log2WordSize)
  99. }
  100. // wordsIndex calculates the index of words in a `uint64`
  101. func wordsIndex(i uint) uint {
  102. return i & (wordSize - 1)
  103. }
  104. // New creates a new BitSet with a hint that length bits will be required
  105. func New(length uint) (bset *BitSet) {
  106. defer func() {
  107. if r := recover(); r != nil {
  108. bset = &BitSet{
  109. 0,
  110. make([]uint64, 0),
  111. }
  112. }
  113. }()
  114. bset = &BitSet{
  115. length,
  116. make([]uint64, wordsNeeded(length)),
  117. }
  118. return bset
  119. }
  120. // Cap returns the total possible capacity, or number of bits
  121. func Cap() uint {
  122. return ^uint(0)
  123. }
  124. // Len returns the number of bits in the BitSet.
  125. // Note the difference to method Count, see example.
  126. func (b *BitSet) Len() uint {
  127. return b.length
  128. }
  129. // extendSet adds additional words to incorporate new bits if needed
  130. func (b *BitSet) extendSet(i uint) {
  131. if i >= Cap() {
  132. panic("You are exceeding the capacity")
  133. }
  134. nsize := wordsNeeded(i + 1)
  135. if b.set == nil {
  136. b.set = make([]uint64, nsize)
  137. } else if cap(b.set) >= nsize {
  138. b.set = b.set[:nsize] // fast resize
  139. } else if len(b.set) < nsize {
  140. newset := make([]uint64, nsize, 2*nsize) // increase capacity 2x
  141. copy(newset, b.set)
  142. b.set = newset
  143. }
  144. b.length = i + 1
  145. }
  146. // Test whether bit i is set.
  147. func (b *BitSet) Test(i uint) bool {
  148. if i >= b.length {
  149. return false
  150. }
  151. return b.set[i>>log2WordSize]&(1<<wordsIndex(i)) != 0
  152. }
  153. // Set bit i to 1, the capacity of the bitset is automatically
  154. // increased accordingly.
  155. // If i>= Cap(), this function will panic.
  156. // Warning: using a very large value for 'i'
  157. // may lead to a memory shortage and a panic: the caller is responsible
  158. // for providing sensible parameters in line with their memory capacity.
  159. func (b *BitSet) Set(i uint) *BitSet {
  160. if i >= b.length { // if we need more bits, make 'em
  161. b.extendSet(i)
  162. }
  163. b.set[i>>log2WordSize] |= 1 << wordsIndex(i)
  164. return b
  165. }
  166. // Clear bit i to 0
  167. func (b *BitSet) Clear(i uint) *BitSet {
  168. if i >= b.length {
  169. return b
  170. }
  171. b.set[i>>log2WordSize] &^= 1 << wordsIndex(i)
  172. return b
  173. }
  174. // SetTo sets bit i to value.
  175. // If i>= Cap(), this function will panic.
  176. // Warning: using a very large value for 'i'
  177. // may lead to a memory shortage and a panic: the caller is responsible
  178. // for providing sensible parameters in line with their memory capacity.
  179. func (b *BitSet) SetTo(i uint, value bool) *BitSet {
  180. if value {
  181. return b.Set(i)
  182. }
  183. return b.Clear(i)
  184. }
  185. // Flip bit at i.
  186. // If i>= Cap(), this function will panic.
  187. // Warning: using a very large value for 'i'
  188. // may lead to a memory shortage and a panic: the caller is responsible
  189. // for providing sensible parameters in line with their memory capacity.
  190. func (b *BitSet) Flip(i uint) *BitSet {
  191. if i >= b.length {
  192. return b.Set(i)
  193. }
  194. b.set[i>>log2WordSize] ^= 1 << wordsIndex(i)
  195. return b
  196. }
  197. // FlipRange bit in [start, end).
  198. // If end>= Cap(), this function will panic.
  199. // Warning: using a very large value for 'end'
  200. // may lead to a memory shortage and a panic: the caller is responsible
  201. // for providing sensible parameters in line with their memory capacity.
  202. func (b *BitSet) FlipRange(start, end uint) *BitSet {
  203. if start >= end {
  204. return b
  205. }
  206. if end-1 >= b.length { // if we need more bits, make 'em
  207. b.extendSet(end - 1)
  208. }
  209. var startWord uint = start >> log2WordSize
  210. var endWord uint = end >> log2WordSize
  211. b.set[startWord] ^= ^(^uint64(0) << wordsIndex(start))
  212. for i := startWord; i < endWord; i++ {
  213. b.set[i] = ^b.set[i]
  214. }
  215. if end&(wordSize-1) != 0 {
  216. b.set[endWord] ^= ^uint64(0) >> wordsIndex(-end)
  217. }
  218. return b
  219. }
  220. // Shrink shrinks BitSet so that the provided value is the last possible
  221. // set value. It clears all bits > the provided index and reduces the size
  222. // and length of the set.
  223. //
  224. // Note that the parameter value is not the new length in bits: it is the
  225. // maximal value that can be stored in the bitset after the function call.
  226. // The new length in bits is the parameter value + 1. Thus it is not possible
  227. // to use this function to set the length to 0, the minimal value of the length
  228. // after this function call is 1.
  229. //
  230. // A new slice is allocated to store the new bits, so you may see an increase in
  231. // memory usage until the GC runs. Normally this should not be a problem, but if you
  232. // have an extremely large BitSet its important to understand that the old BitSet will
  233. // remain in memory until the GC frees it.
  234. func (b *BitSet) Shrink(lastbitindex uint) *BitSet {
  235. length := lastbitindex + 1
  236. idx := wordsNeeded(length)
  237. if idx > len(b.set) {
  238. return b
  239. }
  240. shrunk := make([]uint64, idx)
  241. copy(shrunk, b.set[:idx])
  242. b.set = shrunk
  243. b.length = length
  244. lastWordUsedBits := length % 64
  245. if lastWordUsedBits != 0 {
  246. b.set[idx-1] &= allBits >> uint64(64-wordsIndex(lastWordUsedBits))
  247. }
  248. return b
  249. }
  250. // Compact shrinks BitSet to so that we preserve all set bits, while minimizing
  251. // memory usage. Compact calls Shrink.
  252. func (b *BitSet) Compact() *BitSet {
  253. idx := len(b.set) - 1
  254. for ; idx >= 0 && b.set[idx] == 0; idx-- {
  255. }
  256. newlength := uint((idx + 1) << log2WordSize)
  257. if newlength >= b.length {
  258. return b // nothing to do
  259. }
  260. if newlength > 0 {
  261. return b.Shrink(newlength - 1)
  262. }
  263. // We preserve one word
  264. return b.Shrink(63)
  265. }
  266. // InsertAt takes an index which indicates where a bit should be
  267. // inserted. Then it shifts all the bits in the set to the left by 1, starting
  268. // from the given index position, and sets the index position to 0.
  269. //
  270. // Depending on the size of your BitSet, and where you are inserting the new entry,
  271. // this method could be extremely slow and in some cases might cause the entire BitSet
  272. // to be recopied.
  273. func (b *BitSet) InsertAt(idx uint) *BitSet {
  274. insertAtElement := idx >> log2WordSize
  275. // if length of set is a multiple of wordSize we need to allocate more space first
  276. if b.isLenExactMultiple() {
  277. b.set = append(b.set, uint64(0))
  278. }
  279. var i uint
  280. for i = uint(len(b.set) - 1); i > insertAtElement; i-- {
  281. // all elements above the position where we want to insert can simply by shifted
  282. b.set[i] <<= 1
  283. // we take the most significant bit of the previous element and set it as
  284. // the least significant bit of the current element
  285. b.set[i] |= (b.set[i-1] & 0x8000000000000000) >> 63
  286. }
  287. // generate a mask to extract the data that we need to shift left
  288. // within the element where we insert a bit
  289. dataMask := uint64(1)<<uint64(wordsIndex(idx)) - 1
  290. // extract that data that we'll shift
  291. data := b.set[i] & (^dataMask)
  292. // set the positions of the data mask to 0 in the element where we insert
  293. b.set[i] &= dataMask
  294. // shift data mask to the left and insert its data to the slice element
  295. b.set[i] |= data << 1
  296. // add 1 to length of BitSet
  297. b.length++
  298. return b
  299. }
  300. // String creates a string representation of the Bitmap
  301. func (b *BitSet) String() string {
  302. // follows code from https://github.com/RoaringBitmap/roaring
  303. var buffer bytes.Buffer
  304. start := []byte("{")
  305. buffer.Write(start)
  306. counter := 0
  307. i, e := b.NextSet(0)
  308. for e {
  309. counter = counter + 1
  310. // to avoid exhausting the memory
  311. if counter > 0x40000 {
  312. buffer.WriteString("...")
  313. break
  314. }
  315. buffer.WriteString(strconv.FormatInt(int64(i), 10))
  316. i, e = b.NextSet(i + 1)
  317. if e {
  318. buffer.WriteString(",")
  319. }
  320. }
  321. buffer.WriteString("}")
  322. return buffer.String()
  323. }
  324. // DeleteAt deletes the bit at the given index position from
  325. // within the bitset
  326. // All the bits residing on the left of the deleted bit get
  327. // shifted right by 1
  328. // The running time of this operation may potentially be
  329. // relatively slow, O(length)
  330. func (b *BitSet) DeleteAt(i uint) *BitSet {
  331. // the index of the slice element where we'll delete a bit
  332. deleteAtElement := i >> log2WordSize
  333. // generate a mask for the data that needs to be shifted right
  334. // within that slice element that gets modified
  335. dataMask := ^((uint64(1) << wordsIndex(i)) - 1)
  336. // extract the data that we'll shift right from the slice element
  337. data := b.set[deleteAtElement] & dataMask
  338. // set the masked area to 0 while leaving the rest as it is
  339. b.set[deleteAtElement] &= ^dataMask
  340. // shift the previously extracted data to the right and then
  341. // set it in the previously masked area
  342. b.set[deleteAtElement] |= (data >> 1) & dataMask
  343. // loop over all the consecutive slice elements to copy each
  344. // lowest bit into the highest position of the previous element,
  345. // then shift the entire content to the right by 1
  346. for i := int(deleteAtElement) + 1; i < len(b.set); i++ {
  347. b.set[i-1] |= (b.set[i] & 1) << 63
  348. b.set[i] >>= 1
  349. }
  350. b.length = b.length - 1
  351. return b
  352. }
  353. // NextSet returns the next bit set from the specified index,
  354. // including possibly the current index
  355. // along with an error code (true = valid, false = no set bit found)
  356. // for i,e := v.NextSet(0); e; i,e = v.NextSet(i + 1) {...}
  357. //
  358. // Users concerned with performance may want to use NextSetMany to
  359. // retrieve several values at once.
  360. func (b *BitSet) NextSet(i uint) (uint, bool) {
  361. x := int(i >> log2WordSize)
  362. if x >= len(b.set) {
  363. return 0, false
  364. }
  365. w := b.set[x]
  366. w = w >> wordsIndex(i)
  367. if w != 0 {
  368. return i + trailingZeroes64(w), true
  369. }
  370. x = x + 1
  371. for x < len(b.set) {
  372. if b.set[x] != 0 {
  373. return uint(x)*wordSize + trailingZeroes64(b.set[x]), true
  374. }
  375. x = x + 1
  376. }
  377. return 0, false
  378. }
  379. // NextSetMany returns many next bit sets from the specified index,
  380. // including possibly the current index and up to cap(buffer).
  381. // If the returned slice has len zero, then no more set bits were found
  382. //
  383. // buffer := make([]uint, 256) // this should be reused
  384. // j := uint(0)
  385. // j, buffer = bitmap.NextSetMany(j, buffer)
  386. // for ; len(buffer) > 0; j, buffer = bitmap.NextSetMany(j,buffer) {
  387. // for k := range buffer {
  388. // do something with buffer[k]
  389. // }
  390. // j += 1
  391. // }
  392. //
  393. // It is possible to retrieve all set bits as follow:
  394. //
  395. // indices := make([]uint, bitmap.Count())
  396. // bitmap.NextSetMany(0, indices)
  397. //
  398. // However if bitmap.Count() is large, it might be preferable to
  399. // use several calls to NextSetMany, for performance reasons.
  400. func (b *BitSet) NextSetMany(i uint, buffer []uint) (uint, []uint) {
  401. myanswer := buffer
  402. capacity := cap(buffer)
  403. x := int(i >> log2WordSize)
  404. if x >= len(b.set) || capacity == 0 {
  405. return 0, myanswer[:0]
  406. }
  407. skip := wordsIndex(i)
  408. word := b.set[x] >> skip
  409. myanswer = myanswer[:capacity]
  410. size := int(0)
  411. for word != 0 {
  412. r := trailingZeroes64(word)
  413. t := word & ((^word) + 1)
  414. myanswer[size] = r + i
  415. size++
  416. if size == capacity {
  417. goto End
  418. }
  419. word = word ^ t
  420. }
  421. x++
  422. for idx, word := range b.set[x:] {
  423. for word != 0 {
  424. r := trailingZeroes64(word)
  425. t := word & ((^word) + 1)
  426. myanswer[size] = r + (uint(x+idx) << 6)
  427. size++
  428. if size == capacity {
  429. goto End
  430. }
  431. word = word ^ t
  432. }
  433. }
  434. End:
  435. if size > 0 {
  436. return myanswer[size-1], myanswer[:size]
  437. }
  438. return 0, myanswer[:0]
  439. }
  440. // NextClear returns the next clear bit from the specified index,
  441. // including possibly the current index
  442. // along with an error code (true = valid, false = no bit found i.e. all bits are set)
  443. func (b *BitSet) NextClear(i uint) (uint, bool) {
  444. x := int(i >> log2WordSize)
  445. if x >= len(b.set) {
  446. return 0, false
  447. }
  448. w := b.set[x]
  449. w = w >> wordsIndex(i)
  450. wA := allBits >> wordsIndex(i)
  451. index := i + trailingZeroes64(^w)
  452. if w != wA && index < b.length {
  453. return index, true
  454. }
  455. x++
  456. for x < len(b.set) {
  457. index = uint(x)*wordSize + trailingZeroes64(^b.set[x])
  458. if b.set[x] != allBits && index < b.length {
  459. return index, true
  460. }
  461. x++
  462. }
  463. return 0, false
  464. }
  465. // ClearAll clears the entire BitSet
  466. func (b *BitSet) ClearAll() *BitSet {
  467. if b != nil && b.set != nil {
  468. for i := range b.set {
  469. b.set[i] = 0
  470. }
  471. }
  472. return b
  473. }
  474. // wordCount returns the number of words used in a bit set
  475. func (b *BitSet) wordCount() int {
  476. return wordsNeededUnbound(b.length)
  477. }
  478. // Clone this BitSet
  479. func (b *BitSet) Clone() *BitSet {
  480. c := New(b.length)
  481. if b.set != nil { // Clone should not modify current object
  482. copy(c.set, b.set)
  483. }
  484. return c
  485. }
  486. // Copy into a destination BitSet using the Go array copy semantics:
  487. // the number of bits copied is the minimum of the number of bits in the current
  488. // BitSet (Len()) and the destination Bitset.
  489. // We return the number of bits copied in the destination BitSet.
  490. func (b *BitSet) Copy(c *BitSet) (count uint) {
  491. if c == nil {
  492. return
  493. }
  494. if b.set != nil { // Copy should not modify current object
  495. copy(c.set, b.set)
  496. }
  497. count = c.length
  498. if b.length < c.length {
  499. count = b.length
  500. }
  501. // Cleaning the last word is needed to keep the invariant that other functions, such as Count, require
  502. // that any bits in the last word that would exceed the length of the bitmask are set to 0.
  503. c.cleanLastWord()
  504. return
  505. }
  506. // CopyFull copies into a destination BitSet such that the destination is
  507. // identical to the source after the operation, allocating memory if necessary.
  508. func (b *BitSet) CopyFull(c *BitSet) {
  509. if c == nil {
  510. return
  511. }
  512. c.length = b.length
  513. if len(b.set) == 0 {
  514. if c.set != nil {
  515. c.set = c.set[:0]
  516. }
  517. } else {
  518. if cap(c.set) < len(b.set) {
  519. c.set = make([]uint64, len(b.set))
  520. } else {
  521. c.set = c.set[:len(b.set)]
  522. }
  523. copy(c.set, b.set)
  524. }
  525. }
  526. // Count (number of set bits).
  527. // Also known as "popcount" or "population count".
  528. func (b *BitSet) Count() uint {
  529. if b != nil && b.set != nil {
  530. return uint(popcntSlice(b.set))
  531. }
  532. return 0
  533. }
  534. // Equal tests the equivalence of two BitSets.
  535. // False if they are of different sizes, otherwise true
  536. // only if all the same bits are set
  537. func (b *BitSet) Equal(c *BitSet) bool {
  538. if c == nil || b == nil {
  539. return c == b
  540. }
  541. if b.length != c.length {
  542. return false
  543. }
  544. if b.length == 0 { // if they have both length == 0, then could have nil set
  545. return true
  546. }
  547. wn := b.wordCount()
  548. for p := 0; p < wn; p++ {
  549. if c.set[p] != b.set[p] {
  550. return false
  551. }
  552. }
  553. return true
  554. }
  555. func panicIfNull(b *BitSet) {
  556. if b == nil {
  557. panic(Error("BitSet must not be null"))
  558. }
  559. }
  560. // Difference of base set and other set
  561. // This is the BitSet equivalent of &^ (and not)
  562. func (b *BitSet) Difference(compare *BitSet) (result *BitSet) {
  563. panicIfNull(b)
  564. panicIfNull(compare)
  565. result = b.Clone() // clone b (in case b is bigger than compare)
  566. l := int(compare.wordCount())
  567. if l > int(b.wordCount()) {
  568. l = int(b.wordCount())
  569. }
  570. for i := 0; i < l; i++ {
  571. result.set[i] = b.set[i] &^ compare.set[i]
  572. }
  573. return
  574. }
  575. // DifferenceCardinality computes the cardinality of the differnce
  576. func (b *BitSet) DifferenceCardinality(compare *BitSet) uint {
  577. panicIfNull(b)
  578. panicIfNull(compare)
  579. l := int(compare.wordCount())
  580. if l > int(b.wordCount()) {
  581. l = int(b.wordCount())
  582. }
  583. cnt := uint64(0)
  584. cnt += popcntMaskSlice(b.set[:l], compare.set[:l])
  585. cnt += popcntSlice(b.set[l:])
  586. return uint(cnt)
  587. }
  588. // InPlaceDifference computes the difference of base set and other set
  589. // This is the BitSet equivalent of &^ (and not)
  590. func (b *BitSet) InPlaceDifference(compare *BitSet) {
  591. panicIfNull(b)
  592. panicIfNull(compare)
  593. l := int(compare.wordCount())
  594. if l > int(b.wordCount()) {
  595. l = int(b.wordCount())
  596. }
  597. for i := 0; i < l; i++ {
  598. b.set[i] &^= compare.set[i]
  599. }
  600. }
  601. // Convenience function: return two bitsets ordered by
  602. // increasing length. Note: neither can be nil
  603. func sortByLength(a *BitSet, b *BitSet) (ap *BitSet, bp *BitSet) {
  604. if a.length <= b.length {
  605. ap, bp = a, b
  606. } else {
  607. ap, bp = b, a
  608. }
  609. return
  610. }
  611. // Intersection of base set and other set
  612. // This is the BitSet equivalent of & (and)
  613. func (b *BitSet) Intersection(compare *BitSet) (result *BitSet) {
  614. panicIfNull(b)
  615. panicIfNull(compare)
  616. b, compare = sortByLength(b, compare)
  617. result = New(b.length)
  618. for i, word := range b.set {
  619. result.set[i] = word & compare.set[i]
  620. }
  621. return
  622. }
  623. // IntersectionCardinality computes the cardinality of the union
  624. func (b *BitSet) IntersectionCardinality(compare *BitSet) uint {
  625. panicIfNull(b)
  626. panicIfNull(compare)
  627. b, compare = sortByLength(b, compare)
  628. cnt := popcntAndSlice(b.set, compare.set)
  629. return uint(cnt)
  630. }
  631. // InPlaceIntersection destructively computes the intersection of
  632. // base set and the compare set.
  633. // This is the BitSet equivalent of & (and)
  634. func (b *BitSet) InPlaceIntersection(compare *BitSet) {
  635. panicIfNull(b)
  636. panicIfNull(compare)
  637. l := int(compare.wordCount())
  638. if l > int(b.wordCount()) {
  639. l = int(b.wordCount())
  640. }
  641. for i := 0; i < l; i++ {
  642. b.set[i] &= compare.set[i]
  643. }
  644. for i := l; i < len(b.set); i++ {
  645. b.set[i] = 0
  646. }
  647. if compare.length > 0 {
  648. if compare.length-1 >= b.length {
  649. b.extendSet(compare.length - 1)
  650. }
  651. }
  652. }
  653. // Union of base set and other set
  654. // This is the BitSet equivalent of | (or)
  655. func (b *BitSet) Union(compare *BitSet) (result *BitSet) {
  656. panicIfNull(b)
  657. panicIfNull(compare)
  658. b, compare = sortByLength(b, compare)
  659. result = compare.Clone()
  660. for i, word := range b.set {
  661. result.set[i] = word | compare.set[i]
  662. }
  663. return
  664. }
  665. // UnionCardinality computes the cardinality of the uniton of the base set
  666. // and the compare set.
  667. func (b *BitSet) UnionCardinality(compare *BitSet) uint {
  668. panicIfNull(b)
  669. panicIfNull(compare)
  670. b, compare = sortByLength(b, compare)
  671. cnt := popcntOrSlice(b.set, compare.set)
  672. if len(compare.set) > len(b.set) {
  673. cnt += popcntSlice(compare.set[len(b.set):])
  674. }
  675. return uint(cnt)
  676. }
  677. // InPlaceUnion creates the destructive union of base set and compare set.
  678. // This is the BitSet equivalent of | (or).
  679. func (b *BitSet) InPlaceUnion(compare *BitSet) {
  680. panicIfNull(b)
  681. panicIfNull(compare)
  682. l := int(compare.wordCount())
  683. if l > int(b.wordCount()) {
  684. l = int(b.wordCount())
  685. }
  686. if compare.length > 0 && compare.length-1 >= b.length {
  687. b.extendSet(compare.length - 1)
  688. }
  689. for i := 0; i < l; i++ {
  690. b.set[i] |= compare.set[i]
  691. }
  692. if len(compare.set) > l {
  693. for i := l; i < len(compare.set); i++ {
  694. b.set[i] = compare.set[i]
  695. }
  696. }
  697. }
  698. // SymmetricDifference of base set and other set
  699. // This is the BitSet equivalent of ^ (xor)
  700. func (b *BitSet) SymmetricDifference(compare *BitSet) (result *BitSet) {
  701. panicIfNull(b)
  702. panicIfNull(compare)
  703. b, compare = sortByLength(b, compare)
  704. // compare is bigger, so clone it
  705. result = compare.Clone()
  706. for i, word := range b.set {
  707. result.set[i] = word ^ compare.set[i]
  708. }
  709. return
  710. }
  711. // SymmetricDifferenceCardinality computes the cardinality of the symmetric difference
  712. func (b *BitSet) SymmetricDifferenceCardinality(compare *BitSet) uint {
  713. panicIfNull(b)
  714. panicIfNull(compare)
  715. b, compare = sortByLength(b, compare)
  716. cnt := popcntXorSlice(b.set, compare.set)
  717. if len(compare.set) > len(b.set) {
  718. cnt += popcntSlice(compare.set[len(b.set):])
  719. }
  720. return uint(cnt)
  721. }
  722. // InPlaceSymmetricDifference creates the destructive SymmetricDifference of base set and other set
  723. // This is the BitSet equivalent of ^ (xor)
  724. func (b *BitSet) InPlaceSymmetricDifference(compare *BitSet) {
  725. panicIfNull(b)
  726. panicIfNull(compare)
  727. l := int(compare.wordCount())
  728. if l > int(b.wordCount()) {
  729. l = int(b.wordCount())
  730. }
  731. if compare.length > 0 && compare.length-1 >= b.length {
  732. b.extendSet(compare.length - 1)
  733. }
  734. for i := 0; i < l; i++ {
  735. b.set[i] ^= compare.set[i]
  736. }
  737. if len(compare.set) > l {
  738. for i := l; i < len(compare.set); i++ {
  739. b.set[i] = compare.set[i]
  740. }
  741. }
  742. }
  743. // Is the length an exact multiple of word sizes?
  744. func (b *BitSet) isLenExactMultiple() bool {
  745. return wordsIndex(b.length) == 0
  746. }
  747. // Clean last word by setting unused bits to 0
  748. func (b *BitSet) cleanLastWord() {
  749. if !b.isLenExactMultiple() {
  750. b.set[len(b.set)-1] &= allBits >> (wordSize - wordsIndex(b.length))
  751. }
  752. }
  753. // Complement computes the (local) complement of a bitset (up to length bits)
  754. func (b *BitSet) Complement() (result *BitSet) {
  755. panicIfNull(b)
  756. result = New(b.length)
  757. for i, word := range b.set {
  758. result.set[i] = ^word
  759. }
  760. result.cleanLastWord()
  761. return
  762. }
  763. // All returns true if all bits are set, false otherwise. Returns true for
  764. // empty sets.
  765. func (b *BitSet) All() bool {
  766. panicIfNull(b)
  767. return b.Count() == b.length
  768. }
  769. // None returns true if no bit is set, false otherwise. Returns true for
  770. // empty sets.
  771. func (b *BitSet) None() bool {
  772. panicIfNull(b)
  773. if b != nil && b.set != nil {
  774. for _, word := range b.set {
  775. if word > 0 {
  776. return false
  777. }
  778. }
  779. }
  780. return true
  781. }
  782. // Any returns true if any bit is set, false otherwise
  783. func (b *BitSet) Any() bool {
  784. panicIfNull(b)
  785. return !b.None()
  786. }
  787. // IsSuperSet returns true if this is a superset of the other set
  788. func (b *BitSet) IsSuperSet(other *BitSet) bool {
  789. for i, e := other.NextSet(0); e; i, e = other.NextSet(i + 1) {
  790. if !b.Test(i) {
  791. return false
  792. }
  793. }
  794. return true
  795. }
  796. // IsStrictSuperSet returns true if this is a strict superset of the other set
  797. func (b *BitSet) IsStrictSuperSet(other *BitSet) bool {
  798. return b.Count() > other.Count() && b.IsSuperSet(other)
  799. }
  800. // DumpAsBits dumps a bit set as a string of bits
  801. func (b *BitSet) DumpAsBits() string {
  802. if b.set == nil {
  803. return "."
  804. }
  805. buffer := bytes.NewBufferString("")
  806. i := len(b.set) - 1
  807. for ; i >= 0; i-- {
  808. fmt.Fprintf(buffer, "%064b.", b.set[i])
  809. }
  810. return buffer.String()
  811. }
  812. // BinaryStorageSize returns the binary storage requirements (see WriteTo) in bytes.
  813. func (b *BitSet) BinaryStorageSize() int {
  814. return int(wordBytes + wordBytes*uint(b.wordCount()))
  815. }
  816. func readUint64Array(reader io.Reader, data []uint64) error {
  817. length := len(data)
  818. bufferSize := 128
  819. buffer := make([]byte, bufferSize*int(wordBytes))
  820. for i := 0; i < length; i += bufferSize {
  821. end := i + bufferSize
  822. if end > length {
  823. end = length
  824. buffer = buffer[:wordBytes*uint(end-i)]
  825. }
  826. chunk := data[i:end]
  827. if _, err := io.ReadFull(reader, buffer); err != nil {
  828. return err
  829. }
  830. for i := range chunk {
  831. chunk[i] = uint64(binaryOrder.Uint64(buffer[8*i:]))
  832. }
  833. }
  834. return nil
  835. }
  836. func writeUint64Array(writer io.Writer, data []uint64) error {
  837. bufferSize := 128
  838. buffer := make([]byte, bufferSize*int(wordBytes))
  839. for i := 0; i < len(data); i += bufferSize {
  840. end := i + bufferSize
  841. if end > len(data) {
  842. end = len(data)
  843. buffer = buffer[:wordBytes*uint(end-i)]
  844. }
  845. chunk := data[i:end]
  846. for i, x := range chunk {
  847. binaryOrder.PutUint64(buffer[8*i:], x)
  848. }
  849. _, err := writer.Write(buffer)
  850. if err != nil {
  851. return err
  852. }
  853. }
  854. return nil
  855. }
  856. // WriteTo writes a BitSet to a stream. The format is:
  857. // 1. uint64 length
  858. // 2. []uint64 set
  859. // Upon success, the number of bytes written is returned.
  860. //
  861. // Performance: if this function is used to write to a disk or network
  862. // connection, it might be beneficial to wrap the stream in a bufio.Writer.
  863. // E.g.,
  864. //
  865. // f, err := os.Create("myfile")
  866. // w := bufio.NewWriter(f)
  867. func (b *BitSet) WriteTo(stream io.Writer) (int64, error) {
  868. length := uint64(b.length)
  869. // Write length
  870. err := binary.Write(stream, binaryOrder, &length)
  871. if err != nil {
  872. // Upon failure, we do not guarantee that we
  873. // return the number of bytes written.
  874. return int64(0), err
  875. }
  876. err = writeUint64Array(stream, b.set[:b.wordCount()])
  877. if err != nil {
  878. // Upon failure, we do not guarantee that we
  879. // return the number of bytes written.
  880. return int64(wordBytes), err
  881. }
  882. return int64(b.BinaryStorageSize()), nil
  883. }
  884. // ReadFrom reads a BitSet from a stream written using WriteTo
  885. // The format is:
  886. // 1. uint64 length
  887. // 2. []uint64 set
  888. // Upon success, the number of bytes read is returned.
  889. // If the current BitSet is not large enough to hold the data,
  890. // it is extended. In case of error, the BitSet is either
  891. // left unchanged or made empty if the error occurs too late
  892. // to preserve the content.
  893. //
  894. // Performance: if this function is used to read from a disk or network
  895. // connection, it might be beneficial to wrap the stream in a bufio.Reader.
  896. // E.g.,
  897. //
  898. // f, err := os.Open("myfile")
  899. // r := bufio.NewReader(f)
  900. func (b *BitSet) ReadFrom(stream io.Reader) (int64, error) {
  901. var length uint64
  902. err := binary.Read(stream, binaryOrder, &length)
  903. if err != nil {
  904. if err == io.EOF {
  905. err = io.ErrUnexpectedEOF
  906. }
  907. return 0, err
  908. }
  909. newlength := uint(length)
  910. if uint64(newlength) != length {
  911. return 0, errors.New("unmarshalling error: type mismatch")
  912. }
  913. nWords := wordsNeeded(uint(newlength))
  914. if cap(b.set) >= nWords {
  915. b.set = b.set[:nWords]
  916. } else {
  917. b.set = make([]uint64, nWords)
  918. }
  919. b.length = newlength
  920. err = readUint64Array(stream, b.set)
  921. if err != nil {
  922. if err == io.EOF {
  923. err = io.ErrUnexpectedEOF
  924. }
  925. // We do not want to leave the BitSet partially filled as
  926. // it is error prone.
  927. b.set = b.set[:0]
  928. b.length = 0
  929. return 0, err
  930. }
  931. return int64(b.BinaryStorageSize()), nil
  932. }
  933. // MarshalBinary encodes a BitSet into a binary form and returns the result.
  934. func (b *BitSet) MarshalBinary() ([]byte, error) {
  935. var buf bytes.Buffer
  936. _, err := b.WriteTo(&buf)
  937. if err != nil {
  938. return []byte{}, err
  939. }
  940. return buf.Bytes(), err
  941. }
  942. // UnmarshalBinary decodes the binary form generated by MarshalBinary.
  943. func (b *BitSet) UnmarshalBinary(data []byte) error {
  944. buf := bytes.NewReader(data)
  945. _, err := b.ReadFrom(buf)
  946. return err
  947. }
  948. // MarshalJSON marshals a BitSet as a JSON structure
  949. func (b BitSet) MarshalJSON() ([]byte, error) {
  950. buffer := bytes.NewBuffer(make([]byte, 0, b.BinaryStorageSize()))
  951. _, err := b.WriteTo(buffer)
  952. if err != nil {
  953. return nil, err
  954. }
  955. // URLEncode all bytes
  956. return json.Marshal(base64Encoding.EncodeToString(buffer.Bytes()))
  957. }
  958. // UnmarshalJSON unmarshals a BitSet from JSON created using MarshalJSON
  959. func (b *BitSet) UnmarshalJSON(data []byte) error {
  960. // Unmarshal as string
  961. var s string
  962. err := json.Unmarshal(data, &s)
  963. if err != nil {
  964. return err
  965. }
  966. // URLDecode string
  967. buf, err := base64Encoding.DecodeString(s)
  968. if err != nil {
  969. return err
  970. }
  971. _, err = b.ReadFrom(bytes.NewReader(buf))
  972. return err
  973. }