bitwriter.go 3.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112
  1. // Copyright 2018 Klaus Post. All rights reserved.
  2. // Use of this source code is governed by a BSD-style
  3. // license that can be found in the LICENSE file.
  4. // Based on work Copyright (c) 2013, Yann Collet, released under BSD License.
  5. package zstd
  6. // bitWriter will write bits.
  7. // First bit will be LSB of the first byte of output.
  8. type bitWriter struct {
  9. bitContainer uint64
  10. nBits uint8
  11. out []byte
  12. }
  13. // bitMask16 is bitmasks. Has extra to avoid bounds check.
  14. var bitMask16 = [32]uint16{
  15. 0, 1, 3, 7, 0xF, 0x1F,
  16. 0x3F, 0x7F, 0xFF, 0x1FF, 0x3FF, 0x7FF,
  17. 0xFFF, 0x1FFF, 0x3FFF, 0x7FFF, 0xFFFF, 0xFFFF,
  18. 0xFFFF, 0xFFFF, 0xFFFF, 0xFFFF, 0xFFFF, 0xFFFF,
  19. 0xFFFF, 0xFFFF} /* up to 16 bits */
  20. var bitMask32 = [32]uint32{
  21. 0, 1, 3, 7, 0xF, 0x1F, 0x3F, 0x7F, 0xFF,
  22. 0x1FF, 0x3FF, 0x7FF, 0xFFF, 0x1FFF, 0x3FFF, 0x7FFF, 0xFFFF,
  23. 0x1ffff, 0x3ffff, 0x7FFFF, 0xfFFFF, 0x1fFFFF, 0x3fFFFF, 0x7fFFFF, 0xffFFFF,
  24. 0x1ffFFFF, 0x3ffFFFF, 0x7ffFFFF, 0xfffFFFF, 0x1fffFFFF, 0x3fffFFFF, 0x7fffFFFF,
  25. } // up to 32 bits
  26. // addBits16NC will add up to 16 bits.
  27. // It will not check if there is space for them,
  28. // so the caller must ensure that it has flushed recently.
  29. func (b *bitWriter) addBits16NC(value uint16, bits uint8) {
  30. b.bitContainer |= uint64(value&bitMask16[bits&31]) << (b.nBits & 63)
  31. b.nBits += bits
  32. }
  33. // addBits32NC will add up to 31 bits.
  34. // It will not check if there is space for them,
  35. // so the caller must ensure that it has flushed recently.
  36. func (b *bitWriter) addBits32NC(value uint32, bits uint8) {
  37. b.bitContainer |= uint64(value&bitMask32[bits&31]) << (b.nBits & 63)
  38. b.nBits += bits
  39. }
  40. // addBits64NC will add up to 64 bits.
  41. // There must be space for 32 bits.
  42. func (b *bitWriter) addBits64NC(value uint64, bits uint8) {
  43. if bits <= 31 {
  44. b.addBits32Clean(uint32(value), bits)
  45. return
  46. }
  47. b.addBits32Clean(uint32(value), 32)
  48. b.flush32()
  49. b.addBits32Clean(uint32(value>>32), bits-32)
  50. }
  51. // addBits32Clean will add up to 32 bits.
  52. // It will not check if there is space for them.
  53. // The input must not contain more bits than specified.
  54. func (b *bitWriter) addBits32Clean(value uint32, bits uint8) {
  55. b.bitContainer |= uint64(value) << (b.nBits & 63)
  56. b.nBits += bits
  57. }
  58. // addBits16Clean will add up to 16 bits. value may not contain more set bits than indicated.
  59. // It will not check if there is space for them, so the caller must ensure that it has flushed recently.
  60. func (b *bitWriter) addBits16Clean(value uint16, bits uint8) {
  61. b.bitContainer |= uint64(value) << (b.nBits & 63)
  62. b.nBits += bits
  63. }
  64. // flush32 will flush out, so there are at least 32 bits available for writing.
  65. func (b *bitWriter) flush32() {
  66. if b.nBits < 32 {
  67. return
  68. }
  69. b.out = append(b.out,
  70. byte(b.bitContainer),
  71. byte(b.bitContainer>>8),
  72. byte(b.bitContainer>>16),
  73. byte(b.bitContainer>>24))
  74. b.nBits -= 32
  75. b.bitContainer >>= 32
  76. }
  77. // flushAlign will flush remaining full bytes and align to next byte boundary.
  78. func (b *bitWriter) flushAlign() {
  79. nbBytes := (b.nBits + 7) >> 3
  80. for i := uint8(0); i < nbBytes; i++ {
  81. b.out = append(b.out, byte(b.bitContainer>>(i*8)))
  82. }
  83. b.nBits = 0
  84. b.bitContainer = 0
  85. }
  86. // close will write the alignment bit and write the final byte(s)
  87. // to the output.
  88. func (b *bitWriter) close() {
  89. // End mark
  90. b.addBits16Clean(1, 1)
  91. // flush until next byte.
  92. b.flushAlign()
  93. }
  94. // reset and continue writing by appending to out.
  95. func (b *bitWriter) reset(out []byte) {
  96. b.bitContainer = 0
  97. b.nBits = 0
  98. b.out = out
  99. }