bitwriter.go 5.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210
  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 huff0
  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. // encSymbol will add up to 16 bits. value may not contain more set bits than indicated.
  35. // It will not check if there is space for them, so the caller must ensure that it has flushed recently.
  36. func (b *bitWriter) encSymbol(ct cTable, symbol byte) {
  37. enc := ct[symbol]
  38. b.bitContainer |= uint64(enc.val) << (b.nBits & 63)
  39. if false {
  40. if enc.nBits == 0 {
  41. panic("nbits 0")
  42. }
  43. }
  44. b.nBits += enc.nBits
  45. }
  46. // encTwoSymbols will add up to 32 bits. value may not contain more set bits than indicated.
  47. // It will not check if there is space for them, so the caller must ensure that it has flushed recently.
  48. func (b *bitWriter) encTwoSymbols(ct cTable, av, bv byte) {
  49. encA := ct[av]
  50. encB := ct[bv]
  51. sh := b.nBits & 63
  52. combined := uint64(encA.val) | (uint64(encB.val) << (encA.nBits & 63))
  53. b.bitContainer |= combined << sh
  54. if false {
  55. if encA.nBits == 0 {
  56. panic("nbitsA 0")
  57. }
  58. if encB.nBits == 0 {
  59. panic("nbitsB 0")
  60. }
  61. }
  62. b.nBits += encA.nBits + encB.nBits
  63. }
  64. // addBits16ZeroNC will add up to 16 bits.
  65. // It will not check if there is space for them,
  66. // so the caller must ensure that it has flushed recently.
  67. // This is fastest if bits can be zero.
  68. func (b *bitWriter) addBits16ZeroNC(value uint16, bits uint8) {
  69. if bits == 0 {
  70. return
  71. }
  72. value <<= (16 - bits) & 15
  73. value >>= (16 - bits) & 15
  74. b.bitContainer |= uint64(value) << (b.nBits & 63)
  75. b.nBits += bits
  76. }
  77. // flush will flush all pending full bytes.
  78. // There will be at least 56 bits available for writing when this has been called.
  79. // Using flush32 is faster, but leaves less space for writing.
  80. func (b *bitWriter) flush() {
  81. v := b.nBits >> 3
  82. switch v {
  83. case 0:
  84. return
  85. case 1:
  86. b.out = append(b.out,
  87. byte(b.bitContainer),
  88. )
  89. b.bitContainer >>= 1 << 3
  90. case 2:
  91. b.out = append(b.out,
  92. byte(b.bitContainer),
  93. byte(b.bitContainer>>8),
  94. )
  95. b.bitContainer >>= 2 << 3
  96. case 3:
  97. b.out = append(b.out,
  98. byte(b.bitContainer),
  99. byte(b.bitContainer>>8),
  100. byte(b.bitContainer>>16),
  101. )
  102. b.bitContainer >>= 3 << 3
  103. case 4:
  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. )
  110. b.bitContainer >>= 4 << 3
  111. case 5:
  112. b.out = append(b.out,
  113. byte(b.bitContainer),
  114. byte(b.bitContainer>>8),
  115. byte(b.bitContainer>>16),
  116. byte(b.bitContainer>>24),
  117. byte(b.bitContainer>>32),
  118. )
  119. b.bitContainer >>= 5 << 3
  120. case 6:
  121. b.out = append(b.out,
  122. byte(b.bitContainer),
  123. byte(b.bitContainer>>8),
  124. byte(b.bitContainer>>16),
  125. byte(b.bitContainer>>24),
  126. byte(b.bitContainer>>32),
  127. byte(b.bitContainer>>40),
  128. )
  129. b.bitContainer >>= 6 << 3
  130. case 7:
  131. b.out = append(b.out,
  132. byte(b.bitContainer),
  133. byte(b.bitContainer>>8),
  134. byte(b.bitContainer>>16),
  135. byte(b.bitContainer>>24),
  136. byte(b.bitContainer>>32),
  137. byte(b.bitContainer>>40),
  138. byte(b.bitContainer>>48),
  139. )
  140. b.bitContainer >>= 7 << 3
  141. case 8:
  142. b.out = append(b.out,
  143. byte(b.bitContainer),
  144. byte(b.bitContainer>>8),
  145. byte(b.bitContainer>>16),
  146. byte(b.bitContainer>>24),
  147. byte(b.bitContainer>>32),
  148. byte(b.bitContainer>>40),
  149. byte(b.bitContainer>>48),
  150. byte(b.bitContainer>>56),
  151. )
  152. b.bitContainer = 0
  153. b.nBits = 0
  154. return
  155. default:
  156. panic(fmt.Errorf("bits (%d) > 64", b.nBits))
  157. }
  158. b.nBits &= 7
  159. }
  160. // flush32 will flush out, so there are at least 32 bits available for writing.
  161. func (b *bitWriter) flush32() {
  162. if b.nBits < 32 {
  163. return
  164. }
  165. b.out = append(b.out,
  166. byte(b.bitContainer),
  167. byte(b.bitContainer>>8),
  168. byte(b.bitContainer>>16),
  169. byte(b.bitContainer>>24))
  170. b.nBits -= 32
  171. b.bitContainer >>= 32
  172. }
  173. // flushAlign will flush remaining full bytes and align to next byte boundary.
  174. func (b *bitWriter) flushAlign() {
  175. nbBytes := (b.nBits + 7) >> 3
  176. for i := uint8(0); i < nbBytes; i++ {
  177. b.out = append(b.out, byte(b.bitContainer>>(i*8)))
  178. }
  179. b.nBits = 0
  180. b.bitContainer = 0
  181. }
  182. // close will write the alignment bit and write the final byte(s)
  183. // to the output.
  184. func (b *bitWriter) close() error {
  185. // End mark
  186. b.addBits16Clean(1, 1)
  187. // flush until next byte.
  188. b.flushAlign()
  189. return nil
  190. }
  191. // reset and continue writing by appending to out.
  192. func (b *bitWriter) reset(out []byte) {
  193. b.bitContainer = 0
  194. b.nBits = 0
  195. b.out = out
  196. }