huff0.go 8.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335
  1. // Package huff0 provides fast huffman encoding as used in zstd.
  2. //
  3. // See README.md at https://github.com/klauspost/compress/tree/master/huff0 for details.
  4. package huff0
  5. import (
  6. "errors"
  7. "fmt"
  8. "math"
  9. "math/bits"
  10. "github.com/klauspost/compress/fse"
  11. )
  12. const (
  13. maxSymbolValue = 255
  14. // zstandard limits tablelog to 11, see:
  15. // https://github.com/facebook/zstd/blob/dev/doc/zstd_compression_format.md#huffman-tree-description
  16. tableLogMax = 11
  17. tableLogDefault = 11
  18. minTablelog = 5
  19. huffNodesLen = 512
  20. // BlockSizeMax is maximum input size for a single block uncompressed.
  21. BlockSizeMax = 1<<18 - 1
  22. )
  23. var (
  24. // ErrIncompressible is returned when input is judged to be too hard to compress.
  25. ErrIncompressible = errors.New("input is not compressible")
  26. // ErrUseRLE is returned from the compressor when the input is a single byte value repeated.
  27. ErrUseRLE = errors.New("input is single value repeated")
  28. // ErrTooBig is return if input is too large for a single block.
  29. ErrTooBig = errors.New("input too big")
  30. // ErrMaxDecodedSizeExceeded is return if input is too large for a single block.
  31. ErrMaxDecodedSizeExceeded = errors.New("maximum output size exceeded")
  32. )
  33. type ReusePolicy uint8
  34. const (
  35. // ReusePolicyAllow will allow reuse if it produces smaller output.
  36. ReusePolicyAllow ReusePolicy = iota
  37. // ReusePolicyPrefer will re-use aggressively if possible.
  38. // This will not check if a new table will produce smaller output,
  39. // except if the current table is impossible to use or
  40. // compressed output is bigger than input.
  41. ReusePolicyPrefer
  42. // ReusePolicyNone will disable re-use of tables.
  43. // This is slightly faster than ReusePolicyAllow but may produce larger output.
  44. ReusePolicyNone
  45. // ReusePolicyMust must allow reuse and produce smaller output.
  46. ReusePolicyMust
  47. )
  48. type Scratch struct {
  49. count [maxSymbolValue + 1]uint32
  50. // Per block parameters.
  51. // These can be used to override compression parameters of the block.
  52. // Do not touch, unless you know what you are doing.
  53. // Out is output buffer.
  54. // If the scratch is re-used before the caller is done processing the output,
  55. // set this field to nil.
  56. // Otherwise the output buffer will be re-used for next Compression/Decompression step
  57. // and allocation will be avoided.
  58. Out []byte
  59. // OutTable will contain the table data only, if a new table has been generated.
  60. // Slice of the returned data.
  61. OutTable []byte
  62. // OutData will contain the compressed data.
  63. // Slice of the returned data.
  64. OutData []byte
  65. // MaxDecodedSize will set the maximum allowed output size.
  66. // This value will automatically be set to BlockSizeMax if not set.
  67. // Decoders will return ErrMaxDecodedSizeExceeded is this limit is exceeded.
  68. MaxDecodedSize int
  69. br byteReader
  70. // MaxSymbolValue will override the maximum symbol value of the next block.
  71. MaxSymbolValue uint8
  72. // TableLog will attempt to override the tablelog for the next block.
  73. // Must be <= 11 and >= 5.
  74. TableLog uint8
  75. // Reuse will specify the reuse policy
  76. Reuse ReusePolicy
  77. // WantLogLess allows to specify a log 2 reduction that should at least be achieved,
  78. // otherwise the block will be returned as incompressible.
  79. // The reduction should then at least be (input size >> WantLogLess)
  80. // If WantLogLess == 0 any improvement will do.
  81. WantLogLess uint8
  82. symbolLen uint16 // Length of active part of the symbol table.
  83. maxCount int // count of the most probable symbol
  84. clearCount bool // clear count
  85. actualTableLog uint8 // Selected tablelog.
  86. prevTableLog uint8 // Tablelog for previous table
  87. prevTable cTable // Table used for previous compression.
  88. cTable cTable // compression table
  89. dt dTable // decompression table
  90. nodes []nodeElt
  91. tmpOut [4][]byte
  92. fse *fse.Scratch
  93. huffWeight [maxSymbolValue + 1]byte
  94. }
  95. // TransferCTable will transfer the previously used compression table.
  96. func (s *Scratch) TransferCTable(src *Scratch) {
  97. if cap(s.prevTable) < len(src.prevTable) {
  98. s.prevTable = make(cTable, 0, maxSymbolValue+1)
  99. }
  100. s.prevTable = s.prevTable[:len(src.prevTable)]
  101. copy(s.prevTable, src.prevTable)
  102. s.prevTableLog = src.prevTableLog
  103. }
  104. func (s *Scratch) prepare(in []byte) (*Scratch, error) {
  105. if len(in) > BlockSizeMax {
  106. return nil, ErrTooBig
  107. }
  108. if s == nil {
  109. s = &Scratch{}
  110. }
  111. if s.MaxSymbolValue == 0 {
  112. s.MaxSymbolValue = maxSymbolValue
  113. }
  114. if s.TableLog == 0 {
  115. s.TableLog = tableLogDefault
  116. }
  117. if s.TableLog > tableLogMax || s.TableLog < minTablelog {
  118. return nil, fmt.Errorf(" invalid tableLog %d (%d -> %d)", s.TableLog, minTablelog, tableLogMax)
  119. }
  120. if s.MaxDecodedSize <= 0 || s.MaxDecodedSize > BlockSizeMax {
  121. s.MaxDecodedSize = BlockSizeMax
  122. }
  123. if s.clearCount && s.maxCount == 0 {
  124. for i := range s.count {
  125. s.count[i] = 0
  126. }
  127. s.clearCount = false
  128. }
  129. if cap(s.Out) == 0 {
  130. s.Out = make([]byte, 0, len(in))
  131. }
  132. s.Out = s.Out[:0]
  133. s.OutTable = nil
  134. s.OutData = nil
  135. if cap(s.nodes) < huffNodesLen+1 {
  136. s.nodes = make([]nodeElt, 0, huffNodesLen+1)
  137. }
  138. s.nodes = s.nodes[:0]
  139. if s.fse == nil {
  140. s.fse = &fse.Scratch{}
  141. }
  142. s.br.init(in)
  143. return s, nil
  144. }
  145. type cTable []cTableEntry
  146. func (c cTable) write(s *Scratch) error {
  147. var (
  148. // precomputed conversion table
  149. bitsToWeight [tableLogMax + 1]byte
  150. huffLog = s.actualTableLog
  151. // last weight is not saved.
  152. maxSymbolValue = uint8(s.symbolLen - 1)
  153. huffWeight = s.huffWeight[:256]
  154. )
  155. const (
  156. maxFSETableLog = 6
  157. )
  158. // convert to weight
  159. bitsToWeight[0] = 0
  160. for n := uint8(1); n < huffLog+1; n++ {
  161. bitsToWeight[n] = huffLog + 1 - n
  162. }
  163. // Acquire histogram for FSE.
  164. hist := s.fse.Histogram()
  165. hist = hist[:256]
  166. for i := range hist[:16] {
  167. hist[i] = 0
  168. }
  169. for n := uint8(0); n < maxSymbolValue; n++ {
  170. v := bitsToWeight[c[n].nBits] & 15
  171. huffWeight[n] = v
  172. hist[v]++
  173. }
  174. // FSE compress if feasible.
  175. if maxSymbolValue >= 2 {
  176. huffMaxCnt := uint32(0)
  177. huffMax := uint8(0)
  178. for i, v := range hist[:16] {
  179. if v == 0 {
  180. continue
  181. }
  182. huffMax = byte(i)
  183. if v > huffMaxCnt {
  184. huffMaxCnt = v
  185. }
  186. }
  187. s.fse.HistogramFinished(huffMax, int(huffMaxCnt))
  188. s.fse.TableLog = maxFSETableLog
  189. b, err := fse.Compress(huffWeight[:maxSymbolValue], s.fse)
  190. if err == nil && len(b) < int(s.symbolLen>>1) {
  191. s.Out = append(s.Out, uint8(len(b)))
  192. s.Out = append(s.Out, b...)
  193. return nil
  194. }
  195. // Unable to compress (RLE/uncompressible)
  196. }
  197. // write raw values as 4-bits (max : 15)
  198. if maxSymbolValue > (256 - 128) {
  199. // should not happen : likely means source cannot be compressed
  200. return ErrIncompressible
  201. }
  202. op := s.Out
  203. // special case, pack weights 4 bits/weight.
  204. op = append(op, 128|(maxSymbolValue-1))
  205. // be sure it doesn't cause msan issue in final combination
  206. huffWeight[maxSymbolValue] = 0
  207. for n := uint16(0); n < uint16(maxSymbolValue); n += 2 {
  208. op = append(op, (huffWeight[n]<<4)|huffWeight[n+1])
  209. }
  210. s.Out = op
  211. return nil
  212. }
  213. func (c cTable) estTableSize(s *Scratch) (sz int, err error) {
  214. var (
  215. // precomputed conversion table
  216. bitsToWeight [tableLogMax + 1]byte
  217. huffLog = s.actualTableLog
  218. // last weight is not saved.
  219. maxSymbolValue = uint8(s.symbolLen - 1)
  220. huffWeight = s.huffWeight[:256]
  221. )
  222. const (
  223. maxFSETableLog = 6
  224. )
  225. // convert to weight
  226. bitsToWeight[0] = 0
  227. for n := uint8(1); n < huffLog+1; n++ {
  228. bitsToWeight[n] = huffLog + 1 - n
  229. }
  230. // Acquire histogram for FSE.
  231. hist := s.fse.Histogram()
  232. hist = hist[:256]
  233. for i := range hist[:16] {
  234. hist[i] = 0
  235. }
  236. for n := uint8(0); n < maxSymbolValue; n++ {
  237. v := bitsToWeight[c[n].nBits] & 15
  238. huffWeight[n] = v
  239. hist[v]++
  240. }
  241. // FSE compress if feasible.
  242. if maxSymbolValue >= 2 {
  243. huffMaxCnt := uint32(0)
  244. huffMax := uint8(0)
  245. for i, v := range hist[:16] {
  246. if v == 0 {
  247. continue
  248. }
  249. huffMax = byte(i)
  250. if v > huffMaxCnt {
  251. huffMaxCnt = v
  252. }
  253. }
  254. s.fse.HistogramFinished(huffMax, int(huffMaxCnt))
  255. s.fse.TableLog = maxFSETableLog
  256. b, err := fse.Compress(huffWeight[:maxSymbolValue], s.fse)
  257. if err == nil && len(b) < int(s.symbolLen>>1) {
  258. sz += 1 + len(b)
  259. return sz, nil
  260. }
  261. // Unable to compress (RLE/uncompressible)
  262. }
  263. // write raw values as 4-bits (max : 15)
  264. if maxSymbolValue > (256 - 128) {
  265. // should not happen : likely means source cannot be compressed
  266. return 0, ErrIncompressible
  267. }
  268. // special case, pack weights 4 bits/weight.
  269. sz += 1 + int(maxSymbolValue/2)
  270. return sz, nil
  271. }
  272. // estimateSize returns the estimated size in bytes of the input represented in the
  273. // histogram supplied.
  274. func (c cTable) estimateSize(hist []uint32) int {
  275. nbBits := uint32(7)
  276. for i, v := range c[:len(hist)] {
  277. nbBits += uint32(v.nBits) * hist[i]
  278. }
  279. return int(nbBits >> 3)
  280. }
  281. // minSize returns the minimum possible size considering the shannon limit.
  282. func (s *Scratch) minSize(total int) int {
  283. nbBits := float64(7)
  284. fTotal := float64(total)
  285. for _, v := range s.count[:s.symbolLen] {
  286. n := float64(v)
  287. if n > 0 {
  288. nbBits += math.Log2(fTotal/n) * n
  289. }
  290. }
  291. return int(nbBits) >> 3
  292. }
  293. func highBit32(val uint32) (n uint32) {
  294. return uint32(bits.Len32(val) - 1)
  295. }