blockenc.go 24 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893
  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. "errors"
  7. "fmt"
  8. "math"
  9. "math/bits"
  10. "slices"
  11. "github.com/klauspost/compress/huff0"
  12. )
  13. type blockEnc struct {
  14. size int
  15. literals []byte
  16. sequences []seq
  17. coders seqCoders
  18. litEnc *huff0.Scratch
  19. dictLitEnc *huff0.Scratch
  20. wr bitWriter
  21. extraLits int
  22. output []byte
  23. recentOffsets [3]uint32
  24. prevRecentOffsets [3]uint32
  25. last bool
  26. lowMem bool
  27. }
  28. // init should be used once the block has been created.
  29. // If called more than once, the effect is the same as calling reset.
  30. func (b *blockEnc) init() {
  31. if b.lowMem {
  32. // 1K literals
  33. if cap(b.literals) < 1<<10 {
  34. b.literals = make([]byte, 0, 1<<10)
  35. }
  36. const defSeqs = 20
  37. if cap(b.sequences) < defSeqs {
  38. b.sequences = make([]seq, 0, defSeqs)
  39. }
  40. // 1K
  41. if cap(b.output) < 1<<10 {
  42. b.output = make([]byte, 0, 1<<10)
  43. }
  44. } else {
  45. if cap(b.literals) < maxCompressedBlockSize {
  46. b.literals = make([]byte, 0, maxCompressedBlockSize)
  47. }
  48. const defSeqs = 2000
  49. if cap(b.sequences) < defSeqs {
  50. b.sequences = make([]seq, 0, defSeqs)
  51. }
  52. if cap(b.output) < maxCompressedBlockSize {
  53. b.output = make([]byte, 0, maxCompressedBlockSize)
  54. }
  55. }
  56. if b.coders.mlEnc == nil {
  57. b.coders.mlEnc = &fseEncoder{}
  58. b.coders.mlPrev = &fseEncoder{}
  59. b.coders.ofEnc = &fseEncoder{}
  60. b.coders.ofPrev = &fseEncoder{}
  61. b.coders.llEnc = &fseEncoder{}
  62. b.coders.llPrev = &fseEncoder{}
  63. }
  64. b.litEnc = &huff0.Scratch{WantLogLess: 4}
  65. b.reset(nil)
  66. }
  67. // initNewEncode can be used to reset offsets and encoders to the initial state.
  68. func (b *blockEnc) initNewEncode() {
  69. b.recentOffsets = [3]uint32{1, 4, 8}
  70. b.litEnc.Reuse = huff0.ReusePolicyNone
  71. b.coders.setPrev(nil, nil, nil)
  72. b.dictLitEnc = nil
  73. }
  74. // reset will reset the block for a new encode, but in the same stream,
  75. // meaning that state will be carried over, but the block content is reset.
  76. // If a previous block is provided, the recent offsets are carried over.
  77. func (b *blockEnc) reset(prev *blockEnc) {
  78. b.extraLits = 0
  79. b.literals = b.literals[:0]
  80. b.size = 0
  81. b.sequences = b.sequences[:0]
  82. b.output = b.output[:0]
  83. b.last = false
  84. if prev != nil {
  85. b.recentOffsets = prev.prevRecentOffsets
  86. }
  87. b.dictLitEnc = nil
  88. }
  89. // reset will reset the block for a new encode, but in the same stream,
  90. // meaning that state will be carried over, but the block content is reset.
  91. // If a previous block is provided, the recent offsets are carried over.
  92. func (b *blockEnc) swapEncoders(prev *blockEnc) {
  93. b.coders.swap(&prev.coders)
  94. b.litEnc, prev.litEnc = prev.litEnc, b.litEnc
  95. }
  96. // blockHeader contains the information for a block header.
  97. type blockHeader uint32
  98. // setLast sets the 'last' indicator on a block.
  99. func (h *blockHeader) setLast(b bool) {
  100. if b {
  101. *h = *h | 1
  102. } else {
  103. const mask = (1 << 24) - 2
  104. *h = *h & mask
  105. }
  106. }
  107. // setSize will store the compressed size of a block.
  108. func (h *blockHeader) setSize(v uint32) {
  109. const mask = 7
  110. *h = (*h)&mask | blockHeader(v<<3)
  111. }
  112. // setType sets the block type.
  113. func (h *blockHeader) setType(t blockType) {
  114. const mask = 1 | (((1 << 24) - 1) ^ 7)
  115. *h = (*h & mask) | blockHeader(t<<1)
  116. }
  117. // appendTo will append the block header to a slice.
  118. func (h blockHeader) appendTo(b []byte) []byte {
  119. return append(b, uint8(h), uint8(h>>8), uint8(h>>16))
  120. }
  121. // String returns a string representation of the block.
  122. func (h blockHeader) String() string {
  123. return fmt.Sprintf("Type: %d, Size: %d, Last:%t", (h>>1)&3, h>>3, h&1 == 1)
  124. }
  125. // literalsHeader contains literals header information.
  126. type literalsHeader uint64
  127. // setType can be used to set the type of literal block.
  128. func (h *literalsHeader) setType(t literalsBlockType) {
  129. const mask = math.MaxUint64 - 3
  130. *h = (*h & mask) | literalsHeader(t)
  131. }
  132. // setSize can be used to set a single size, for uncompressed and RLE content.
  133. func (h *literalsHeader) setSize(regenLen int) {
  134. inBits := bits.Len32(uint32(regenLen))
  135. // Only retain 2 bits
  136. const mask = 3
  137. lh := uint64(*h & mask)
  138. switch {
  139. case inBits < 5:
  140. lh |= (uint64(regenLen) << 3) | (1 << 60)
  141. if debugEncoder {
  142. got := int(lh>>3) & 0xff
  143. if got != regenLen {
  144. panic(fmt.Sprint("litRegenSize = ", regenLen, "(want) != ", got, "(got)"))
  145. }
  146. }
  147. case inBits < 12:
  148. lh |= (1 << 2) | (uint64(regenLen) << 4) | (2 << 60)
  149. case inBits < 20:
  150. lh |= (3 << 2) | (uint64(regenLen) << 4) | (3 << 60)
  151. default:
  152. panic(fmt.Errorf("internal error: block too big (%d)", regenLen))
  153. }
  154. *h = literalsHeader(lh)
  155. }
  156. // setSizes will set the size of a compressed literals section and the input length.
  157. func (h *literalsHeader) setSizes(compLen, inLen int, single bool) {
  158. compBits, inBits := bits.Len32(uint32(compLen)), bits.Len32(uint32(inLen))
  159. // Only retain 2 bits
  160. const mask = 3
  161. lh := uint64(*h & mask)
  162. switch {
  163. case compBits <= 10 && inBits <= 10:
  164. if !single {
  165. lh |= 1 << 2
  166. }
  167. lh |= (uint64(inLen) << 4) | (uint64(compLen) << (10 + 4)) | (3 << 60)
  168. if debugEncoder {
  169. const mmask = (1 << 24) - 1
  170. n := (lh >> 4) & mmask
  171. if int(n&1023) != inLen {
  172. panic(fmt.Sprint("regensize:", int(n&1023), "!=", inLen, inBits))
  173. }
  174. if int(n>>10) != compLen {
  175. panic(fmt.Sprint("compsize:", int(n>>10), "!=", compLen, compBits))
  176. }
  177. }
  178. case compBits <= 14 && inBits <= 14:
  179. lh |= (2 << 2) | (uint64(inLen) << 4) | (uint64(compLen) << (14 + 4)) | (4 << 60)
  180. if single {
  181. panic("single stream used with more than 10 bits length.")
  182. }
  183. case compBits <= 18 && inBits <= 18:
  184. lh |= (3 << 2) | (uint64(inLen) << 4) | (uint64(compLen) << (18 + 4)) | (5 << 60)
  185. if single {
  186. panic("single stream used with more than 10 bits length.")
  187. }
  188. default:
  189. panic("internal error: block too big")
  190. }
  191. *h = literalsHeader(lh)
  192. }
  193. // appendTo will append the literals header to a byte slice.
  194. func (h literalsHeader) appendTo(b []byte) []byte {
  195. size := uint8(h >> 60)
  196. switch size {
  197. case 1:
  198. b = append(b, uint8(h))
  199. case 2:
  200. b = append(b, uint8(h), uint8(h>>8))
  201. case 3:
  202. b = append(b, uint8(h), uint8(h>>8), uint8(h>>16))
  203. case 4:
  204. b = append(b, uint8(h), uint8(h>>8), uint8(h>>16), uint8(h>>24))
  205. case 5:
  206. b = append(b, uint8(h), uint8(h>>8), uint8(h>>16), uint8(h>>24), uint8(h>>32))
  207. default:
  208. panic(fmt.Errorf("internal error: literalsHeader has invalid size (%d)", size))
  209. }
  210. return b
  211. }
  212. // size returns the output size with currently set values.
  213. func (h literalsHeader) size() int {
  214. return int(h >> 60)
  215. }
  216. func (h literalsHeader) String() string {
  217. return fmt.Sprintf("Type: %d, SizeFormat: %d, Size: 0x%d, Bytes:%d", literalsBlockType(h&3), (h>>2)&3, h&((1<<60)-1)>>4, h>>60)
  218. }
  219. // pushOffsets will push the recent offsets to the backup store.
  220. func (b *blockEnc) pushOffsets() {
  221. b.prevRecentOffsets = b.recentOffsets
  222. }
  223. // pushOffsets will push the recent offsets to the backup store.
  224. func (b *blockEnc) popOffsets() {
  225. b.recentOffsets = b.prevRecentOffsets
  226. }
  227. // matchOffset will adjust recent offsets and return the adjusted one,
  228. // if it matches a previous offset.
  229. func (b *blockEnc) matchOffset(offset, lits uint32) uint32 {
  230. // Check if offset is one of the recent offsets.
  231. // Adjusts the output offset accordingly.
  232. // Gives a tiny bit of compression, typically around 1%.
  233. if true {
  234. if lits > 0 {
  235. switch offset {
  236. case b.recentOffsets[0]:
  237. offset = 1
  238. case b.recentOffsets[1]:
  239. b.recentOffsets[1] = b.recentOffsets[0]
  240. b.recentOffsets[0] = offset
  241. offset = 2
  242. case b.recentOffsets[2]:
  243. b.recentOffsets[2] = b.recentOffsets[1]
  244. b.recentOffsets[1] = b.recentOffsets[0]
  245. b.recentOffsets[0] = offset
  246. offset = 3
  247. default:
  248. b.recentOffsets[2] = b.recentOffsets[1]
  249. b.recentOffsets[1] = b.recentOffsets[0]
  250. b.recentOffsets[0] = offset
  251. offset += 3
  252. }
  253. } else {
  254. switch offset {
  255. case b.recentOffsets[1]:
  256. b.recentOffsets[1] = b.recentOffsets[0]
  257. b.recentOffsets[0] = offset
  258. offset = 1
  259. case b.recentOffsets[2]:
  260. b.recentOffsets[2] = b.recentOffsets[1]
  261. b.recentOffsets[1] = b.recentOffsets[0]
  262. b.recentOffsets[0] = offset
  263. offset = 2
  264. case b.recentOffsets[0] - 1:
  265. b.recentOffsets[2] = b.recentOffsets[1]
  266. b.recentOffsets[1] = b.recentOffsets[0]
  267. b.recentOffsets[0] = offset
  268. offset = 3
  269. default:
  270. b.recentOffsets[2] = b.recentOffsets[1]
  271. b.recentOffsets[1] = b.recentOffsets[0]
  272. b.recentOffsets[0] = offset
  273. offset += 3
  274. }
  275. }
  276. } else {
  277. offset += 3
  278. }
  279. return offset
  280. }
  281. // encodeRaw can be used to set the output to a raw representation of supplied bytes.
  282. func (b *blockEnc) encodeRaw(a []byte) {
  283. var bh blockHeader
  284. bh.setLast(b.last)
  285. bh.setSize(uint32(len(a)))
  286. bh.setType(blockTypeRaw)
  287. b.output = bh.appendTo(b.output[:0])
  288. b.output = append(b.output, a...)
  289. if debugEncoder {
  290. println("Adding RAW block, length", len(a), "last:", b.last)
  291. }
  292. }
  293. // encodeRaw can be used to set the output to a raw representation of supplied bytes.
  294. func (b *blockEnc) encodeRawTo(dst, src []byte) []byte {
  295. var bh blockHeader
  296. bh.setLast(b.last)
  297. bh.setSize(uint32(len(src)))
  298. bh.setType(blockTypeRaw)
  299. dst = bh.appendTo(dst)
  300. dst = append(dst, src...)
  301. if debugEncoder {
  302. println("Adding RAW block, length", len(src), "last:", b.last)
  303. }
  304. return dst
  305. }
  306. // encodeLits can be used if the block is only litLen.
  307. func (b *blockEnc) encodeLits(lits []byte, raw bool) error {
  308. var bh blockHeader
  309. bh.setLast(b.last)
  310. bh.setSize(uint32(len(lits)))
  311. // Don't compress extremely small blocks
  312. if len(lits) < 8 || (len(lits) < 32 && b.dictLitEnc == nil) || raw {
  313. if debugEncoder {
  314. println("Adding RAW block, length", len(lits), "last:", b.last)
  315. }
  316. bh.setType(blockTypeRaw)
  317. b.output = bh.appendTo(b.output)
  318. b.output = append(b.output, lits...)
  319. return nil
  320. }
  321. var (
  322. out []byte
  323. reUsed, single bool
  324. err error
  325. )
  326. if b.dictLitEnc != nil {
  327. b.litEnc.TransferCTable(b.dictLitEnc)
  328. b.litEnc.Reuse = huff0.ReusePolicyAllow
  329. b.dictLitEnc = nil
  330. }
  331. if len(lits) >= 1024 {
  332. // Use 4 Streams.
  333. out, reUsed, err = huff0.Compress4X(lits, b.litEnc)
  334. } else if len(lits) > 16 {
  335. // Use 1 stream
  336. single = true
  337. out, reUsed, err = huff0.Compress1X(lits, b.litEnc)
  338. } else {
  339. err = huff0.ErrIncompressible
  340. }
  341. if err == nil && len(out)+5 > len(lits) {
  342. // If we are close, we may still be worse or equal to raw.
  343. var lh literalsHeader
  344. lh.setSizes(len(out), len(lits), single)
  345. if len(out)+lh.size() >= len(lits) {
  346. err = huff0.ErrIncompressible
  347. }
  348. }
  349. switch err {
  350. case huff0.ErrIncompressible:
  351. if debugEncoder {
  352. println("Adding RAW block, length", len(lits), "last:", b.last)
  353. }
  354. bh.setType(blockTypeRaw)
  355. b.output = bh.appendTo(b.output)
  356. b.output = append(b.output, lits...)
  357. return nil
  358. case huff0.ErrUseRLE:
  359. if debugEncoder {
  360. println("Adding RLE block, length", len(lits))
  361. }
  362. bh.setType(blockTypeRLE)
  363. b.output = bh.appendTo(b.output)
  364. b.output = append(b.output, lits[0])
  365. return nil
  366. case nil:
  367. default:
  368. return err
  369. }
  370. // Compressed...
  371. // Now, allow reuse
  372. b.litEnc.Reuse = huff0.ReusePolicyAllow
  373. bh.setType(blockTypeCompressed)
  374. var lh literalsHeader
  375. if reUsed {
  376. if debugEncoder {
  377. println("Reused tree, compressed to", len(out))
  378. }
  379. lh.setType(literalsBlockTreeless)
  380. } else {
  381. if debugEncoder {
  382. println("New tree, compressed to", len(out), "tree size:", len(b.litEnc.OutTable))
  383. }
  384. lh.setType(literalsBlockCompressed)
  385. }
  386. // Set sizes
  387. lh.setSizes(len(out), len(lits), single)
  388. bh.setSize(uint32(len(out) + lh.size() + 1))
  389. // Write block headers.
  390. b.output = bh.appendTo(b.output)
  391. b.output = lh.appendTo(b.output)
  392. // Add compressed data.
  393. b.output = append(b.output, out...)
  394. // No sequences.
  395. b.output = append(b.output, 0)
  396. return nil
  397. }
  398. // encodeRLE will encode an RLE block.
  399. func (b *blockEnc) encodeRLE(val byte, length uint32) {
  400. var bh blockHeader
  401. bh.setLast(b.last)
  402. bh.setSize(length)
  403. bh.setType(blockTypeRLE)
  404. b.output = bh.appendTo(b.output)
  405. b.output = append(b.output, val)
  406. }
  407. // fuzzFseEncoder can be used to fuzz the FSE encoder.
  408. func fuzzFseEncoder(data []byte) int {
  409. if len(data) > maxSequences || len(data) < 2 {
  410. return 0
  411. }
  412. enc := fseEncoder{}
  413. hist := enc.Histogram()
  414. maxSym := uint8(0)
  415. for i, v := range data {
  416. v = v & 63
  417. data[i] = v
  418. hist[v]++
  419. if v > maxSym {
  420. maxSym = v
  421. }
  422. }
  423. if maxSym == 0 {
  424. // All 0
  425. return 0
  426. }
  427. cnt := int(slices.Max(hist[:maxSym]))
  428. if cnt == len(data) {
  429. // RLE
  430. return 0
  431. }
  432. enc.HistogramFinished(maxSym, cnt)
  433. err := enc.normalizeCount(len(data))
  434. if err != nil {
  435. return 0
  436. }
  437. _, err = enc.writeCount(nil)
  438. if err != nil {
  439. panic(err)
  440. }
  441. return 1
  442. }
  443. // encode will encode the block and append the output in b.output.
  444. // Previous offset codes must be pushed if more blocks are expected.
  445. func (b *blockEnc) encode(org []byte, raw, rawAllLits bool) error {
  446. if len(b.sequences) == 0 {
  447. return b.encodeLits(b.literals, rawAllLits)
  448. }
  449. if len(b.sequences) == 1 && len(org) > 0 && len(b.literals) <= 1 {
  450. // Check common RLE cases.
  451. seq := b.sequences[0]
  452. if seq.litLen == uint32(len(b.literals)) && seq.offset-3 == 1 {
  453. // Offset == 1 and 0 or 1 literals.
  454. b.encodeRLE(org[0], b.sequences[0].matchLen+zstdMinMatch+seq.litLen)
  455. return nil
  456. }
  457. }
  458. // We want some difference to at least account for the headers.
  459. saved := b.size - len(b.literals) - (b.size >> 6)
  460. if saved < 16 {
  461. if org == nil {
  462. return errIncompressible
  463. }
  464. b.popOffsets()
  465. return b.encodeLits(org, rawAllLits)
  466. }
  467. var bh blockHeader
  468. var lh literalsHeader
  469. bh.setLast(b.last)
  470. bh.setType(blockTypeCompressed)
  471. // Store offset of the block header. Needed when we know the size.
  472. bhOffset := len(b.output)
  473. b.output = bh.appendTo(b.output)
  474. var (
  475. out []byte
  476. reUsed, single bool
  477. err error
  478. )
  479. if b.dictLitEnc != nil {
  480. b.litEnc.TransferCTable(b.dictLitEnc)
  481. b.litEnc.Reuse = huff0.ReusePolicyAllow
  482. b.dictLitEnc = nil
  483. }
  484. if len(b.literals) >= 1024 && !raw {
  485. // Use 4 Streams.
  486. out, reUsed, err = huff0.Compress4X(b.literals, b.litEnc)
  487. } else if len(b.literals) > 16 && !raw {
  488. // Use 1 stream
  489. single = true
  490. out, reUsed, err = huff0.Compress1X(b.literals, b.litEnc)
  491. } else {
  492. err = huff0.ErrIncompressible
  493. }
  494. if err == nil && len(out)+5 > len(b.literals) {
  495. // If we are close, we may still be worse or equal to raw.
  496. var lh literalsHeader
  497. lh.setSize(len(b.literals))
  498. szRaw := lh.size()
  499. lh.setSizes(len(out), len(b.literals), single)
  500. szComp := lh.size()
  501. if len(out)+szComp >= len(b.literals)+szRaw {
  502. err = huff0.ErrIncompressible
  503. }
  504. }
  505. switch err {
  506. case huff0.ErrIncompressible:
  507. lh.setType(literalsBlockRaw)
  508. lh.setSize(len(b.literals))
  509. b.output = lh.appendTo(b.output)
  510. b.output = append(b.output, b.literals...)
  511. if debugEncoder {
  512. println("Adding literals RAW, length", len(b.literals))
  513. }
  514. case huff0.ErrUseRLE:
  515. lh.setType(literalsBlockRLE)
  516. lh.setSize(len(b.literals))
  517. b.output = lh.appendTo(b.output)
  518. b.output = append(b.output, b.literals[0])
  519. if debugEncoder {
  520. println("Adding literals RLE")
  521. }
  522. case nil:
  523. // Compressed litLen...
  524. if reUsed {
  525. if debugEncoder {
  526. println("reused tree")
  527. }
  528. lh.setType(literalsBlockTreeless)
  529. } else {
  530. if debugEncoder {
  531. println("new tree, size:", len(b.litEnc.OutTable))
  532. }
  533. lh.setType(literalsBlockCompressed)
  534. if debugEncoder {
  535. _, _, err := huff0.ReadTable(out, nil)
  536. if err != nil {
  537. panic(err)
  538. }
  539. }
  540. }
  541. lh.setSizes(len(out), len(b.literals), single)
  542. if debugEncoder {
  543. printf("Compressed %d literals to %d bytes", len(b.literals), len(out))
  544. println("Adding literal header:", lh)
  545. }
  546. b.output = lh.appendTo(b.output)
  547. b.output = append(b.output, out...)
  548. b.litEnc.Reuse = huff0.ReusePolicyAllow
  549. if debugEncoder {
  550. println("Adding literals compressed")
  551. }
  552. default:
  553. if debugEncoder {
  554. println("Adding literals ERROR:", err)
  555. }
  556. return err
  557. }
  558. // Sequence compression
  559. // Write the number of sequences
  560. switch {
  561. case len(b.sequences) < 128:
  562. b.output = append(b.output, uint8(len(b.sequences)))
  563. case len(b.sequences) < 0x7f00: // TODO: this could be wrong
  564. n := len(b.sequences)
  565. b.output = append(b.output, 128+uint8(n>>8), uint8(n))
  566. default:
  567. n := len(b.sequences) - 0x7f00
  568. b.output = append(b.output, 255, uint8(n), uint8(n>>8))
  569. }
  570. if debugEncoder {
  571. println("Encoding", len(b.sequences), "sequences")
  572. }
  573. b.genCodes()
  574. llEnc := b.coders.llEnc
  575. ofEnc := b.coders.ofEnc
  576. mlEnc := b.coders.mlEnc
  577. err = llEnc.normalizeCount(len(b.sequences))
  578. if err != nil {
  579. return err
  580. }
  581. err = ofEnc.normalizeCount(len(b.sequences))
  582. if err != nil {
  583. return err
  584. }
  585. err = mlEnc.normalizeCount(len(b.sequences))
  586. if err != nil {
  587. return err
  588. }
  589. // Choose the best compression mode for each type.
  590. // Will evaluate the new vs predefined and previous.
  591. chooseComp := func(cur, prev, preDef *fseEncoder) (*fseEncoder, seqCompMode) {
  592. // See if predefined/previous is better
  593. hist := cur.count[:cur.symbolLen]
  594. nSize := cur.approxSize(hist) + cur.maxHeaderSize()
  595. predefSize := preDef.approxSize(hist)
  596. prevSize := prev.approxSize(hist)
  597. // Add a small penalty for new encoders.
  598. // Don't bother with extremely small (<2 byte gains).
  599. nSize = nSize + (nSize+2*8*16)>>4
  600. switch {
  601. case predefSize <= prevSize && predefSize <= nSize || forcePreDef:
  602. if debugEncoder {
  603. println("Using predefined", predefSize>>3, "<=", nSize>>3)
  604. }
  605. return preDef, compModePredefined
  606. case prevSize <= nSize:
  607. if debugEncoder {
  608. println("Using previous", prevSize>>3, "<=", nSize>>3)
  609. }
  610. return prev, compModeRepeat
  611. default:
  612. if debugEncoder {
  613. println("Using new, predef", predefSize>>3, ". previous:", prevSize>>3, ">", nSize>>3, "header max:", cur.maxHeaderSize()>>3, "bytes")
  614. println("tl:", cur.actualTableLog, "symbolLen:", cur.symbolLen, "norm:", cur.norm[:cur.symbolLen], "hist", cur.count[:cur.symbolLen])
  615. }
  616. return cur, compModeFSE
  617. }
  618. }
  619. // Write compression mode
  620. var mode uint8
  621. if llEnc.useRLE {
  622. mode |= uint8(compModeRLE) << 6
  623. llEnc.setRLE(b.sequences[0].llCode)
  624. if debugEncoder {
  625. println("llEnc.useRLE")
  626. }
  627. } else {
  628. var m seqCompMode
  629. llEnc, m = chooseComp(llEnc, b.coders.llPrev, &fsePredefEnc[tableLiteralLengths])
  630. mode |= uint8(m) << 6
  631. }
  632. if ofEnc.useRLE {
  633. mode |= uint8(compModeRLE) << 4
  634. ofEnc.setRLE(b.sequences[0].ofCode)
  635. if debugEncoder {
  636. println("ofEnc.useRLE")
  637. }
  638. } else {
  639. var m seqCompMode
  640. ofEnc, m = chooseComp(ofEnc, b.coders.ofPrev, &fsePredefEnc[tableOffsets])
  641. mode |= uint8(m) << 4
  642. }
  643. if mlEnc.useRLE {
  644. mode |= uint8(compModeRLE) << 2
  645. mlEnc.setRLE(b.sequences[0].mlCode)
  646. if debugEncoder {
  647. println("mlEnc.useRLE, code: ", b.sequences[0].mlCode, "value", b.sequences[0].matchLen)
  648. }
  649. } else {
  650. var m seqCompMode
  651. mlEnc, m = chooseComp(mlEnc, b.coders.mlPrev, &fsePredefEnc[tableMatchLengths])
  652. mode |= uint8(m) << 2
  653. }
  654. b.output = append(b.output, mode)
  655. if debugEncoder {
  656. printf("Compression modes: 0b%b", mode)
  657. }
  658. b.output, err = llEnc.writeCount(b.output)
  659. if err != nil {
  660. return err
  661. }
  662. start := len(b.output)
  663. b.output, err = ofEnc.writeCount(b.output)
  664. if err != nil {
  665. return err
  666. }
  667. if false {
  668. println("block:", b.output[start:], "tablelog", ofEnc.actualTableLog, "maxcount:", ofEnc.maxCount)
  669. fmt.Printf("selected TableLog: %d, Symbol length: %d\n", ofEnc.actualTableLog, ofEnc.symbolLen)
  670. for i, v := range ofEnc.norm[:ofEnc.symbolLen] {
  671. fmt.Printf("%3d: %5d -> %4d \n", i, ofEnc.count[i], v)
  672. }
  673. }
  674. b.output, err = mlEnc.writeCount(b.output)
  675. if err != nil {
  676. return err
  677. }
  678. // Maybe in block?
  679. wr := &b.wr
  680. wr.reset(b.output)
  681. var ll, of, ml cState
  682. // Current sequence
  683. seq := len(b.sequences) - 1
  684. s := b.sequences[seq]
  685. llEnc.setBits(llBitsTable[:])
  686. mlEnc.setBits(mlBitsTable[:])
  687. ofEnc.setBits(nil)
  688. llTT, ofTT, mlTT := llEnc.ct.symbolTT[:256], ofEnc.ct.symbolTT[:256], mlEnc.ct.symbolTT[:256]
  689. // We have 3 bounds checks here (and in the loop).
  690. // Since we are iterating backwards it is kinda hard to avoid.
  691. llB, ofB, mlB := llTT[s.llCode], ofTT[s.ofCode], mlTT[s.mlCode]
  692. ll.init(wr, &llEnc.ct, llB)
  693. of.init(wr, &ofEnc.ct, ofB)
  694. wr.flush32()
  695. ml.init(wr, &mlEnc.ct, mlB)
  696. // Each of these lookups also generates a bounds check.
  697. wr.addBits32NC(s.litLen, llB.outBits)
  698. wr.addBits32NC(s.matchLen, mlB.outBits)
  699. wr.flush32()
  700. wr.addBits32NC(s.offset, ofB.outBits)
  701. if debugSequences {
  702. println("Encoded seq", seq, s, "codes:", s.llCode, s.mlCode, s.ofCode, "states:", ll.state, ml.state, of.state, "bits:", llB, mlB, ofB)
  703. }
  704. seq--
  705. // Store sequences in reverse...
  706. for seq >= 0 {
  707. s = b.sequences[seq]
  708. ofB := ofTT[s.ofCode]
  709. wr.flush32() // tablelog max is below 8 for each, so it will fill max 24 bits.
  710. //of.encode(ofB)
  711. nbBitsOut := (uint32(of.state) + ofB.deltaNbBits) >> 16
  712. dstState := int32(of.state>>(nbBitsOut&15)) + int32(ofB.deltaFindState)
  713. wr.addBits16NC(of.state, uint8(nbBitsOut))
  714. of.state = of.stateTable[dstState]
  715. // Accumulate extra bits.
  716. outBits := ofB.outBits & 31
  717. extraBits := uint64(s.offset & bitMask32[outBits])
  718. extraBitsN := outBits
  719. mlB := mlTT[s.mlCode]
  720. //ml.encode(mlB)
  721. nbBitsOut = (uint32(ml.state) + mlB.deltaNbBits) >> 16
  722. dstState = int32(ml.state>>(nbBitsOut&15)) + int32(mlB.deltaFindState)
  723. wr.addBits16NC(ml.state, uint8(nbBitsOut))
  724. ml.state = ml.stateTable[dstState]
  725. outBits = mlB.outBits & 31
  726. extraBits = extraBits<<outBits | uint64(s.matchLen&bitMask32[outBits])
  727. extraBitsN += outBits
  728. llB := llTT[s.llCode]
  729. //ll.encode(llB)
  730. nbBitsOut = (uint32(ll.state) + llB.deltaNbBits) >> 16
  731. dstState = int32(ll.state>>(nbBitsOut&15)) + int32(llB.deltaFindState)
  732. wr.addBits16NC(ll.state, uint8(nbBitsOut))
  733. ll.state = ll.stateTable[dstState]
  734. outBits = llB.outBits & 31
  735. extraBits = extraBits<<outBits | uint64(s.litLen&bitMask32[outBits])
  736. extraBitsN += outBits
  737. wr.flush32()
  738. wr.addBits64NC(extraBits, extraBitsN)
  739. if debugSequences {
  740. println("Encoded seq", seq, s)
  741. }
  742. seq--
  743. }
  744. ml.flush(mlEnc.actualTableLog)
  745. of.flush(ofEnc.actualTableLog)
  746. ll.flush(llEnc.actualTableLog)
  747. wr.close()
  748. b.output = wr.out
  749. // Maybe even add a bigger margin.
  750. if len(b.output)-3-bhOffset >= b.size {
  751. // Discard and encode as raw block.
  752. b.output = b.encodeRawTo(b.output[:bhOffset], org)
  753. b.popOffsets()
  754. b.litEnc.Reuse = huff0.ReusePolicyNone
  755. return nil
  756. }
  757. // Size is output minus block header.
  758. bh.setSize(uint32(len(b.output)-bhOffset) - 3)
  759. if debugEncoder {
  760. println("Rewriting block header", bh)
  761. }
  762. _ = bh.appendTo(b.output[bhOffset:bhOffset])
  763. b.coders.setPrev(llEnc, mlEnc, ofEnc)
  764. return nil
  765. }
  766. var errIncompressible = errors.New("incompressible")
  767. func (b *blockEnc) genCodes() {
  768. if len(b.sequences) == 0 {
  769. // nothing to do
  770. return
  771. }
  772. if len(b.sequences) > math.MaxUint16 {
  773. panic("can only encode up to 64K sequences")
  774. }
  775. // No bounds checks after here:
  776. llH := b.coders.llEnc.Histogram()
  777. ofH := b.coders.ofEnc.Histogram()
  778. mlH := b.coders.mlEnc.Histogram()
  779. for i := range llH {
  780. llH[i] = 0
  781. }
  782. for i := range ofH {
  783. ofH[i] = 0
  784. }
  785. for i := range mlH {
  786. mlH[i] = 0
  787. }
  788. var llMax, ofMax, mlMax uint8
  789. for i := range b.sequences {
  790. seq := &b.sequences[i]
  791. v := llCode(seq.litLen)
  792. seq.llCode = v
  793. llH[v]++
  794. if v > llMax {
  795. llMax = v
  796. }
  797. v = ofCode(seq.offset)
  798. seq.ofCode = v
  799. ofH[v]++
  800. if v > ofMax {
  801. ofMax = v
  802. }
  803. v = mlCode(seq.matchLen)
  804. seq.mlCode = v
  805. mlH[v]++
  806. if v > mlMax {
  807. mlMax = v
  808. if debugAsserts && mlMax > maxMatchLengthSymbol {
  809. panic(fmt.Errorf("mlMax > maxMatchLengthSymbol (%d), matchlen: %d", mlMax, seq.matchLen))
  810. }
  811. }
  812. }
  813. if debugAsserts && mlMax > maxMatchLengthSymbol {
  814. panic(fmt.Errorf("mlMax > maxMatchLengthSymbol (%d)", mlMax))
  815. }
  816. if debugAsserts && ofMax > maxOffsetBits {
  817. panic(fmt.Errorf("ofMax > maxOffsetBits (%d)", ofMax))
  818. }
  819. if debugAsserts && llMax > maxLiteralLengthSymbol {
  820. panic(fmt.Errorf("llMax > maxLiteralLengthSymbol (%d)", llMax))
  821. }
  822. b.coders.mlEnc.HistogramFinished(mlMax, int(slices.Max(mlH[:mlMax+1])))
  823. b.coders.ofEnc.HistogramFinished(ofMax, int(slices.Max(ofH[:ofMax+1])))
  824. b.coders.llEnc.HistogramFinished(llMax, int(slices.Max(llH[:llMax+1])))
  825. }