huff0.go 8.9 KB

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