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

141 рядки
3.4 KiB

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