rangecodec.go 5.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222
  1. // Copyright 2014-2022 Ulrich Kunitz. 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. package lzma
  5. import (
  6. "errors"
  7. "io"
  8. )
  9. // rangeEncoder implements range encoding of single bits. The low value can
  10. // overflow therefore we need uint64. The cache value is used to handle
  11. // overflows.
  12. type rangeEncoder struct {
  13. lbw *LimitedByteWriter
  14. nrange uint32
  15. low uint64
  16. cacheLen int64
  17. cache byte
  18. }
  19. // maxInt64 provides the maximal value of the int64 type
  20. const maxInt64 = 1<<63 - 1
  21. // newRangeEncoder creates a new range encoder.
  22. func newRangeEncoder(bw io.ByteWriter) (re *rangeEncoder, err error) {
  23. lbw, ok := bw.(*LimitedByteWriter)
  24. if !ok {
  25. lbw = &LimitedByteWriter{BW: bw, N: maxInt64}
  26. }
  27. return &rangeEncoder{
  28. lbw: lbw,
  29. nrange: 0xffffffff,
  30. cacheLen: 1}, nil
  31. }
  32. // Available returns the number of bytes that still can be written. The
  33. // method takes the bytes that will be currently written by Close into
  34. // account.
  35. func (e *rangeEncoder) Available() int64 {
  36. return e.lbw.N - (e.cacheLen + 4)
  37. }
  38. // writeByte writes a single byte to the underlying writer. An error is
  39. // returned if the limit is reached. The written byte will be counted if
  40. // the underlying writer doesn't return an error.
  41. func (e *rangeEncoder) writeByte(c byte) error {
  42. if e.Available() < 1 {
  43. return ErrLimit
  44. }
  45. return e.lbw.WriteByte(c)
  46. }
  47. // DirectEncodeBit encodes the least-significant bit of b with probability 1/2.
  48. func (e *rangeEncoder) DirectEncodeBit(b uint32) error {
  49. e.nrange >>= 1
  50. e.low += uint64(e.nrange) & (0 - (uint64(b) & 1))
  51. // normalize
  52. const top = 1 << 24
  53. if e.nrange >= top {
  54. return nil
  55. }
  56. e.nrange <<= 8
  57. return e.shiftLow()
  58. }
  59. // EncodeBit encodes the least significant bit of b. The p value will be
  60. // updated by the function depending on the bit encoded.
  61. func (e *rangeEncoder) EncodeBit(b uint32, p *prob) error {
  62. bound := p.bound(e.nrange)
  63. if b&1 == 0 {
  64. e.nrange = bound
  65. p.inc()
  66. } else {
  67. e.low += uint64(bound)
  68. e.nrange -= bound
  69. p.dec()
  70. }
  71. // normalize
  72. const top = 1 << 24
  73. if e.nrange >= top {
  74. return nil
  75. }
  76. e.nrange <<= 8
  77. return e.shiftLow()
  78. }
  79. // Close writes a complete copy of the low value.
  80. func (e *rangeEncoder) Close() error {
  81. for i := 0; i < 5; i++ {
  82. if err := e.shiftLow(); err != nil {
  83. return err
  84. }
  85. }
  86. return nil
  87. }
  88. // shiftLow shifts the low value for 8 bit. The shifted byte is written into
  89. // the byte writer. The cache value is used to handle overflows.
  90. func (e *rangeEncoder) shiftLow() error {
  91. if uint32(e.low) < 0xff000000 || (e.low>>32) != 0 {
  92. tmp := e.cache
  93. for {
  94. err := e.writeByte(tmp + byte(e.low>>32))
  95. if err != nil {
  96. return err
  97. }
  98. tmp = 0xff
  99. e.cacheLen--
  100. if e.cacheLen <= 0 {
  101. if e.cacheLen < 0 {
  102. panic("negative cacheLen")
  103. }
  104. break
  105. }
  106. }
  107. e.cache = byte(uint32(e.low) >> 24)
  108. }
  109. e.cacheLen++
  110. e.low = uint64(uint32(e.low) << 8)
  111. return nil
  112. }
  113. // rangeDecoder decodes single bits of the range encoding stream.
  114. type rangeDecoder struct {
  115. br io.ByteReader
  116. nrange uint32
  117. code uint32
  118. }
  119. // newRangeDecoder initializes a range decoder. It reads five bytes from the
  120. // reader and therefore may return an error.
  121. func newRangeDecoder(br io.ByteReader) (d *rangeDecoder, err error) {
  122. d = &rangeDecoder{br: br, nrange: 0xffffffff}
  123. b, err := d.br.ReadByte()
  124. if err != nil {
  125. return nil, err
  126. }
  127. if b != 0 {
  128. return nil, errors.New("newRangeDecoder: first byte not zero")
  129. }
  130. for i := 0; i < 4; i++ {
  131. if err = d.updateCode(); err != nil {
  132. return nil, err
  133. }
  134. }
  135. if d.code >= d.nrange {
  136. return nil, errors.New("newRangeDecoder: d.code >= d.nrange")
  137. }
  138. return d, nil
  139. }
  140. // possiblyAtEnd checks whether the decoder may be at the end of the stream.
  141. func (d *rangeDecoder) possiblyAtEnd() bool {
  142. return d.code == 0
  143. }
  144. // DirectDecodeBit decodes a bit with probability 1/2. The return value b will
  145. // contain the bit at the least-significant position. All other bits will be
  146. // zero.
  147. func (d *rangeDecoder) DirectDecodeBit() (b uint32, err error) {
  148. d.nrange >>= 1
  149. d.code -= d.nrange
  150. t := 0 - (d.code >> 31)
  151. d.code += d.nrange & t
  152. b = (t + 1) & 1
  153. // d.code will stay less then d.nrange
  154. // normalize
  155. // assume d.code < d.nrange
  156. const top = 1 << 24
  157. if d.nrange >= top {
  158. return b, nil
  159. }
  160. d.nrange <<= 8
  161. // d.code < d.nrange will be maintained
  162. return b, d.updateCode()
  163. }
  164. // decodeBit decodes a single bit. The bit will be returned at the
  165. // least-significant position. All other bits will be zero. The probability
  166. // value will be updated.
  167. func (d *rangeDecoder) DecodeBit(p *prob) (b uint32, err error) {
  168. bound := p.bound(d.nrange)
  169. if d.code < bound {
  170. d.nrange = bound
  171. p.inc()
  172. b = 0
  173. } else {
  174. d.code -= bound
  175. d.nrange -= bound
  176. p.dec()
  177. b = 1
  178. }
  179. // normalize
  180. // assume d.code < d.nrange
  181. const top = 1 << 24
  182. if d.nrange >= top {
  183. return b, nil
  184. }
  185. d.nrange <<= 8
  186. // d.code < d.nrange will be maintained
  187. return b, d.updateCode()
  188. }
  189. // updateCode reads a new byte into the code.
  190. func (d *rangeDecoder) updateCode() error {
  191. b, err := d.br.ReadByte()
  192. if err != nil {
  193. return err
  194. }
  195. d.code = (d.code << 8) | uint32(b)
  196. return nil
  197. }