bitreader.go 3.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136
  1. // Copyright 2019+ Klaus Post. All rights reserved.
  2. // License information can be found in the LICENSE file.
  3. // Based on work by Yann Collet, released under BSD License.
  4. package zstd
  5. import (
  6. "encoding/binary"
  7. "errors"
  8. "fmt"
  9. "io"
  10. "math/bits"
  11. )
  12. // bitReader reads a bitstream in reverse.
  13. // The last set bit indicates the start of the stream and is used
  14. // for aligning the input.
  15. type bitReader struct {
  16. in []byte
  17. value uint64 // Maybe use [16]byte, but shifting is awkward.
  18. bitsRead uint8
  19. }
  20. // init initializes and resets the bit reader.
  21. func (b *bitReader) init(in []byte) error {
  22. if len(in) < 1 {
  23. return errors.New("corrupt stream: too short")
  24. }
  25. b.in = in
  26. // The highest bit of the last byte indicates where to start
  27. v := in[len(in)-1]
  28. if v == 0 {
  29. return errors.New("corrupt stream, did not find end of stream")
  30. }
  31. b.bitsRead = 64
  32. b.value = 0
  33. if len(in) >= 8 {
  34. b.fillFastStart()
  35. } else {
  36. b.fill()
  37. b.fill()
  38. }
  39. b.bitsRead += 8 - uint8(highBits(uint32(v)))
  40. return nil
  41. }
  42. // getBits will return n bits. n can be 0.
  43. func (b *bitReader) getBits(n uint8) int {
  44. if n == 0 /*|| b.bitsRead >= 64 */ {
  45. return 0
  46. }
  47. return int(b.get32BitsFast(n))
  48. }
  49. // get32BitsFast requires that at least one bit is requested every time.
  50. // There are no checks if the buffer is filled.
  51. func (b *bitReader) get32BitsFast(n uint8) uint32 {
  52. const regMask = 64 - 1
  53. v := uint32((b.value << (b.bitsRead & regMask)) >> ((regMask + 1 - n) & regMask))
  54. b.bitsRead += n
  55. return v
  56. }
  57. // fillFast() will make sure at least 32 bits are available.
  58. // There must be at least 4 bytes available.
  59. func (b *bitReader) fillFast() {
  60. if b.bitsRead < 32 {
  61. return
  62. }
  63. v := b.in[len(b.in)-4:]
  64. b.in = b.in[:len(b.in)-4]
  65. low := (uint32(v[0])) | (uint32(v[1]) << 8) | (uint32(v[2]) << 16) | (uint32(v[3]) << 24)
  66. b.value = (b.value << 32) | uint64(low)
  67. b.bitsRead -= 32
  68. }
  69. // fillFastStart() assumes the bitreader is empty and there is at least 8 bytes to read.
  70. func (b *bitReader) fillFastStart() {
  71. v := b.in[len(b.in)-8:]
  72. b.in = b.in[:len(b.in)-8]
  73. b.value = binary.LittleEndian.Uint64(v)
  74. b.bitsRead = 0
  75. }
  76. // fill() will make sure at least 32 bits are available.
  77. func (b *bitReader) fill() {
  78. if b.bitsRead < 32 {
  79. return
  80. }
  81. if len(b.in) >= 4 {
  82. v := b.in[len(b.in)-4:]
  83. b.in = b.in[:len(b.in)-4]
  84. low := (uint32(v[0])) | (uint32(v[1]) << 8) | (uint32(v[2]) << 16) | (uint32(v[3]) << 24)
  85. b.value = (b.value << 32) | uint64(low)
  86. b.bitsRead -= 32
  87. return
  88. }
  89. b.bitsRead -= uint8(8 * len(b.in))
  90. for len(b.in) > 0 {
  91. b.value = (b.value << 8) | uint64(b.in[len(b.in)-1])
  92. b.in = b.in[:len(b.in)-1]
  93. }
  94. }
  95. // finished returns true if all bits have been read from the bit stream.
  96. func (b *bitReader) finished() bool {
  97. return len(b.in) == 0 && b.bitsRead >= 64
  98. }
  99. // overread returns true if more bits have been requested than is on the stream.
  100. func (b *bitReader) overread() bool {
  101. return b.bitsRead > 64
  102. }
  103. // remain returns the number of bits remaining.
  104. func (b *bitReader) remain() uint {
  105. return 8*uint(len(b.in)) + 64 - uint(b.bitsRead)
  106. }
  107. // close the bitstream and returns an error if out-of-buffer reads occurred.
  108. func (b *bitReader) close() error {
  109. // Release reference.
  110. b.in = nil
  111. if !b.finished() {
  112. return fmt.Errorf("%d extra bits on block, should be 0", b.remain())
  113. }
  114. if b.bitsRead > 64 {
  115. return io.ErrUnexpectedEOF
  116. }
  117. return nil
  118. }
  119. func highBits(val uint32) (n uint32) {
  120. return uint32(bits.Len32(val) - 1)
  121. }