Ви не можете вибрати більше 25 тем Теми мають розпочинатися з літери або цифри, можуть містити дефіси (-) і не повинні перевищувати 35 символів.

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759
  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. "bufio"
  33. "bytes"
  34. "encoding/base64"
  35. "encoding/binary"
  36. "encoding/json"
  37. "errors"
  38. "fmt"
  39. "io"
  40. "strconv"
  41. )
  42. // the wordSize of a bit set
  43. const wordSize = uint(64)
  44. // log2WordSize is lg(wordSize)
  45. const log2WordSize = uint(6)
  46. // allBits has every bit set
  47. const allBits uint64 = 0xffffffffffffffff
  48. // A BitSet is a set of bits. The zero value of a BitSet is an empty set of length 0.
  49. type BitSet struct {
  50. length uint
  51. set []uint64
  52. }
  53. // Error is used to distinguish errors (panics) generated in this package.
  54. type Error string
  55. // safeSet will fixup b.set to be non-nil and return the field value
  56. func (b *BitSet) safeSet() []uint64 {
  57. if b.set == nil {
  58. b.set = make([]uint64, wordsNeeded(0))
  59. }
  60. return b.set
  61. }
  62. // From is a constructor used to create a BitSet from an array of integers
  63. func From(buf []uint64) *BitSet {
  64. return &BitSet{uint(len(buf)) * 64, buf}
  65. }
  66. // Bytes returns the bitset as array of integers
  67. func (b *BitSet) Bytes() []uint64 {
  68. return b.set
  69. }
  70. // wordsNeeded calculates the number of words needed for i bits
  71. func wordsNeeded(i uint) int {
  72. if i > (Cap() - wordSize + 1) {
  73. return int(Cap() >> log2WordSize)
  74. }
  75. return int((i + (wordSize - 1)) >> log2WordSize)
  76. }
  77. // New creates a new BitSet with a hint that length bits will be required
  78. func New(length uint) (bset *BitSet) {
  79. defer func() {
  80. if r := recover(); r != nil {
  81. bset = &BitSet{
  82. 0,
  83. make([]uint64, 0),
  84. }
  85. }
  86. }()
  87. bset = &BitSet{
  88. length,
  89. make([]uint64, wordsNeeded(length)),
  90. }
  91. return bset
  92. }
  93. // Cap returns the total possible capacity, or number of bits
  94. func Cap() uint {
  95. return ^uint(0)
  96. }
  97. // Len returns the length of the BitSet in words
  98. func (b *BitSet) Len() uint {
  99. return b.length
  100. }
  101. // extendSetMaybe adds additional words to incorporate new bits if needed
  102. func (b *BitSet) extendSetMaybe(i uint) {
  103. if i >= b.length { // if we need more bits, make 'em
  104. nsize := wordsNeeded(i + 1)
  105. if b.set == nil {
  106. b.set = make([]uint64, nsize)
  107. } else if cap(b.set) >= nsize {
  108. b.set = b.set[:nsize] // fast resize
  109. } else if len(b.set) < nsize {
  110. newset := make([]uint64, nsize, 2*nsize) // increase capacity 2x
  111. copy(newset, b.set)
  112. b.set = newset
  113. }
  114. b.length = i + 1
  115. }
  116. }
  117. // Test whether bit i is set.
  118. func (b *BitSet) Test(i uint) bool {
  119. if i >= b.length {
  120. return false
  121. }
  122. return b.set[i>>log2WordSize]&(1<<(i&(wordSize-1))) != 0
  123. }
  124. // Set bit i to 1
  125. func (b *BitSet) Set(i uint) *BitSet {
  126. b.extendSetMaybe(i)
  127. b.set[i>>log2WordSize] |= 1 << (i & (wordSize - 1))
  128. return b
  129. }
  130. // Clear bit i to 0
  131. func (b *BitSet) Clear(i uint) *BitSet {
  132. if i >= b.length {
  133. return b
  134. }
  135. b.set[i>>log2WordSize] &^= 1 << (i & (wordSize - 1))
  136. return b
  137. }
  138. // SetTo sets bit i to value
  139. func (b *BitSet) SetTo(i uint, value bool) *BitSet {
  140. if value {
  141. return b.Set(i)
  142. }
  143. return b.Clear(i)
  144. }
  145. // Flip bit at i
  146. func (b *BitSet) Flip(i uint) *BitSet {
  147. if i >= b.length {
  148. return b.Set(i)
  149. }
  150. b.set[i>>log2WordSize] ^= 1 << (i & (wordSize - 1))
  151. return b
  152. }
  153. // String creates a string representation of the Bitmap
  154. func (b *BitSet) String() string {
  155. // follows code from https://github.com/RoaringBitmap/roaring
  156. var buffer bytes.Buffer
  157. start := []byte("{")
  158. buffer.Write(start)
  159. counter := 0
  160. i, e := b.NextSet(0)
  161. for e {
  162. counter = counter + 1
  163. // to avoid exhausting the memory
  164. if counter > 0x40000 {
  165. buffer.WriteString("...")
  166. break
  167. }
  168. buffer.WriteString(strconv.FormatInt(int64(i), 10))
  169. i, e = b.NextSet(i + 1)
  170. if e {
  171. buffer.WriteString(",")
  172. }
  173. }
  174. buffer.WriteString("}")
  175. return buffer.String()
  176. }
  177. // NextSet returns the next bit set from the specified index,
  178. // including possibly the current index
  179. // along with an error code (true = valid, false = no set bit found)
  180. // for i,e := v.NextSet(0); e; i,e = v.NextSet(i + 1) {...}
  181. func (b *BitSet) NextSet(i uint) (uint, bool) {
  182. x := int(i >> log2WordSize)
  183. if x >= len(b.set) {
  184. return 0, false
  185. }
  186. w := b.set[x]
  187. w = w >> (i & (wordSize - 1))
  188. if w != 0 {
  189. return i + trailingZeroes64(w), true
  190. }
  191. x = x + 1
  192. for x < len(b.set) {
  193. if b.set[x] != 0 {
  194. return uint(x)*wordSize + trailingZeroes64(b.set[x]), true
  195. }
  196. x = x + 1
  197. }
  198. return 0, false
  199. }
  200. // NextSetMany returns many next bit sets from the specified index,
  201. // including possibly the current index and up to cap(buffer).
  202. // If the returned slice has len zero, then no more set bits were found
  203. //
  204. // buffer := make([]uint, 256)
  205. // j := uint(0)
  206. // j, buffer = bitmap.NextSetMany(j, buffer)
  207. // for ; len(buffer) > 0; j, buffer = bitmap.NextSetMany(j,buffer) {
  208. // for k := range buffer {
  209. // do something with buffer[k]
  210. // }
  211. // j += 1
  212. // }
  213. //
  214. func (b *BitSet) NextSetMany(i uint, buffer []uint) (uint, []uint) {
  215. myanswer := buffer[:0]
  216. x := int(i >> log2WordSize)
  217. if x >= len(b.set) {
  218. return 0, myanswer
  219. }
  220. w := b.set[x]
  221. w = w >> (i & (wordSize - 1))
  222. base := uint(x << 6)
  223. capacity := cap(buffer)
  224. for len(myanswer) < capacity {
  225. for w != 0 {
  226. t := w & ((^w) + 1)
  227. r := trailingZeroes64(w)
  228. myanswer = append(myanswer, r+base)
  229. if len(myanswer) == capacity {
  230. goto End
  231. }
  232. w = w ^ t
  233. }
  234. x += 1
  235. if x == len(b.set) {
  236. break
  237. }
  238. base += 64
  239. w = b.set[x]
  240. }
  241. End:
  242. if len(myanswer) > 0 {
  243. return myanswer[len(myanswer)-1], myanswer
  244. } else {
  245. return 0, myanswer
  246. }
  247. }
  248. // NextClear returns the next clear bit from the specified index,
  249. // including possibly the current index
  250. // along with an error code (true = valid, false = no bit found i.e. all bits are set)
  251. func (b *BitSet) NextClear(i uint) (uint, bool) {
  252. x := int(i >> log2WordSize)
  253. if x >= len(b.set) {
  254. return 0, false
  255. }
  256. w := b.set[x]
  257. w = w >> (i & (wordSize - 1))
  258. wA := allBits >> (i & (wordSize - 1))
  259. index := i + trailingZeroes64(^w)
  260. if w != wA && index < b.length {
  261. return index, true
  262. }
  263. x++
  264. for x < len(b.set) {
  265. index = uint(x)*wordSize + trailingZeroes64(^b.set[x])
  266. if b.set[x] != allBits && index < b.length {
  267. return index, true
  268. }
  269. x++
  270. }
  271. return 0, false
  272. }
  273. // ClearAll clears the entire BitSet
  274. func (b *BitSet) ClearAll() *BitSet {
  275. if b != nil && b.set != nil {
  276. for i := range b.set {
  277. b.set[i] = 0
  278. }
  279. }
  280. return b
  281. }
  282. // wordCount returns the number of words used in a bit set
  283. func (b *BitSet) wordCount() int {
  284. return len(b.set)
  285. }
  286. // Clone this BitSet
  287. func (b *BitSet) Clone() *BitSet {
  288. c := New(b.length)
  289. if b.set != nil { // Clone should not modify current object
  290. copy(c.set, b.set)
  291. }
  292. return c
  293. }
  294. // Copy into a destination BitSet
  295. // Returning the size of the destination BitSet
  296. // like array copy
  297. func (b *BitSet) Copy(c *BitSet) (count uint) {
  298. if c == nil {
  299. return
  300. }
  301. if b.set != nil { // Copy should not modify current object
  302. copy(c.set, b.set)
  303. }
  304. count = c.length
  305. if b.length < c.length {
  306. count = b.length
  307. }
  308. return
  309. }
  310. // Count (number of set bits)
  311. func (b *BitSet) Count() uint {
  312. if b != nil && b.set != nil {
  313. return uint(popcntSlice(b.set))
  314. }
  315. return 0
  316. }
  317. // Equal tests the equvalence of two BitSets.
  318. // False if they are of different sizes, otherwise true
  319. // only if all the same bits are set
  320. func (b *BitSet) Equal(c *BitSet) bool {
  321. if c == nil {
  322. return false
  323. }
  324. if b.length != c.length {
  325. return false
  326. }
  327. if b.length == 0 { // if they have both length == 0, then could have nil set
  328. return true
  329. }
  330. // testing for equality shoud not transform the bitset (no call to safeSet)
  331. for p, v := range b.set {
  332. if c.set[p] != v {
  333. return false
  334. }
  335. }
  336. return true
  337. }
  338. func panicIfNull(b *BitSet) {
  339. if b == nil {
  340. panic(Error("BitSet must not be null"))
  341. }
  342. }
  343. // Difference of base set and other set
  344. // This is the BitSet equivalent of &^ (and not)
  345. func (b *BitSet) Difference(compare *BitSet) (result *BitSet) {
  346. panicIfNull(b)
  347. panicIfNull(compare)
  348. result = b.Clone() // clone b (in case b is bigger than compare)
  349. l := int(compare.wordCount())
  350. if l > int(b.wordCount()) {
  351. l = int(b.wordCount())
  352. }
  353. for i := 0; i < l; i++ {
  354. result.set[i] = b.set[i] &^ compare.set[i]
  355. }
  356. return
  357. }
  358. // DifferenceCardinality computes the cardinality of the differnce
  359. func (b *BitSet) DifferenceCardinality(compare *BitSet) uint {
  360. panicIfNull(b)
  361. panicIfNull(compare)
  362. l := int(compare.wordCount())
  363. if l > int(b.wordCount()) {
  364. l = int(b.wordCount())
  365. }
  366. cnt := uint64(0)
  367. cnt += popcntMaskSlice(b.set[:l], compare.set[:l])
  368. cnt += popcntSlice(b.set[l:])
  369. return uint(cnt)
  370. }
  371. // InPlaceDifference computes the difference of base set and other set
  372. // This is the BitSet equivalent of &^ (and not)
  373. func (b *BitSet) InPlaceDifference(compare *BitSet) {
  374. panicIfNull(b)
  375. panicIfNull(compare)
  376. l := int(compare.wordCount())
  377. if l > int(b.wordCount()) {
  378. l = int(b.wordCount())
  379. }
  380. for i := 0; i < l; i++ {
  381. b.set[i] &^= compare.set[i]
  382. }
  383. }
  384. // Convenience function: return two bitsets ordered by
  385. // increasing length. Note: neither can be nil
  386. func sortByLength(a *BitSet, b *BitSet) (ap *BitSet, bp *BitSet) {
  387. if a.length <= b.length {
  388. ap, bp = a, b
  389. } else {
  390. ap, bp = b, a
  391. }
  392. return
  393. }
  394. // Intersection of base set and other set
  395. // This is the BitSet equivalent of & (and)
  396. func (b *BitSet) Intersection(compare *BitSet) (result *BitSet) {
  397. panicIfNull(b)
  398. panicIfNull(compare)
  399. b, compare = sortByLength(b, compare)
  400. result = New(b.length)
  401. for i, word := range b.set {
  402. result.set[i] = word & compare.set[i]
  403. }
  404. return
  405. }
  406. // IntersectionCardinality computes the cardinality of the union
  407. func (b *BitSet) IntersectionCardinality(compare *BitSet) uint {
  408. panicIfNull(b)
  409. panicIfNull(compare)
  410. b, compare = sortByLength(b, compare)
  411. cnt := popcntAndSlice(b.set, compare.set)
  412. return uint(cnt)
  413. }
  414. // InPlaceIntersection destructively computes the intersection of
  415. // base set and the compare set.
  416. // This is the BitSet equivalent of & (and)
  417. func (b *BitSet) InPlaceIntersection(compare *BitSet) {
  418. panicIfNull(b)
  419. panicIfNull(compare)
  420. l := int(compare.wordCount())
  421. if l > int(b.wordCount()) {
  422. l = int(b.wordCount())
  423. }
  424. for i := 0; i < l; i++ {
  425. b.set[i] &= compare.set[i]
  426. }
  427. for i := l; i < len(b.set); i++ {
  428. b.set[i] = 0
  429. }
  430. if compare.length > 0 {
  431. b.extendSetMaybe(compare.length - 1)
  432. }
  433. }
  434. // Union of base set and other set
  435. // This is the BitSet equivalent of | (or)
  436. func (b *BitSet) Union(compare *BitSet) (result *BitSet) {
  437. panicIfNull(b)
  438. panicIfNull(compare)
  439. b, compare = sortByLength(b, compare)
  440. result = compare.Clone()
  441. for i, word := range b.set {
  442. result.set[i] = word | compare.set[i]
  443. }
  444. return
  445. }
  446. // UnionCardinality computes the cardinality of the uniton of the base set
  447. // and the compare set.
  448. func (b *BitSet) UnionCardinality(compare *BitSet) uint {
  449. panicIfNull(b)
  450. panicIfNull(compare)
  451. b, compare = sortByLength(b, compare)
  452. cnt := popcntOrSlice(b.set, compare.set)
  453. if len(compare.set) > len(b.set) {
  454. cnt += popcntSlice(compare.set[len(b.set):])
  455. }
  456. return uint(cnt)
  457. }
  458. // InPlaceUnion creates the destructive union of base set and compare set.
  459. // This is the BitSet equivalent of | (or).
  460. func (b *BitSet) InPlaceUnion(compare *BitSet) {
  461. panicIfNull(b)
  462. panicIfNull(compare)
  463. l := int(compare.wordCount())
  464. if l > int(b.wordCount()) {
  465. l = int(b.wordCount())
  466. }
  467. if compare.length > 0 {
  468. b.extendSetMaybe(compare.length - 1)
  469. }
  470. for i := 0; i < l; i++ {
  471. b.set[i] |= compare.set[i]
  472. }
  473. if len(compare.set) > l {
  474. for i := l; i < len(compare.set); i++ {
  475. b.set[i] = compare.set[i]
  476. }
  477. }
  478. }
  479. // SymmetricDifference of base set and other set
  480. // This is the BitSet equivalent of ^ (xor)
  481. func (b *BitSet) SymmetricDifference(compare *BitSet) (result *BitSet) {
  482. panicIfNull(b)
  483. panicIfNull(compare)
  484. b, compare = sortByLength(b, compare)
  485. // compare is bigger, so clone it
  486. result = compare.Clone()
  487. for i, word := range b.set {
  488. result.set[i] = word ^ compare.set[i]
  489. }
  490. return
  491. }
  492. // SymmetricDifferenceCardinality computes the cardinality of the symmetric difference
  493. func (b *BitSet) SymmetricDifferenceCardinality(compare *BitSet) uint {
  494. panicIfNull(b)
  495. panicIfNull(compare)
  496. b, compare = sortByLength(b, compare)
  497. cnt := popcntXorSlice(b.set, compare.set)
  498. if len(compare.set) > len(b.set) {
  499. cnt += popcntSlice(compare.set[len(b.set):])
  500. }
  501. return uint(cnt)
  502. }
  503. // InPlaceSymmetricDifference creates the destructive SymmetricDifference of base set and other set
  504. // This is the BitSet equivalent of ^ (xor)
  505. func (b *BitSet) InPlaceSymmetricDifference(compare *BitSet) {
  506. panicIfNull(b)
  507. panicIfNull(compare)
  508. l := int(compare.wordCount())
  509. if l > int(b.wordCount()) {
  510. l = int(b.wordCount())
  511. }
  512. if compare.length > 0 {
  513. b.extendSetMaybe(compare.length - 1)
  514. }
  515. for i := 0; i < l; i++ {
  516. b.set[i] ^= compare.set[i]
  517. }
  518. if len(compare.set) > l {
  519. for i := l; i < len(compare.set); i++ {
  520. b.set[i] = compare.set[i]
  521. }
  522. }
  523. }
  524. // Is the length an exact multiple of word sizes?
  525. func (b *BitSet) isLenExactMultiple() bool {
  526. return b.length%wordSize == 0
  527. }
  528. // Clean last word by setting unused bits to 0
  529. func (b *BitSet) cleanLastWord() {
  530. if !b.isLenExactMultiple() {
  531. b.set[len(b.set)-1] &= allBits >> (wordSize - b.length%wordSize)
  532. }
  533. }
  534. // Complement computes the (local) complement of a biset (up to length bits)
  535. func (b *BitSet) Complement() (result *BitSet) {
  536. panicIfNull(b)
  537. result = New(b.length)
  538. for i, word := range b.set {
  539. result.set[i] = ^word
  540. }
  541. result.cleanLastWord()
  542. return
  543. }
  544. // All returns true if all bits are set, false otherwise. Returns true for
  545. // empty sets.
  546. func (b *BitSet) All() bool {
  547. panicIfNull(b)
  548. return b.Count() == b.length
  549. }
  550. // None returns true if no bit is set, false otherwise. Retursn true for
  551. // empty sets.
  552. func (b *BitSet) None() bool {
  553. panicIfNull(b)
  554. if b != nil && b.set != nil {
  555. for _, word := range b.set {
  556. if word > 0 {
  557. return false
  558. }
  559. }
  560. return true
  561. }
  562. return true
  563. }
  564. // Any returns true if any bit is set, false otherwise
  565. func (b *BitSet) Any() bool {
  566. panicIfNull(b)
  567. return !b.None()
  568. }
  569. // IsSuperSet returns true if this is a superset of the other set
  570. func (b *BitSet) IsSuperSet(other *BitSet) bool {
  571. for i, e := other.NextSet(0); e; i, e = other.NextSet(i + 1) {
  572. if !b.Test(i) {
  573. return false
  574. }
  575. }
  576. return true
  577. }
  578. // IsStrictSuperSet returns true if this is a strict superset of the other set
  579. func (b *BitSet) IsStrictSuperSet(other *BitSet) bool {
  580. return b.Count() > other.Count() && b.IsSuperSet(other)
  581. }
  582. // DumpAsBits dumps a bit set as a string of bits
  583. func (b *BitSet) DumpAsBits() string {
  584. if b.set == nil {
  585. return "."
  586. }
  587. buffer := bytes.NewBufferString("")
  588. i := len(b.set) - 1
  589. for ; i >= 0; i-- {
  590. fmt.Fprintf(buffer, "%064b.", b.set[i])
  591. }
  592. return string(buffer.Bytes())
  593. }
  594. // BinaryStorageSize returns the binary storage requirements
  595. func (b *BitSet) BinaryStorageSize() int {
  596. return binary.Size(uint64(0)) + binary.Size(b.set)
  597. }
  598. // WriteTo writes a BitSet to a stream
  599. func (b *BitSet) WriteTo(stream io.Writer) (int64, error) {
  600. length := uint64(b.length)
  601. // Write length
  602. err := binary.Write(stream, binary.BigEndian, length)
  603. if err != nil {
  604. return 0, err
  605. }
  606. // Write set
  607. err = binary.Write(stream, binary.BigEndian, b.set)
  608. return int64(b.BinaryStorageSize()), err
  609. }
  610. // ReadFrom reads a BitSet from a stream written using WriteTo
  611. func (b *BitSet) ReadFrom(stream io.Reader) (int64, error) {
  612. var length uint64
  613. // Read length first
  614. err := binary.Read(stream, binary.BigEndian, &length)
  615. if err != nil {
  616. return 0, err
  617. }
  618. newset := New(uint(length))
  619. if uint64(newset.length) != length {
  620. return 0, errors.New("Unmarshalling error: type mismatch")
  621. }
  622. // Read remaining bytes as set
  623. err = binary.Read(stream, binary.BigEndian, newset.set)
  624. if err != nil {
  625. return 0, err
  626. }
  627. *b = *newset
  628. return int64(b.BinaryStorageSize()), nil
  629. }
  630. // MarshalBinary encodes a BitSet into a binary form and returns the result.
  631. func (b *BitSet) MarshalBinary() ([]byte, error) {
  632. var buf bytes.Buffer
  633. writer := bufio.NewWriter(&buf)
  634. _, err := b.WriteTo(writer)
  635. if err != nil {
  636. return []byte{}, err
  637. }
  638. err = writer.Flush()
  639. return buf.Bytes(), err
  640. }
  641. // UnmarshalBinary decodes the binary form generated by MarshalBinary.
  642. func (b *BitSet) UnmarshalBinary(data []byte) error {
  643. buf := bytes.NewReader(data)
  644. reader := bufio.NewReader(buf)
  645. _, err := b.ReadFrom(reader)
  646. return err
  647. }
  648. // MarshalJSON marshals a BitSet as a JSON structure
  649. func (b *BitSet) MarshalJSON() ([]byte, error) {
  650. buffer := bytes.NewBuffer(make([]byte, 0, b.BinaryStorageSize()))
  651. _, err := b.WriteTo(buffer)
  652. if err != nil {
  653. return nil, err
  654. }
  655. // URLEncode all bytes
  656. return json.Marshal(base64.URLEncoding.EncodeToString(buffer.Bytes()))
  657. }
  658. // UnmarshalJSON unmarshals a BitSet from JSON created using MarshalJSON
  659. func (b *BitSet) UnmarshalJSON(data []byte) error {
  660. // Unmarshal as string
  661. var s string
  662. err := json.Unmarshal(data, &s)
  663. if err != nil {
  664. return err
  665. }
  666. // URLDecode string
  667. buf, err := base64.URLEncoding.DecodeString(s)
  668. if err != nil {
  669. return err
  670. }
  671. _, err = b.ReadFrom(bytes.NewReader(buf))
  672. return err
  673. }