bitwriter.go 4.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167
  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 fse
  6. import "fmt"
  7. // bitWriter will write bits.
  8. // First bit will be LSB of the first byte of output.
  9. type bitWriter struct {
  10. bitContainer uint64
  11. nBits uint8
  12. out []byte
  13. }
  14. // bitMask16 is bitmasks. Has extra to avoid bounds check.
  15. var bitMask16 = [32]uint16{
  16. 0, 1, 3, 7, 0xF, 0x1F,
  17. 0x3F, 0x7F, 0xFF, 0x1FF, 0x3FF, 0x7FF,
  18. 0xFFF, 0x1FFF, 0x3FFF, 0x7FFF, 0xFFFF, 0xFFFF,
  19. 0xFFFF, 0xFFFF, 0xFFFF, 0xFFFF, 0xFFFF, 0xFFFF,
  20. 0xFFFF, 0xFFFF} /* up to 16 bits */
  21. // addBits16NC will add up to 16 bits.
  22. // It will not check if there is space for them,
  23. // so the caller must ensure that it has flushed recently.
  24. func (b *bitWriter) addBits16NC(value uint16, bits uint8) {
  25. b.bitContainer |= uint64(value&bitMask16[bits&31]) << (b.nBits & 63)
  26. b.nBits += bits
  27. }
  28. // addBits16Clean will add up to 16 bits. value may not contain more set bits than indicated.
  29. // It will not check if there is space for them, so the caller must ensure that it has flushed recently.
  30. func (b *bitWriter) addBits16Clean(value uint16, bits uint8) {
  31. b.bitContainer |= uint64(value) << (b.nBits & 63)
  32. b.nBits += bits
  33. }
  34. // addBits16ZeroNC will add up to 16 bits.
  35. // It will not check if there is space for them,
  36. // so the caller must ensure that it has flushed recently.
  37. // This is fastest if bits can be zero.
  38. func (b *bitWriter) addBits16ZeroNC(value uint16, bits uint8) {
  39. if bits == 0 {
  40. return
  41. }
  42. value <<= (16 - bits) & 15
  43. value >>= (16 - bits) & 15
  44. b.bitContainer |= uint64(value) << (b.nBits & 63)
  45. b.nBits += bits
  46. }
  47. // flush will flush all pending full bytes.
  48. // There will be at least 56 bits available for writing when this has been called.
  49. // Using flush32 is faster, but leaves less space for writing.
  50. func (b *bitWriter) flush() {
  51. v := b.nBits >> 3
  52. switch v {
  53. case 0:
  54. case 1:
  55. b.out = append(b.out,
  56. byte(b.bitContainer),
  57. )
  58. case 2:
  59. b.out = append(b.out,
  60. byte(b.bitContainer),
  61. byte(b.bitContainer>>8),
  62. )
  63. case 3:
  64. b.out = append(b.out,
  65. byte(b.bitContainer),
  66. byte(b.bitContainer>>8),
  67. byte(b.bitContainer>>16),
  68. )
  69. case 4:
  70. b.out = append(b.out,
  71. byte(b.bitContainer),
  72. byte(b.bitContainer>>8),
  73. byte(b.bitContainer>>16),
  74. byte(b.bitContainer>>24),
  75. )
  76. case 5:
  77. b.out = append(b.out,
  78. byte(b.bitContainer),
  79. byte(b.bitContainer>>8),
  80. byte(b.bitContainer>>16),
  81. byte(b.bitContainer>>24),
  82. byte(b.bitContainer>>32),
  83. )
  84. case 6:
  85. b.out = append(b.out,
  86. byte(b.bitContainer),
  87. byte(b.bitContainer>>8),
  88. byte(b.bitContainer>>16),
  89. byte(b.bitContainer>>24),
  90. byte(b.bitContainer>>32),
  91. byte(b.bitContainer>>40),
  92. )
  93. case 7:
  94. b.out = append(b.out,
  95. byte(b.bitContainer),
  96. byte(b.bitContainer>>8),
  97. byte(b.bitContainer>>16),
  98. byte(b.bitContainer>>24),
  99. byte(b.bitContainer>>32),
  100. byte(b.bitContainer>>40),
  101. byte(b.bitContainer>>48),
  102. )
  103. case 8:
  104. b.out = append(b.out,
  105. byte(b.bitContainer),
  106. byte(b.bitContainer>>8),
  107. byte(b.bitContainer>>16),
  108. byte(b.bitContainer>>24),
  109. byte(b.bitContainer>>32),
  110. byte(b.bitContainer>>40),
  111. byte(b.bitContainer>>48),
  112. byte(b.bitContainer>>56),
  113. )
  114. default:
  115. panic(fmt.Errorf("bits (%d) > 64", b.nBits))
  116. }
  117. b.bitContainer >>= v << 3
  118. b.nBits &= 7
  119. }
  120. // flush32 will flush out, so there are at least 32 bits available for writing.
  121. func (b *bitWriter) flush32() {
  122. if b.nBits < 32 {
  123. return
  124. }
  125. b.out = append(b.out,
  126. byte(b.bitContainer),
  127. byte(b.bitContainer>>8),
  128. byte(b.bitContainer>>16),
  129. byte(b.bitContainer>>24))
  130. b.nBits -= 32
  131. b.bitContainer >>= 32
  132. }
  133. // flushAlign will flush remaining full bytes and align to next byte boundary.
  134. func (b *bitWriter) flushAlign() {
  135. nbBytes := (b.nBits + 7) >> 3
  136. for i := uint8(0); i < nbBytes; i++ {
  137. b.out = append(b.out, byte(b.bitContainer>>(i*8)))
  138. }
  139. b.nBits = 0
  140. b.bitContainer = 0
  141. }
  142. // close will write the alignment bit and write the final byte(s)
  143. // to the output.
  144. func (b *bitWriter) close() {
  145. // End mark
  146. b.addBits16Clean(1, 1)
  147. // flush until next byte.
  148. b.flushAlign()
  149. }
  150. // reset and continue writing by appending to out.
  151. func (b *bitWriter) reset(out []byte) {
  152. b.bitContainer = 0
  153. b.nBits = 0
  154. b.out = out
  155. }