bitreader.go 3.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122
  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 (
  7. "encoding/binary"
  8. "errors"
  9. "io"
  10. )
  11. // bitReader reads a bitstream in reverse.
  12. // The last set bit indicates the start of the stream and is used
  13. // for aligning the input.
  14. type bitReader struct {
  15. in []byte
  16. off uint // next byte to read is at in[off - 1]
  17. value uint64
  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. b.off = uint(len(in))
  27. // The highest bit of the last byte indicates where to start
  28. v := in[len(in)-1]
  29. if v == 0 {
  30. return errors.New("corrupt stream, did not find end of stream")
  31. }
  32. b.bitsRead = 64
  33. b.value = 0
  34. if len(in) >= 8 {
  35. b.fillFastStart()
  36. } else {
  37. b.fill()
  38. b.fill()
  39. }
  40. b.bitsRead += 8 - uint8(highBits(uint32(v)))
  41. return nil
  42. }
  43. // getBits will return n bits. n can be 0.
  44. func (b *bitReader) getBits(n uint8) uint16 {
  45. if n == 0 || b.bitsRead >= 64 {
  46. return 0
  47. }
  48. return b.getBitsFast(n)
  49. }
  50. // getBitsFast requires that at least one bit is requested every time.
  51. // There are no checks if the buffer is filled.
  52. func (b *bitReader) getBitsFast(n uint8) uint16 {
  53. const regMask = 64 - 1
  54. v := uint16((b.value << (b.bitsRead & regMask)) >> ((regMask + 1 - n) & regMask))
  55. b.bitsRead += n
  56. return v
  57. }
  58. // fillFast() will make sure at least 32 bits are available.
  59. // There must be at least 4 bytes available.
  60. func (b *bitReader) fillFast() {
  61. if b.bitsRead < 32 {
  62. return
  63. }
  64. // 2 bounds checks.
  65. v := b.in[b.off-4:]
  66. v = v[:4]
  67. low := (uint32(v[0])) | (uint32(v[1]) << 8) | (uint32(v[2]) << 16) | (uint32(v[3]) << 24)
  68. b.value = (b.value << 32) | uint64(low)
  69. b.bitsRead -= 32
  70. b.off -= 4
  71. }
  72. // fill() will make sure at least 32 bits are available.
  73. func (b *bitReader) fill() {
  74. if b.bitsRead < 32 {
  75. return
  76. }
  77. if b.off > 4 {
  78. v := b.in[b.off-4:]
  79. v = v[:4]
  80. low := (uint32(v[0])) | (uint32(v[1]) << 8) | (uint32(v[2]) << 16) | (uint32(v[3]) << 24)
  81. b.value = (b.value << 32) | uint64(low)
  82. b.bitsRead -= 32
  83. b.off -= 4
  84. return
  85. }
  86. for b.off > 0 {
  87. b.value = (b.value << 8) | uint64(b.in[b.off-1])
  88. b.bitsRead -= 8
  89. b.off--
  90. }
  91. }
  92. // fillFastStart() assumes the bitreader is empty and there is at least 8 bytes to read.
  93. func (b *bitReader) fillFastStart() {
  94. // Do single re-slice to avoid bounds checks.
  95. b.value = binary.LittleEndian.Uint64(b.in[b.off-8:])
  96. b.bitsRead = 0
  97. b.off -= 8
  98. }
  99. // finished returns true if all bits have been read from the bit stream.
  100. func (b *bitReader) finished() bool {
  101. return b.bitsRead >= 64 && b.off == 0
  102. }
  103. // close the bitstream and returns an error if out-of-buffer reads occurred.
  104. func (b *bitReader) close() error {
  105. // Release reference.
  106. b.in = nil
  107. if b.bitsRead > 64 {
  108. return io.ErrUnexpectedEOF
  109. }
  110. return nil
  111. }